"""Tests for PydanticAI Graph → Areteos Workflow converter."""

from __future__ import annotations

import importlib
import sys
from dataclasses import dataclass
from unittest.mock import MagicMock, patch

import pytest
from pydantic_graph import BaseNode, End, Graph, GraphRunContext
from pytest_httpx import HTTPXMock

from areteos import AreteosClient, GraphConversionError, Workflow
from areteos.pydantic_ai import _build_linear_chain, _generate_node_script, from_pydantic_graph

# --- Test constants ---

BASE_URL = "https://areteos.test"
API_KEY = "areteos_test_key"
ENGAGEMENT_ID = "550e8400-0000-0000-0000-000000000001"
WORKFLOW_ID = "660e8400-0000-0000-0000-000000000001"
VERSION_ID = "770e8400-0000-0000-0000-000000000001"

NODE_MODULE = "tests.test_pydantic_ai"


def _workflow_response(workflow_id: str = WORKFLOW_ID, version_id: str = VERSION_ID) -> dict:
    return {
        "data": {
            "type": "workflows",
            "id": workflow_id,
            "attributes": {
                "name": "Test Workflow",
                "description": None,
                "scope_type": "engagement",
                "current_version_id": version_id,
                "engagement_id": ENGAGEMENT_ID,
            },
        }
    }


# --- Test PydanticAI Node classes (defined at module level for import-by-path) ---


@dataclass
class NodeA(BaseNode[dict, None, str]):
    value: str = "hello"

    async def run(self, ctx: GraphRunContext[dict, None]) -> NodeB:
        ctx.state["a_ran"] = True
        return NodeB()


@dataclass
class NodeB(BaseNode[dict, None, str]):
    async def run(self, ctx: GraphRunContext[dict, None]) -> End[str]:
        ctx.state["b_ran"] = True
        return End("done")


@dataclass
class NodeC(BaseNode[dict, None, str]):
    async def run(self, ctx: GraphRunContext[dict, None]) -> End[str]:
        return End("c_done")


@dataclass
class Chain1(BaseNode[dict, None, str]):
    async def run(self, ctx: GraphRunContext[dict, None]) -> Chain2:
        return Chain2()


@dataclass
class Chain2(BaseNode[dict, None, str]):
    async def run(self, ctx: GraphRunContext[dict, None]) -> Chain3:
        return Chain3()


@dataclass
class Chain3(BaseNode[dict, None, str]):
    async def run(self, ctx: GraphRunContext[dict, None]) -> End[str]:
        return End("final")


@dataclass
class BranchingNode(BaseNode[dict, None, str]):
    async def run(self, ctx: GraphRunContext[dict, None]) -> NodeB | NodeC:
        return NodeB()


# --- Tests ---


class TestLinearGraphProducesStepsAndEdges:
    """areteos-0xmi.5: test_linear_graph_produces_steps_and_edges"""

    def test_two_node_graph(self, httpx_mock: HTTPXMock) -> None:
        httpx_mock.add_response(json=_workflow_response())
        graph = Graph(nodes=[NodeA, NodeB])

        with AreteosClient(
            base_url=BASE_URL, api_key=API_KEY, engagement_id=ENGAGEMENT_ID
        ) as client:
            workflow = from_pydantic_graph(
                graph,
                client=client,
                name="Test",
                start_node="NodeA",
                node_module=NODE_MODULE,
            )

        assert isinstance(workflow, Workflow)
        assert str(workflow.id) == WORKFLOW_ID

        # Verify the request sent to the API
        request = httpx_mock.get_requests()[0]
        import json

        body = json.loads(request.read())
        attrs = body["data"]["attributes"]

        steps = attrs["steps"]
        edges = attrs["edges"]

        assert len(steps) == 2
        assert steps[0]["id"] == "NodeA"
        assert steps[1]["id"] == "NodeB"

        assert len(edges) == 1
        assert edges[0]["from_step_id"] == "NodeA"
        assert edges[0]["to_step_id"] == "NodeB"


