#!/usr/bin/env python3
"""셰이더 그래프 → GLSL 컴퓨트 셰이더 코드 생성.

그래프(configs/shader_graphs/<이름>.json)는 노드와 간선이다. 노드의 연산(op)은 노드 라이브러리
(src/vkgs/shaders/nodes/*.glsl)의 함수 n_<op> 이고, 그 **시그니처를 읽어** 입력 포트·매개변수·타입을 얻는다
(라이브러리 쪽 '리플렉션' — 노드를 더할 때 이 스크립트는 고치지 않는다). 특수 노드 두 개:
  input  — 이미지 픽셀을 읽는다 (vec3)
  output — 입력 하나를 이미지에 쓴다
노드를 위상 정렬해 한 줄씩 함수 호출로 펼치고, 간선 타입이 맞는지·순환·빠진 입력을 검사한다.

결과 (OUT 디렉터리):
  <이름>.comp         렌더 그래프의 compute 패스가 쓰는 셰이더 (CMake 가 SPIR-V 로 굽는다)
  <이름>.params.json  매개변수 표 (이름·버퍼 위치·기본값·범위) — 뷰어가 슬라이더와 초깃값에 쓴다
  <이름>.dot          그림 (dot -Tsvg)

사용: shader_graph.py GRAPH.json --shaders src/vkgs/shaders --out build/shaders
시스템 파이썬(3.8)에서도 돈다 — 표준 라이브러리만 쓴다.
"""

from __future__ import annotations

import argparse
import json
import re
import sys
from pathlib import Path

SPECIAL = ("input", "output")
_FUNC = re.compile(r"^\s*(\w+)\s+n_(\w+)\s*\(([^)]*)\)", re.M)
_PARAM_DOC = re.compile(r"^\s*//\s*@param\s+(\w+)\s+(\S+)\s+(\S+)\s+(\S+)")


class GraphError(Exception):
    pass


def parse_library(text: str) -> dict:
    """노드 라이브러리 GLSL → {op: {ret, inputs: [(이름, 타입)], ctx, params: [(이름, 기본, 최소, 최대)]}}."""
    lines = text.splitlines()
    ops = {}
    for m in _FUNC.finditer(text):
        ret, op, args = m.group(1), m.group(2), m.group(3)
        if ret == "return":  # 본문 안의 호출 (return n_x(...))
            continue
        line_no = text.count("\n", 0, m.start())
        # 바로 위 주석 블록의 @param
        docs = {}
        k = line_no - 1
        while k >= 0 and lines[k].strip().startswith("//"):
            d = _PARAM_DOC.match(lines[k])
            if d:
                docs[d.group(1)] = tuple(float(d.group(i)) for i in (2, 3, 4))
            k -= 1
        inputs, params, ctx = [], [], False
        for a in [s.strip() for s in args.split(",") if s.strip()]:
            parts = a.split()
            if len(parts) != 2:
                raise GraphError(f"노드 n_{op}: 인자를 읽을 수 없음 '{a}'")
            typ, name = parts
            if typ == "Ctx":
                ctx = True
            elif ctx:
                if typ != "float":
                    raise GraphError(f"노드 n_{op}: 매개변수 {name} 는 float 여야 함 ({typ})")
                params.append((name,) + docs.get(name, (0.0, 0.0, 1.0)))
            else:
                inputs.append((name, typ))
        if op in ops:
            raise GraphError(f"노드 n_{op} 가 두 번 정의됨")
        ops[op] = {"ret": ret, "inputs": inputs, "ctx": ctx, "params": params}
    return ops


def _topo(nodes: dict) -> list:
    """Kahn — 같은 단계는 선언 순서. 순환이면 GraphError."""
    ids = list(nodes)
    indeg = {i: 0 for i in ids}
    outs = {i: [] for i in ids}
    for i, n in nodes.items():
        for src in n.get("in", []):
            if src not in nodes:
                raise GraphError(f"노드 {i}: 없는 노드 '{src}' 를 입력으로 씀")
            outs[src].append(i)
            indeg[i] += 1
    ready = [i for i in ids if indeg[i] == 0]
    order = []
    while ready:
        ready.sort(key=ids.index)
        i = ready.pop(0)
        order.append(i)
        for j in outs[i]:
            indeg[j] -= 1
            if indeg[j] == 0:
                ready.append(j)
    if len(order) != len(ids):
        raise GraphError("셰이더 그래프에 순환이 있음: " + ", ".join(i for i in ids if i not in order))
    return order


