from __future__ import annotations

import json
import keyword
import re
from pprint import pformat
from typing import Any, Mapping

from ..workflow import WorkflowRunError, infer_entrypoint, validate_workflow


def export_framework_scaffold(
    data: Mapping[str, Any],
    target_id: str,
    label: str,
) -> 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)
    projection = _workflow_projection(data, {"node_id": entry_node, "port": entry_port})
    builder_name = {
        "crewai": "build_crewai_tasks",
        "autogen": "build_autogen_agents",
        "swarm": "build_swarm_handoffs",
        "nvidia-agent-stack": "build_nvidia_agent_stack_manifest",
    }[target_id]
    body = {
        "crewai": _crewai_builder(),
        "autogen": _autogen_builder(),
        "swarm": _swarm_builder(),
        "nvidia-agent-stack": _nvidia_agent_stack_builder(),
    }[target_id]
    workflow_name = str(projection["name"])
    entry_name = str(projection.get("entrypoint_node_name") or projection["entrypoint"]["node_id"])

    code = "\n".join([
        "from __future__ import annotations",
        "",
        f"# Generated by Blacknode from workflow: {workflow_name}",
        f"# Target: {label}",
        f"# Entrypoint: {entry_name}.{projection['entrypoint']['port']}",
        "# Keep the visual graph as the source of truth, then regenerate this file when the graph changes.",
        "",
        "from pprint import pprint",
        "",
        "",
        f"WORKFLOW = {pformat(projection, width=100, sort_dicts=False)}",
        "",
        "",
        body,
        "",
        "",
        "if __name__ == '__main__':",
        f"    pprint({builder_name}())",
        "",
    ])
    return code, [f"{label} export carries graph structure; add framework-specific model clients and policies in the generated file."]


def _workflow_projection(data: Mapping[str, Any], entrypoint: dict[str, str]) -> dict[str, Any]:
    node_meta = data.get("node_meta") or {}
    nodes = {
        str(node_id): dict(meta)
        for node_id, meta in node_meta.items()
        if isinstance(meta, Mapping)
    }
    node_names = _unique_node_names(nodes)
    raw_edges = [dict(edge) for edge in data.get("edges") or []]
    incoming = _incoming_edges(raw_edges)
    outgoing = _outgoing_edges(raw_edges)
    return {
        "name": data.get("name") or "Blacknode Workflow",
        "entrypoint": entrypoint,
        "entrypoint_node_name": node_names.get(entrypoint["node_id"], entrypoint["node_id"]),
        "mcp": {
            "stdio": "blacknode mcp",
            "streamable_http": "blacknode mcp --transport streamable-http --host 127.0.0.1 --port 9901 --path /mcp",
        },
        "nvidia": {
            "hosted_nim_env": "NVIDIA_API_KEY",
            "local_nim_env": "NGC_API_KEY",
            "node_names": [
                node_names[node_id]
                for node_id, meta in nodes.items()
                if _is_nvidia_node(meta)
            ],
            "model_refs": _model_refs(nodes),
        },
        "nodes": [
            {
                "step": index,
                "id": str(node_id),
                "name": node_names[str(node_id)],
                "label": _node_comment_label(str(node_id), meta),
                "type": str(meta.get("type")),
                "params": dict(meta.get("params") or {}),
                "inputs": list(meta.get("inputs") or []),
                "outputs": list(meta.get("outputs") or []),
                "upstream": _edge_refs(incoming.get(str(node_id), []), node_names, direction="upstream"),
                "downstream": _edge_refs(outgoing.get(str(node_id), []), node_names, direction="downstream"),
            }
            for index, (node_id, meta) in enumerate(nodes.items(), start=1)
        ],
        "edges": [
            {
                **edge,
                "from_name": node_names.get(str(edge.get("from")), str(edge.get("from"))),
                "to_name": node_names.get(str(edge.get("to")), str(edge.get("to"))),
            }
            for edge in raw_edges
        ],
    }


def _crewai_builder() -> str:
    return '''def build_crewai_tasks() -> list[dict]:
    """Return CrewAI task descriptors mapped from readable Blacknode steps."""
    tasks = []
    for node in WORKFLOW["nodes"]:
        upstream = [edge["name"] for edge in node["upstream"]]
        tasks.append({
            "name": node["name"],
            "description": f"Step {node['step']}: {node['label']}",
            "expected_output": ", ".join(node["outputs"]) or "node result",
            "blacknode_id": node["id"],
            "blacknode_type": node["type"],
            "params": node["params"],
            "context": upstream,
        })
    return tasks'''