class TestThreeNodeChainStateThreading:
    """areteos-0xmi.5: test_three_node_chain_state_threading"""

    def test_three_node_chain(self, httpx_mock: HTTPXMock) -> None:
        graph = Graph(nodes=[Chain1, Chain2, Chain3])

        httpx_mock.add_response(json=_workflow_response())

        with AreteosClient(
            base_url=BASE_URL, api_key=API_KEY, engagement_id=ENGAGEMENT_ID
        ) as client:
            from_pydantic_graph(
                graph,
                client=client,
                name="Three Node Chain",
                start_node="Chain1",
                node_module=NODE_MODULE,
            )

        import json

        request = httpx_mock.get_requests()[0]
        body = json.loads(request.read())
        attrs = body["data"]["attributes"]
        steps = attrs["steps"]
        edges = attrs["edges"]

        # 3 steps, 2 edges
        assert len(steps) == 3
        assert len(edges) == 2

        # Chain1 gets input from run input (source_step_id=None)
        step1_mappings = steps[0]["input_mappings"]
        assert any(
            m["source_field"] == "state"
            and m["target_field"] == "_state"
            and m.get("source_step_id") is None
            for m in step1_mappings
        )

        # Chain2 gets _state from Chain1 (source_step_id="Chain1")
        step2_mappings = steps[1]["input_mappings"]
        assert any(
            m["source_step_id"] == "Chain1"
            and m["source_field"] == "_state"
            and m["target_field"] == "_state"
            for m in step2_mappings
        )

        # Chain3 gets _state from Chain2
        step3_mappings = steps[2]["input_mappings"]
        assert any(
            m["source_step_id"] == "Chain2"
            and m["source_field"] == "_state"
            and m["target_field"] == "_state"
            for m in step3_mappings
        )


class TestEndNodeProducesResult:
    """areteos-0xmi.5: test_end_node_produces_result"""

    def test_terminal_node_no_outgoing_edge(self, httpx_mock: HTTPXMock) -> None:
        graph = Graph(nodes=[NodeA, NodeB])
        httpx_mock.add_response(json=_workflow_response())

        with AreteosClient(
            base_url=BASE_URL, api_key=API_KEY, engagement_id=ENGAGEMENT_ID
        ) as client:
            from_pydantic_graph(
                graph,
                client=client,
                name="Test",
                start_node="NodeA",
                node_module=NODE_MODULE,
            )

        import json

        request = httpx_mock.get_requests()[0]
        body = json.loads(request.read())
        edges = body["data"]["attributes"]["edges"]

        # NodeB is terminal — only 1 edge (NodeA -> NodeB), no edge from NodeB
        assert len(edges) == 1
        assert edges[0]["from_step_id"] == "NodeA"
        assert edges[0]["to_step_id"] == "NodeB"


class TestBranchingGraphRejected:
    """areteos-0xmi.5: test_branching_graph_rejected"""

    def test_branching_raises_error(self) -> None:
        graph = Graph(nodes=[BranchingNode, NodeB, NodeC])

        with (
            AreteosClient(
                base_url=BASE_URL, api_key=API_KEY, engagement_id=ENGAGEMENT_ID
            ) as client,
            pytest.raises(GraphConversionError, match="branching"),
        ):
            from_pydantic_graph(
                graph,
                client=client,
                name="Branching",
                start_node="BranchingNode",
                node_module=NODE_MODULE,
            )


class TestReturnsBaseNodeRejected:
    """areteos-0xmi.5: test_returns_base_node_rejected"""

    def test_returns_base_node_raises_error(self) -> None:
        graph = Graph(nodes=[NodeA, NodeB])

        # Simulate returns_base_node by patching the node_def
        graph.node_defs["NodeA"].returns_base_node = True

        with (
            AreteosClient(
                base_url=BASE_URL, api_key=API_KEY, engagement_id=ENGAGEMENT_ID
            ) as client,
            pytest.raises(GraphConversionError, match="returns_base_node"),
        ):
            from_pydantic_graph(
                graph,
                client=client,
                name="Dynamic",
                start_node="NodeA",
                node_module=NODE_MODULE,
            )


class TestGeneratedScriptImportsByPath:
    """areteos-0xmi.5: test_generated_script_imports_by_path"""

    def test_script_contains_import_statement(self) -> None:
        script = _generate_node_script(
            node_class_name="ResearchNode",
            node_module="my_project.pipeline",
        )
        assert "from my_project.pipeline import ResearchNode" in script

    def test_script_instantiates_node(self) -> None:
        script = _generate_node_script(
            node_class_name="ResearchNode",
            node_module="my_project.pipeline",
        )
        assert "node = ResearchNode(**node_args)" in script