def generate(graph: dict, lib_ops: dict, name: str, includes: list) -> tuple:
    """그래프 → (GLSL 소스, 매개변수 표, DOT)."""
    nodes = {k: v for k, v in graph["nodes"].items() if not k.startswith("_")}
    target = graph.get("target", {})
    image = target.get("image", "kImgPost")
    order = _topo(nodes)
    types, params, body = {}, [], []
    n_out = 0
    for i in order:
        n = nodes[i]
        op = n.get("op")
        ins = n.get("in", [])
        if op == "input":
            if ins:
                raise GraphError(f"노드 {i}: input 은 입력이 없다")
            img = n.get("image", image)
            body.append(f"    vec3 v_{i} = imageLoad(gImgF4[{img}], x.px).rgb;")
            types[i] = "vec3"
            continue
        if op == "output":
            if len(ins) != 1 or types[ins[0]] != "vec3":
                raise GraphError(f"노드 {i}: output 은 vec3 입력 하나")
            img = n.get("image", image)
            body.append(f"    imageStore(gImgF4[{img}], x.px, vec4(v_{ins[0]}, imageLoad(gImgF4[{img}], x.px).a));")
            n_out += 1
            continue
        if op not in lib_ops:
            raise GraphError(f"노드 {i}: 라이브러리에 없는 연산 '{op}' (있는 것: {', '.join(sorted(lib_ops))})")
        spec = lib_ops[op]
        if len(ins) != len(spec["inputs"]):
            raise GraphError(f"노드 {i} ({op}): 입력 {len(spec['inputs'])}개가 필요한데 {len(ins)}개")
        for src, (port, typ) in zip(ins, spec["inputs"]):
            if types[src] != typ:
                raise GraphError(f"노드 {i} ({op}): 입력 {port} 는 {typ} 인데 {src} 는 {types[src]}")
        given = n.get("params", {})
        unknown = set(given) - {p[0] for p in spec["params"]}
        if unknown:
            raise GraphError(f"노드 {i} ({op}): 모르는 매개변수 {sorted(unknown)}")
        args = [f"v_{s}" for s in ins]
        if spec["ctx"]:
            args.append("x")
        for pname, default, lo, hi in spec["params"]:
            idx = len(params)
            params.append({"name": f"{i}.{pname}", "node": i, "op": op, "index": idx,
                           "default": float(given.get(pname, default)), "min": lo, "max": hi})
            args.append(f"P({idx})")
        body.append(f"    {spec['ret']} v_{i} = n_{op}({', '.join(args)});")
        types[i] = spec["ret"]
    if n_out == 0:
        raise GraphError("output 노드가 없음")

    inc = "\n".join(f'#include "{p}"' for p in includes)
    src = f"""#version 450
#extension GL_GOOGLE_include_directive : require
#extension GL_EXT_shader_image_load_formatted : require
// 자동 생성 — 셰이더 그래프 '{name}' (scripts/shader_graph.py). 고치지 말고 그래프 JSON·노드 라이브러리를 고칠 것.
#include "bindless.glsl"
layout(local_size_x = 8, local_size_y = 8) in;
layout(set = 0, binding = 4) uniform image2D gImgF4[];  // 형식 없음 — 렌더 그래프 JSON 이 형식을 정한다
layout(std430, set = 0, binding = 0) readonly buffer GraphParams {{ float v[]; }} gGraphParams[];
// slots.x = 매개변수 버퍼의 SSBO 슬롯, ctx = (시각 s, 프레임 번호 비트)
layout(push_constant) uniform Push {{ uvec4 slots; vec4 ctx; }} pc;
#define P(i) gGraphParams[pc.slots.x].v[i]
{inc}

void main() {{
    ivec2 size = imageSize(gImgF4[{image}]);
    Ctx x;
    x.px = ivec2(gl_GlobalInvocationID.xy);
    if (x.px.x >= size.x || x.px.y >= size.y) return;
    x.size = vec2(size);
    x.uv = (vec2(x.px) + 0.5) / x.size;
    x.time = pc.ctx.x;
    x.frame = floatBitsToUint(pc.ctx.y);
{chr(10).join(body)}
}}
"""
    table = {"graph": name, "resource": target.get("params_resource"), "count": len(params), "params": params}
    return src, table, to_dot(name, nodes, order, lib_ops)


def to_dot(name: str, nodes: dict, order: list, lib_ops: dict) -> str:
    out = [f'digraph "{name}" {{', "  rankdir=LR;", '  node [fontname="NanumGothic", shape=record, style=filled];']
    for i in order:
        n = nodes[i]
        op = n.get("op")
        spec = lib_ops.get(op, {"params": []})
        ps = "\\l".join(f"{p[0]} = {n.get('params', {}).get(p[0], p[1]):g}" for p in spec["params"])
        color = "#d5f5e3" if op in SPECIAL else "#fde8c8"
        label = f"{i}\\n[{op}]" + (f"|{ps}\\l" if ps else "")
        out.append(f'  "{i}" [label="{{{label}}}", fillcolor="{color}"];')
        for s in n.get("in", []):
            out.append(f'  "{s}" -> "{i}";')
    out.append("}")
    return "\n".join(out) + "\n"


def build(graph_path: Path, shaders: Path, out: Path) -> Path:
    graph = json.loads(graph_path.read_text())
    name = graph.get("name", graph_path.stem)
    includes = graph.get("library", ["nodes/isp.glsl"])
    lib_ops = {}
    for inc in includes:
        for op, spec in parse_library((shaders / inc).read_text()).items():
            if op in lib_ops:
                raise GraphError(f"노드 n_{op} 가 라이브러리 여러 곳에 있음")
            lib_ops[op] = spec
    src, table, dot = generate(graph, lib_ops, name, includes)
    out.mkdir(parents=True, exist_ok=True)
    comp = out / f"{name}.comp"
    # 내용이 같으면 쓰지 않는다 (불필요한 재컴파일 방지)
    if not comp.exists() or comp.read_text() != src:
        comp.write_text(src)
    (out / f"{name}.params.json").write_text(json.dumps(table, indent=2, ensure_ascii=False) + "\n")
    (out / f"{name}.dot").write_text(dot)
    return comp


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("graph", type=Path)
    ap.add_argument("--shaders", type=Path, required=True, help="노드 라이브러리 기준 디렉터리 (src/vkgs/shaders)")
    ap.add_argument("--out", type=Path, required=True)
    a = ap.parse_args(argv)
    try:
        comp = build(a.graph, a.shaders, a.out)
    except GraphError as e:
        print(f"셰이더 그래프 오류 ({a.graph}): {e}", file=sys.stderr)
        return 1
    print(f"셰이더 그래프 {a.graph.name} → {comp}")
    return 0


if __name__ == "__main__":
    sys.exit(main())