"""Convert PydanticAI Graphs to Areteos workflows.

Usage::

    from areteos import AreteosClient
    from areteos.pydantic_ai import from_pydantic_graph
    from my_project.pipeline import graph

    client = AreteosClient(base_url="...", api_key="...", engagement_id="...")

    workflow = from_pydantic_graph(
        graph,
        client=client,
        name="Research Pipeline",
        start_node="ResearchNode",
        node_module="my_project.pipeline",
    )

Requires the ``pydantic-ai`` optional dependency::

    pip install areteos[pydantic-ai]
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

from .exceptions import GraphConversionError
from .models import Edge, InputMapping, Step, StepType, Workflow

if TYPE_CHECKING:
    from .client import AreteosClient

try:
    from pydantic_graph import Graph as PydanticGraph  # noqa: F401
except ImportError as _err:
    raise ImportError(
        "pydantic-ai is required for PydanticAI Graph conversion. "
        "Install with: pip install areteos[pydantic-ai]"
    ) from _err


def _generate_node_script(
    node_class_name: str,
    node_module: str,
    requires: list[str] | None = None,
) -> str:
    """Generate an inline Python script that imports and runs a PydanticAI node.

    The generated script uses PEP 723 inline metadata so ``uv run`` can
    resolve dependencies automatically.
    """
    pep723_deps = ['"pydantic-ai>=1.0,<2.0"']
    for dep in requires or []:
        pep723_deps.append(f'"{dep}"')
    deps_lines = "\n".join(f"#     {d}," for d in pep723_deps)

    return (
        "# /// script\n"
        '# requires-python = ">=3.11"\n'
        "# dependencies = [\n"
        f"{deps_lines}\n"
        "# ]\n"
        "# ///\n"
        "import json, sys, asyncio, dataclasses\n"
        "from pydantic_graph import GraphRunContext, End\n"
        f"from {node_module} import {node_class_name}\n"
        "\n"
        "def _main():\n"
        "    if len(sys.argv) > 1:\n"
        "        with open(sys.argv[1]) as _f:\n"
        "            data = json.load(_f)\n"
        "    else:\n"
        "        data = {}\n"
        '    state = data.get("_state", {})\n'
        '    deps_data = data.get("_deps", {})\n'
        '    node_args = data.get("_node_args", {})\n'
        f"    node = {node_class_name}(**node_args)\n"
        "    ctx = GraphRunContext(state=state, deps=deps_data)\n"
        "    result = asyncio.run(node.run(ctx))\n"
        "    if isinstance(result, End):\n"
        '        output = {"_state": state, "_ended": True, "_result": result.data}\n'
        "    else:\n"
        "        output = {\n"
        '            "_state": state,\n'
        '            "_ended": False,\n'
        '            "_next_node": type(result).__name__,\n'
        '            "_node_args": dataclasses.asdict(result),\n'
        "        }\n"
        "    print(json.dumps(output, default=str))\n"
        "_main()\n"
    )


def _build_linear_chain(
    node_defs: dict[str, Any],
    start_node: str,
) -> list[str]:
    """Walk the graph from *start_node* following single successors.

    Returns an ordered list of node IDs forming the linear chain.
    Raises :class:`GraphConversionError` if a node has more than one successor
    or uses ``returns_base_node``.
    """
    if start_node not in node_defs:
        raise GraphConversionError(
            f"Start node '{start_node}' not found in graph. "
            f"Available nodes: {', '.join(sorted(node_defs.keys()))}"
        )

    chain: list[str] = []
    visited: set[str] = set()
    current: str | None = start_node

    while current is not None:
        if current in visited:
            raise GraphConversionError(
                f"Cycle detected: node '{current}' was already visited. "
                "Only acyclic linear graphs are supported."
            )
        visited.add(current)
        chain.append(current)

        ndef = node_defs[current]

        if ndef.returns_base_node:
            raise GraphConversionError(
                f"Node '{current}' uses returns_base_node=True (dynamic dispatch). "
                "This is not supported. Each node must return a specific next node type or End."
            )

        successors = list(ndef.next_node_edges.keys())

        if len(successors) > 1:
            raise GraphConversionError(
                f"Graph has branching at node '{current}' "
                f"(returns {' | '.join(successors)}). "
                "Conditional routing is not yet supported. "
                "Restructure as a linear chain, "
                "or use @python_step to wrap individual nodes manually."
            )

        # Follow the single successor, or stop at terminal nodes
        current = successors[0] if len(successors) == 1 else None

    return chain


def from_pydantic_graph(
    graph: PydanticGraph[Any, Any, Any],
    *,
    client: AreteosClient,
    name: str,
    start_node: str,
    node_module: str,
    description: str | None = None,
    requires: list[str] | None = None,
) -> Workflow:
    """Convert a PydanticAI Graph to an Areteos workflow.

    Inspects the graph's ``node_defs`` to extract nodes and edges, generates
    an inline Python script for each node (using import-by-path), and registers
    the workflow via *client*.

    Only **linear** graphs are supported (each node has at most one successor).
    Branching graphs and ``returns_base_node`` raise :class:`GraphConversionError`.

    Args:
        graph: A PydanticAI ``Graph`` instance.
        client: An :class:`AreteosClient` used to register the workflow.
        name: Display name for the created workflow.
        start_node: The class name of the first node to execute (e.g. ``"ResearchNode"``).
        node_module: Fully qualified Python module where node classes are importable
            (e.g. ``"my_project.pipeline"``).
        description: Optional workflow description.
        requires: Extra pip packages to include in PEP 723 script metadata
            (e.g. ``["my-project>=1.0"]``).

    Returns:
        The created :class:`Workflow`.

    Raises:
        GraphConversionError: If the graph is not a supported linear chain.
    """
    node_defs = graph.node_defs
    chain = _build_linear_chain(node_defs, start_node)

    # --- Build Steps ---
    steps: list[Step] = []
    for i, node_id in enumerate(chain):
        ndef = node_defs[node_id]
        node_class_name = ndef.node.__name__

        script = _generate_node_script(
            node_class_name=node_class_name,
            node_module=node_module,
            requires=requires,
        )

        # Build input mappings for state threading
        input_mappings: list[InputMapping] = []

        if i == 0:
            # Start node: map from run input
            input_mappings.extend(
                [
                    InputMapping(source_step_id=None, source_field="state", target_field="_state"),
                    InputMapping(source_step_id=None, source_field="deps", target_field="_deps"),
                    InputMapping(
                        source_step_id=None,
                        source_field="start_node_args",
                        target_field="_node_args",
                    ),
                ]
            )
        else:
            # Subsequent nodes: chain state and node args from predecessor
            predecessor_id = chain[i - 1]
            input_mappings.extend(
                [
                    InputMapping(
                        source_step_id=predecessor_id,
                        source_field="_state",
                        target_field="_state",
                    ),
                    InputMapping(source_step_id=None, source_field="deps", target_field="_deps"),
                    InputMapping(
                        source_step_id=predecessor_id,
                        source_field="_node_args",
                        target_field="_node_args",
                    ),
                ]
            )

        step = Step(
            id=node_id,
            type=StepType.TOOL,
            name=node_class_name,
            config={"tool": "python", "script_content": script},
            timeout_ms=300_000,
            max_retries=2,
            input_mappings=input_mappings,
        )
        steps.append(step)

    # --- Build Edges ---
    edges: list[Edge] = []
    for i in range(len(chain) - 1):
        edges.append(Edge(from_step_id=chain[i], to_step_id=chain[i + 1]))

    # --- Register Workflow ---
    return client.create_workflow(
        name,
        steps=steps,
        edges=edges,
        description=description,
    )