class TestPep723IncludesPydanticAi:
    """areteos-0xmi.5: test_pep723_includes_pydantic_ai"""

    def test_pep723_header_present(self) -> None:
        script = _generate_node_script(
            node_class_name="MyNode",
            node_module="my_module",
        )
        assert "# /// script" in script
        assert '"pydantic-ai>=1.0,<2.0"' in script
        assert '# requires-python = ">=3.11"' in script

    def test_pep723_includes_user_deps(self) -> None:
        script = _generate_node_script(
            node_class_name="MyNode",
            node_module="my_module",
            requires=["my-package>=1.0", "numpy"],
        )
        assert '"my-package>=1.0"' in script
        assert '"numpy"' in script
        assert '"pydantic-ai>=1.0,<2.0"' in script


class TestMissingPydanticAiImportError:
    """areteos-0xmi.5: test_missing_pydantic_ai_import_error"""

    def test_import_error_without_pydantic_ai(self) -> None:
        # Remove cached module so re-import triggers the top-level guard
        saved_modules = {}
        modules_to_remove = [
            k for k in sys.modules if k == "areteos.pydantic_ai" or k.startswith("pydantic_graph")
        ]
        for mod_name in modules_to_remove:
            saved_modules[mod_name] = sys.modules.pop(mod_name)

        try:
            # Patch pydantic_graph to None so the import raises ImportError
            with (
                patch.dict("sys.modules", {"pydantic_graph": None}),
                pytest.raises(ImportError, match="pip install areteos\\[pydantic-ai\\]"),
            ):
                importlib.import_module("areteos.pydantic_ai")
        finally:
            # Restore original modules so other tests are unaffected
            sys.modules.update(saved_modules)


class TestBuildLinearChain:
    """Unit tests for the _build_linear_chain helper."""

    def test_simple_chain(self) -> None:
        graph = Graph(nodes=[NodeA, NodeB])
        chain = _build_linear_chain(graph.node_defs, "NodeA")
        assert chain == ["NodeA", "NodeB"]

    def test_start_node_not_found(self) -> None:
        graph = Graph(nodes=[NodeA, NodeB])
        with pytest.raises(GraphConversionError, match="not found"):
            _build_linear_chain(graph.node_defs, "NonExistent")

    def test_branching_rejected(self) -> None:
        graph = Graph(nodes=[BranchingNode, NodeB, NodeC])
        with pytest.raises(GraphConversionError, match="branching"):
            _build_linear_chain(graph.node_defs, "BranchingNode")


class TestBuildLinearChainCycleDetection:
    """areteos-geem: Cycle detection in _build_linear_chain."""

    def test_cycle_raises_error(self) -> None:
        """Create node_defs with a cycle (A->B->A) and verify detection."""
        graph = Graph(nodes=[NodeA, NodeB])

        # Patch NodeB's next_node_edges to point back to NodeA (creating a cycle)
        # instead of ending. We mock at the node_defs level since PydanticAI
        # validates return types at Graph.__init__ time.
        original_edges = graph.node_defs["NodeB"].next_node_edges
        graph.node_defs["NodeB"].next_node_edges = {"NodeA": MagicMock()}

        try:
            with pytest.raises(GraphConversionError, match="Cycle detected"):
                _build_linear_chain(graph.node_defs, "NodeA")
        finally:
            graph.node_defs["NodeB"].next_node_edges = original_edges


