from __future__ import annotations

import json
import keyword
import pprint
import re
from collections import defaultdict
from typing import Any, Mapping

from ..workflow import WORKFLOW_KIND, WORKFLOW_SCHEMA_VERSION, WorkflowRunError, infer_entrypoint, validate_workflow


def export_langgraph_workflow(data: Mapping[str, Any]) -> tuple[str, list[str]]:
    report = validate_workflow(data)
    if not report.ok:
        raise WorkflowRunError(json.dumps(report.to_dict(), indent=2))

    entry_node, entry_port = infer_entrypoint(data)
    node_meta = _node_meta(data)
    edges = [dict(edge) for edge in data.get("edges") or []]
    node_names = _unique_node_names(node_meta)

    output_ports = {
        node_id: _output_ports(node_id, meta, edges, (entry_node, entry_port))
        for node_id, meta in node_meta.items()
    }
    state_keys = _state_keys(node_names, output_ports)
    incoming = _incoming_edges(edges)
    predecessors = _predecessors(edges)
    successors = _successors(edges)
    runtime_nodes = _runtime_nodes(node_meta)
    workflow_name = str(data.get("name") or "Blacknode Workflow")
    entry_name = node_names[entry_node]
    blacknode_workflow = _blacknode_workflow_payload(data, {"node_id": entry_node, "port": entry_port})

    lines: list[str] = [
        "from __future__ import annotations",
        "",
        f"# Generated by Blacknode from workflow: {workflow_name}",
        f"# Entrypoint: {entry_name}.{entry_port}",
        "# Keep the visual graph as the source of truth, then regenerate this file when the graph changes.",
        "",
        "from typing import Any, TypedDict",
        "",
        "import blacknode as bn",
        "from blacknode.node import _NODE_REGISTRY",
        "from langgraph.graph import END, START, StateGraph",
        "",
        "",
        f"_BLACKNODE_WORKFLOW = {_pretty_literal(blacknode_workflow)}",
        "",
        f"_RUNTIME_NODES = {_pretty_literal(runtime_nodes)}",
        f"_RUNTIME_EDGES = {_pretty_literal(edges)}",
        "",
        "# Blacknode runtime graph used inside each LangGraph node function.",
        "_RUNTIME_GRAPH = bn.Graph()",
        "_RUNTIME_GRAPH._nodes = _RUNTIME_NODES",
        "_RUNTIME_GRAPH._edges = _RUNTIME_EDGES",
        "_RUNTIME_GRAPH._dirty = set(_RUNTIME_NODES)",
        "_RUNTIME_GRAPH._cache = {}",
        "",
        "",
        "class WorkflowState(TypedDict, total=False):",
    ]

    for key in sorted({key for outputs in state_keys.values() for key in outputs.values()}):
        lines.append(f"    {key}: Any")
    if not any(state_keys.values()):
        lines.append("    value: Any")

    lines.extend([
        "",
        "",
        "def _run_blacknode_node(",
        "    node_id: str,",
        "    type_name: str,",
        "    ctx: dict[str, Any],",
        "    output_ports: list[str],",
        ") -> dict[str, Any]:",
        "    ctx['__graph__'] = _RUNTIME_GRAPH",
        "    ctx['__node_id__'] = node_id",
        "    if type_name == 'Subnet':",
        "        result: dict[str, Any] = {}",
        "        for output_port in output_ports:",
        "            result.update(_RUNTIME_GRAPH._cook_subnet(node_id, output_port, ctx))",
        "        return result",
        "    value = _NODE_REGISTRY[type_name](ctx)",
        "    return value if isinstance(value, dict) else {'output': value}",
        "",
        "",
    ])

    for step, (node_id, meta) in enumerate(node_meta.items(), start=1):
        function_name = f"run_{node_names[node_id]}"
        type_name = str(meta.get("type"))
        params = dict(meta.get("params") or {})
        label = _node_comment_label(node_id, meta)
        lines.append(f"# Step {step}: {label}")
        lines.append(f"# Blacknode node id: {node_id!r}, type: {type_name!r}")
        lines.append(f"def {function_name}(state: WorkflowState) -> dict[str, Any]:")
        lines.append(f"    ctx: dict[str, Any] = {params!r}")
        for edge in incoming.get(node_id, []):
            from_id = str(edge.get("from"))
            from_port = str(edge.get("from_port", "output"))
            to_port = str(edge.get("to_port", "input"))
            source_key = state_keys[from_id][from_port]
            source_name = node_names[from_id]
            lines.append(f"    # {source_name}.{from_port} -> {node_names[node_id]}.{to_port}")
            lines.append(f"    ctx[{to_port!r}] = state.get({source_key!r})")
        lines.append(
            f"    result = _run_blacknode_node({node_id!r}, {type_name!r}, ctx, {output_ports[node_id]!r})"
        )
        if output_ports[node_id]:
            returned = {
                state_keys[node_id][output_port]: f"result.get({output_port!r})"
                for output_port in output_ports[node_id]
            }
            lines.append("    return {")
            for key, expr in returned.items():
                lines.append(f"        {key!r}: {expr},")
            lines.append("    }")
        else:
            lines.append("    return {}")
        lines.append("")

    lines.extend([
        "",
        "# Build the LangGraph StateGraph from the visual Blacknode DAG.",
        "workflow = StateGraph(WorkflowState)",
    ])

    lines.append("# Register nodes.")
    for node_id in node_meta:
        graph_name = node_names[node_id]
        lines.append(f"workflow.add_node({graph_name!r}, run_{graph_name})")

    lines.append("")
    lines.append("# Preserve upstream data dependencies.")
    for node_id in node_meta:
        graph_name = node_names[node_id]
        preds = [node_names[pred] for pred in predecessors.get(node_id, []) if pred in node_names]
        if not preds:
            lines.append(f"workflow.add_edge(START, {graph_name!r})")
        elif len(preds) == 1:
            lines.append(f"workflow.add_edge({preds[0]!r}, {graph_name!r})")
        else:
            lines.append(f"workflow.add_edge({preds!r}, {graph_name!r})")

    lines.append("")
    lines.append("# Finish branches that have no downstream nodes.")
    for node_id in node_meta:
        if not successors.get(node_id):
            lines.append(f"workflow.add_edge({node_names[node_id]!r}, END)")

    entry_key = state_keys[entry_node][entry_port]
    lines.extend([
        "",
        "compiled = workflow.compile()",
        "",
        "",
        "if __name__ == '__main__':",
        "    result = compiled.invoke({})",
        f"    print(result.get({entry_key!r}))",
        "",
    ])

    return "\n".join(lines), ["Install LangGraph before running this export: pip install langgraph"]