def _autogen_builder() -> str:
    return '''def build_autogen_agents() -> list[dict]:
    """Return AutoGen agent descriptors mapped from readable Blacknode steps."""
    agents = []
    for node in WORKFLOW["nodes"]:
        outbound = [edge["name"] for edge in node["downstream"]]
        agents.append({
            "name": node["name"],
            "system_message": f"Execute Step {node['step']}: {node['label']}. Preserve typed inputs and outputs from the Blacknode graph.",
            "blacknode_id": node["id"],
            "blacknode_type": node["type"],
            "params": node["params"],
            "handoffs": outbound,
        })
    return agents'''


def _swarm_builder() -> str:
    return '''def build_swarm_handoffs() -> list[dict]:
    """Return OpenAI Swarm-style handoff descriptors mapped from readable Blacknode steps."""
    handoffs = []
    for node in WORKFLOW["nodes"]:
        targets = [edge["name"] for edge in node["downstream"]]
        handoffs.append({
            "agent": node["name"],
            "instructions": f"Handle Step {node['step']}: {node['label']}. Use the original Blacknode id for traceability.",
            "blacknode_id": node["id"],
            "blacknode_type": node["type"],
            "params": node["params"],
            "handoff_targets": targets,
        })
    return handoffs'''


def _nvidia_agent_stack_builder() -> str:
    return '''def build_nvidia_agent_stack_manifest() -> dict:
    """Return an NVIDIA-facing integration manifest for AI-Q, NeMo Agent Toolkit, and NIM demos."""
    return {
        "workflow": WORKFLOW["name"],
        "entrypoint": {
            "node": WORKFLOW["entrypoint_node_name"],
            "port": WORKFLOW["entrypoint"]["port"],
        },
        "positioning": "Blacknode is the visual workflow editor for the NVIDIA agent stack.",
        "mcp": WORKFLOW["mcp"],
        "nvidia": WORKFLOW["nvidia"],
        "steps": [
            {
                "step": node["step"],
                "name": node["name"],
                "label": node["label"],
                "blacknode_id": node["id"],
                "blacknode_type": node["type"],
                "upstream": [edge["name"] for edge in node["upstream"]],
                "downstream": [edge["name"] for edge in node["downstream"]],
            }
            for node in WORKFLOW["nodes"]
        ],
        "notes": [
            "Use the MCP streamable HTTP command when connecting from AI-Q or NeMo Agent Toolkit.",
            "Use hosted NIM with NVIDIA_API_KEY for the quickest demo path.",
            "Use local NIM launch nodes or Docker commands when demonstrating on-prem deployment.",
            "This export is an integration manifest; the visual Blacknode workflow remains the executable source.",
        ],
    }'''


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


def _outgoing_edges(edges: list[dict[str, Any]]) -> dict[str, list[dict[str, Any]]]:
    outgoing: dict[str, list[dict[str, Any]]] = {}
    for edge in edges:
        outgoing.setdefault(str(edge.get("from")), []).append(edge)
    return outgoing


def _edge_refs(
    edges: list[dict[str, Any]],
    node_names: Mapping[str, str],
    *,
    direction: str,
) -> list[dict[str, str]]:
    refs: list[dict[str, str]] = []
    for edge in edges:
        if direction == "upstream":
            node_id = str(edge.get("from"))
            port = str(edge.get("from_port", "output"))
            local_port = str(edge.get("to_port", "input"))
        else:
            node_id = str(edge.get("to"))
            port = str(edge.get("to_port", "input"))
            local_port = str(edge.get("from_port", "output"))
        refs.append({
            "id": node_id,
            "name": node_names.get(node_id, node_id),
            "port": port,
            "local_port": local_port,
        })
    return refs


def _is_nvidia_node(meta: Mapping[str, Any]) -> bool:
    type_name = str(meta.get("type") or "")
    return (
        type_name.startswith("NVIDIA")
        or type_name.startswith("NIM")
        or type_name in {"VideoFolderInput"}
        or any(str(value).lower().startswith("nim:") for value in (meta.get("params") or {}).values())
    )


def _model_refs(nodes: Mapping[str, Mapping[str, Any]]) -> list[str]:
    refs: list[str] = []
    for meta in nodes.values():
        params = meta.get("params") if isinstance(meta.get("params"), Mapping) else {}
        for key in ("model", "embedding_model", "rerank_model"):
            value = params.get(key) if isinstance(params, Mapping) else None
            if isinstance(value, str) and value and value not in refs:
                refs.append(value)
    return refs


def _unique_node_names(node_meta: Mapping[str, Mapping[str, Any]]) -> dict[str, str]:
    used: set[str] = {"workflow", "tasks", "agents", "handoffs", "node", "edge"}
    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