class TestNodeArgsThreading:
    """areteos-3doz: Verify _node_args are serialized and threaded between steps."""

    def test_generated_script_serializes_node_args(self) -> None:
        """The non-End branch must include dataclasses.asdict(result) as _node_args."""
        script = _generate_node_script(
            node_class_name="MyNode",
            node_module="my_module",
        )
        assert "import json, sys, asyncio, dataclasses" in script
        assert '"_node_args": dataclasses.asdict(result)' in script

    def test_node_args_mapped_from_predecessor(self, httpx_mock: HTTPXMock) -> None:
        """Non-start nodes must have an InputMapping for _node_args from predecessor."""
        graph = Graph(nodes=[NodeA, NodeB])
        httpx_mock.add_response(json=_workflow_response())

        with AreteosClient(
            base_url=BASE_URL, api_key=API_KEY, engagement_id=ENGAGEMENT_ID
        ) as client:
            from_pydantic_graph(
                graph,
                client=client,
                name="Test",
                start_node="NodeA",
                node_module=NODE_MODULE,
            )

        import json

        request = httpx_mock.get_requests()[0]
        body = json.loads(request.read())
        steps = body["data"]["attributes"]["steps"]

        # NodeB (second step) should have _node_args mapped from NodeA
        step_b_mappings = steps[1]["input_mappings"]
        assert any(
            m.get("source_step_id") == "NodeA"
            and m["source_field"] == "_node_args"
            and m["target_field"] == "_node_args"
            for m in step_b_mappings
        ), f"Expected _node_args mapping from NodeA in NodeB's input_mappings: {step_b_mappings}"

    def test_node_args_data_preserved_in_chain(self, httpx_mock: HTTPXMock) -> None:
        """In a three-node chain, each successor gets _node_args from its predecessor."""
        graph = Graph(nodes=[Chain1, Chain2, Chain3])
        httpx_mock.add_response(json=_workflow_response())

        with AreteosClient(
            base_url=BASE_URL, api_key=API_KEY, engagement_id=ENGAGEMENT_ID
        ) as client:
            from_pydantic_graph(
                graph,
                client=client,
                name="Chain Test",
                start_node="Chain1",
                node_module=NODE_MODULE,
            )

        import json

        request = httpx_mock.get_requests()[0]
        body = json.loads(request.read())
        steps = body["data"]["attributes"]["steps"]

        # Chain1 (start): gets start_node_args from run input
        step1_mappings = steps[0]["input_mappings"]
        assert any(
            m.get("source_step_id") is None
            and m["source_field"] == "start_node_args"
            and m["target_field"] == "_node_args"
            for m in step1_mappings
        )

        # Chain2: gets _node_args from Chain1
        step2_mappings = steps[1]["input_mappings"]
        assert any(
            m.get("source_step_id") == "Chain1"
            and m["source_field"] == "_node_args"
            and m["target_field"] == "_node_args"
            for m in step2_mappings
        )

        # Chain3: gets _node_args from Chain2
        step3_mappings = steps[2]["input_mappings"]
        assert any(
            m.get("source_step_id") == "Chain2"
            and m["source_field"] == "_node_args"
            and m["target_field"] == "_node_args"
            for m in step3_mappings
        )


class TestGeneratedScriptFileHandling:
    """areteos-ksv3: Generated script uses context manager for file I/O."""

    def test_script_uses_context_manager(self) -> None:
        script = _generate_node_script(
            node_class_name="MyNode",
            node_module="my_module",
        )
        assert "with open(sys.argv[1]) as _f:" in script
        assert "json.load(_f)" in script
        # Ensure the old unsafe pattern is NOT present
        assert "json.load(open(" not in script


class TestFromPydanticGraphStepConfig:
    """Verify step config details (timeout, type, etc.)."""

    def test_steps_are_tool_type_with_python(self, httpx_mock: HTTPXMock) -> None:
        graph = Graph(nodes=[NodeA, NodeB])
        httpx_mock.add_response(json=_workflow_response())

        with AreteosClient(
            base_url=BASE_URL, api_key=API_KEY, engagement_id=ENGAGEMENT_ID
        ) as client:
            from_pydantic_graph(
                graph,
                client=client,
                name="Test",
                start_node="NodeA",
                node_module=NODE_MODULE,
            )

        import json

        request = httpx_mock.get_requests()[0]
        body = json.loads(request.read())
        steps = body["data"]["attributes"]["steps"]

        for step in steps:
            assert step["type"] == "tool"
            assert step["config"]["tool"] == "python"
            assert "script_content" in step["config"]
            assert step["timeout_ms"] == 300_000

    def test_description_passed_through(self, httpx_mock: HTTPXMock) -> None:
        graph = Graph(nodes=[NodeA, NodeB])
        httpx_mock.add_response(json=_workflow_response())

        with AreteosClient(
            base_url=BASE_URL, api_key=API_KEY, engagement_id=ENGAGEMENT_ID
        ) as client:
            from_pydantic_graph(
                graph,
                client=client,
                name="Test",
                start_node="NodeA",
                node_module=NODE_MODULE,
                description="My description",
            )

        import json

        request = httpx_mock.get_requests()[0]
        body = json.loads(request.read())
        assert body["data"]["attributes"]["description"] == "My description"