def _node_meta(data: Mapping[str, Any]) -> dict[str, dict[str, Any]]:
    return {
        str(node_id): dict(meta)
        for node_id, meta in (data.get("node_meta") or {}).items()
        if isinstance(meta, Mapping)
    }


def _runtime_nodes(node_meta: Mapping[str, Mapping[str, Any]]) -> dict[str, dict[str, Any]]:
    runtime: dict[str, dict[str, Any]] = {}
    for node_id, meta in node_meta.items():
        node_def = {
            "type": str(meta.get("type")),
            "params": dict(meta.get("params") or {}),
        }
        if isinstance(meta.get("subgraph"), Mapping):
            node_def["subgraph"] = dict(meta["subgraph"])
        runtime[str(node_id)] = node_def
    return runtime


def _blacknode_workflow_payload(data: Mapping[str, Any], entrypoint: dict[str, str]) -> dict[str, Any]:
    payload: dict[str, Any] = {
        "kind": WORKFLOW_KIND,
        "schema_version": WORKFLOW_SCHEMA_VERSION,
        "name": str(data.get("name") or "Blacknode Workflow"),
        "node_meta": _node_meta(data),
        "edges": [dict(edge) for edge in data.get("edges") or []],
        "entrypoint": dict(entrypoint),
    }
    saved_at = data.get("saved_at")
    if saved_at:
        payload["saved_at"] = saved_at
    metadata = data.get("metadata")
    if isinstance(metadata, Mapping):
        payload["metadata"] = dict(metadata)
    return payload


def _incoming_edges(edges: list[dict[str, Any]]) -> dict[str, list[dict[str, Any]]]:
    incoming: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for edge in edges:
        incoming[str(edge.get("to"))].append(edge)
    return incoming


def _predecessors(edges: list[dict[str, Any]]) -> dict[str, list[str]]:
    result: dict[str, list[str]] = defaultdict(list)
    for edge in edges:
        from_id = str(edge.get("from"))
        to_id = str(edge.get("to"))
        if from_id not in result[to_id]:
            result[to_id].append(from_id)
    return result


