"""
constraint_graph_svg.py — Ahmad's ConstraintGraph compiler extension
Compiles ConstraintGraph XML to both SVG visualization and executable runtime dict.
"""

import xml.etree.ElementTree as ET
from collections import defaultdict, deque
from typing import List, Dict, Any


# ─── Node type emoji icons ────────────────────────────────────────────────────

NODE_ICONS = {
    "Input":       "\U0001f9e0",  # 🧠
    "Memory":      "\U0001f4da",  # 📚
    "Retrieval":   "\U0001f50d",  # 🔍
    "Transform":   "⚙",      # ⚙
    "Constraint":  "⚖",      # ⚖
    "Proof":       "\U0001f510",  # 🔐
    "Output":      "\U0001f310",  # 🌐
}

# ─── SVG layout constants ─────────────────────────────────────────────────────

CANVAS_W = 400
CANVAS_H = 600
NODE_W = 160
NODE_H = 40
NODE_RX = 8
CENTER_X = 200


def _node_top_y(index: int, total: int) -> int:
    """Return the y-coordinate of the top edge of node at position index."""
    margin_top = 30
    spacing = (CANVAS_H - margin_top - NODE_H) // max(total - 1, 1)
    return margin_top + index * spacing


# ─── Topological sort (Kahn's algorithm) ─────────────────────────────────────

def _topological_sort(nodes: List[dict], edges: List[dict]) -> List[str]:
    in_degree: Dict[str, int] = {n["id"]: 0 for n in nodes}
    adjacency: Dict[str, List[str]] = defaultdict(list)

    for e in edges:
        adjacency[e["from"]].append(e["to"])
        in_degree[e["to"]] += 1

    queue = deque(nid for nid, deg in in_degree.items() if deg == 0)
    order = []

    while queue:
        nid = queue.popleft()
        order.append(nid)
        for neighbor in adjacency[nid]:
            in_degree[neighbor] -= 1
            if in_degree[neighbor] == 0:
                queue.append(neighbor)

    if len(order) != len(nodes):
        raise ValueError("ConstraintGraph contains a cycle — topological sort failed.")

    return order


# ─── XML parsing ─────────────────────────────────────────────────────────────

def _parse_graph_xml(xml_source: str):
    root = ET.fromstring(xml_source)
    if root.tag != "graph":
        raise ValueError(f"Root element must be <graph>, got <{root.tag}>")

    nodes = []
    edges = []

    for child in root:
        if child.tag == "node":
            nodes.append({"id": child.attrib["id"], "type": child.attrib["type"]})
        elif child.tag == "edge":
            edges.append({"from": child.attrib["from"], "to": child.attrib["to"]})

    return nodes, edges


# ─── SVG renderer ────────────────────────────────────────────────────────────

def _render_svg(nodes: List[dict], edges: List[dict], pipeline: List[str]) -> str:
    # Build index map: node id → position in pipeline
    pos_map = {nid: i for i, nid in enumerate(pipeline)}
    type_map = {n["id"]: n["type"] for n in nodes}
    total = len(pipeline)

    svg_lines = [
        f'<svg xmlns="http://www.w3.org/2000/svg" '
        f'width="{CANVAS_W}" height="{CANVAS_H}" '
        f'viewBox="0 0 {CANVAS_W} {CANVAS_H}">',
        "  <!-- ConstraintGraph SVG — generated by constraint_graph_svg.py -->",
        "  <defs>",
        "    <marker id=\"arrowhead\" markerWidth=\"10\" markerHeight=\"7\" "
        "refX=\"10\" refY=\"3.5\" orient=\"auto\">",
        "      <polygon points=\"0 0, 10 3.5, 0 7\" fill=\"#444\"/>",
        "    </marker>",
        "  </defs>",
    ]

    # Compute node rects
    node_rects: Dict[str, dict] = {}
    for nid in pipeline:
        idx = pos_map[nid]
        x = CENTER_X - NODE_W // 2
        y = _node_top_y(idx, total)
        node_rects[nid] = {"x": x, "y": y}

    # Draw edges first (behind nodes)
    for edge in edges:
        src = edge["from"]
        dst = edge["to"]
        if src not in node_rects or dst not in node_rects:
            continue
        sx = CENTER_X
        sy = node_rects[src]["y"] + NODE_H
        dx = CENTER_X
        dy = node_rects[dst]["y"]
        svg_lines.append(
            f'  <line x1="{sx}" y1="{sy}" x2="{dx}" y2="{dy}" '
            f'stroke="#444" stroke-width="2" marker-end="url(#arrowhead)"/>'
        )

    # Draw nodes
    for nid in pipeline:
        node_type = type_map[nid]
        icon = NODE_ICONS.get(node_type, "?")
        x = node_rects[nid]["x"]
        y = node_rects[nid]["y"]
        label = f"{icon} {node_type}"
        cx = x + NODE_W // 2
        cy = y + NODE_H // 2 + 5  # vertical center of text

        svg_lines.append(
            f'  <rect x="{x}" y="{y}" width="{NODE_W}" height="{NODE_H}" '
            f'rx="{NODE_RX}" fill="white" stroke="#333" stroke-width="1.5"/>'
        )
        svg_lines.append(
            f'  <text x="{cx}" y="{cy}" text-anchor="middle" '
            f'font-family="sans-serif" font-size="13" fill="#111">'
            f'{label}</text>'
        )

    svg_lines.append("</svg>")
    return "\n".join(svg_lines)


# ─── ConstraintGraphCompiler ──────────────────────────────────────────────────

class ConstraintGraphCompiler:
    """
    Ahmad's ConstraintGraph compiler.

    Takes ConstraintGraph XML and compiles to:
      - SVG visualization (compile_to_svg)
      - Executable runtime dict (compile_to_runtime)
    """

    def compile_to_svg(self, xml_source: str) -> str:
        """
        Compile ConstraintGraph XML to an SVG string.
        Nodes are laid out vertically in topological order with arrowhead edges.
        """
        nodes, edges = _parse_graph_xml(xml_source)
        pipeline = _topological_sort(nodes, edges)
        return _render_svg(nodes, edges, pipeline)

    def compile_to_runtime(self, xml_source: str) -> Dict[str, Any]:
        """
        Compile ConstraintGraph XML to an executable runtime representation.

        Returns:
            {
                "nodes":    [{"id": ..., "type": ...}, ...],
                "edges":    [{"from": ..., "to": ...}, ...],
                "pipeline": [<topologically sorted node ids>]
            }
        """
        nodes, edges = _parse_graph_xml(xml_source)
        pipeline = _topological_sort(nodes, edges)
        return {
            "nodes": nodes,
            "edges": edges,
            "pipeline": pipeline,
        }


# ─── CLI ──────────────────────────────────────────────────────────────────────

def _cli():
    import sys
    if len(sys.argv) < 2 or sys.argv[1] in ("-h", "--help"):
        print("Usage: constraint-graph <input.xml> [output.svg|--runtime]")
        print("  Compiles a <graph> XML to SVG or prints the runtime dict.")
        sys.exit(0)
    src = open(sys.argv[1], encoding="utf-8").read()
    compiler = ConstraintGraphCompiler()
    if len(sys.argv) >= 3 and sys.argv[2] == "--runtime":
        import json
        print(json.dumps(compiler.compile_to_runtime(src), indent=2))
    else:
        svg = compiler.compile_to_svg(src)
        if len(sys.argv) >= 3:
            open(sys.argv[2], "w", encoding="utf-8").write(svg)
        else:
            print(svg)


if __name__ == "__main__":
    _cli()