def _successors(edges: list[dict[str, Any]]) -> dict[str, list[str]]:
    result: dict[str, list[str]] = defaultdict(list)
    for edge in edges:
        from_id = str(edge.get("from"))
        to_id = str(edge.get("to"))
        if to_id not in result[from_id]:
            result[from_id].append(to_id)
    return result


def _output_ports(
    node_id: str,
    meta: Mapping[str, Any],
    edges: list[dict[str, Any]],
    entrypoint: tuple[str, str],
) -> list[str]:
    ports: list[str] = [str(port) for port in meta.get("outputs") or []]
    for edge in edges:
        if str(edge.get("from")) == node_id:
            port = str(edge.get("from_port", "output"))
            if port not in ports:
                ports.append(port)
    if entrypoint[0] == node_id and entrypoint[1] not in ports:
        ports.append(entrypoint[1])
    return ports or ["output"]


def _state_keys(
    node_names: Mapping[str, str],
    output_ports: Mapping[str, list[str]],
) -> dict[str, dict[str, str]]:
    used: set[str] = set()
    keys: dict[str, dict[str, str]] = {}
    for node_id, ports in output_ports.items():
        node_keys: dict[str, str] = {}
        for port in ports:
            base = f"{node_names[node_id]}__{_identifier(port) or 'output'}"
            node_keys[port] = _dedupe(base, used)
        keys[node_id] = node_keys
    return keys


def _unique_node_names(node_meta: Mapping[str, Mapping[str, Any]]) -> dict[str, str]:
    used: set[str] = {"workflow", "compiled", "result", "state", "start", "end"}
    result: dict[str, str] = {}
    for index, (node_id, meta) in enumerate(node_meta.items(), start=1):
        base = _identifier(_node_name_source(str(node_id), meta)) or f"node_{index}"
        result[str(node_id)] = _dedupe(base, used)
    return result


def _node_name_source(node_id: str, meta: Mapping[str, Any]) -> str:
    type_name = str(meta.get("type") or "")
    if type_name in {"Output", "SubnetOutput"}:
        return "output"
    if _descriptive_node_id(node_id):
        return node_id
    return _node_label(meta) or node_id


def _node_comment_label(node_id: str, meta: Mapping[str, Any]) -> str:
    type_name = str(meta.get("type") or "Node")
    params = meta.get("params") if isinstance(meta.get("params"), Mapping) else {}
    if type_name == "Text":
        value = params.get("value") if isinstance(params, Mapping) else None
        if isinstance(value, str) and value.strip():
            return f"{type_name} {value.strip()!r}"
    label = _node_label(meta)
    if label and label != type_name:
        return f"{label} ({type_name})"
    return f"{type_name} ({node_id})"


def _node_label(meta: Mapping[str, Any]) -> str:
    params = meta.get("params") if isinstance(meta.get("params"), Mapping) else {}
    if isinstance(params, Mapping):
        for key in ("label", "name", "title"):
            value = params.get(key)
            if isinstance(value, str) and value.strip():
                return value.strip()
        type_name = str(meta.get("type") or "")
        value = params.get("value")
        if type_name == "Text" and isinstance(value, str) and value.strip() and len(value.strip()) <= 40:
            return f"{value.strip()} text"
    return str(meta.get("type") or "")


def _descriptive_node_id(node_id: str) -> bool:
    clean = _identifier(node_id)
    if len(clean) <= 2:
        return False
    if re.fullmatch(r"(node_?)?\d+", clean):
        return False
    if re.fullmatch(r"[a-z]\d*", clean):
        return False
    return True


def _dedupe(base: str, used: set[str]) -> str:
    name = base
    suffix = 2
    while name in used:
        name = f"{base}_{suffix}"
        suffix += 1
    used.add(name)
    return name


def _identifier(value: str) -> str:
    name = re.sub(r"\W+", "_", value.strip()).strip("_").lower()
    if not name:
        return ""
    if name[0].isdigit():
        name = f"node_{name}"
    if keyword.iskeyword(name):
        name = f"{name}_node"
    return name


def _pretty_literal(value: Any) -> str:
    return pprint.pformat(value, width=100, sort_dicts=False)
