"""Tests for GEOPipeline"""

import pytest

from agentkit.skills.geo_pipeline import (
    GEOPipeline,
    PipelineStep,
    PipelineStepResult,
    PipelineResult,
)


class MockAgent:
    """Mock Agent for pipeline testing"""

    def __init__(self, name: str, output_data: dict | None = None):
        self.name = name
        self.agent_type = "mock"
        self._output_data = output_data or {"result": f"output from {name}"}

    async def execute(self, task):
        from agentkit.core.protocol import TaskResult, TaskStatus
        from datetime import datetime, timezone
        now = datetime.now(timezone.utc)
        return TaskResult(
            task_id=task.task_id,
            agent_name=self.name,
            status=TaskStatus.COMPLETED,
            output_data=self._output_data,
            error_message=None,
            started_at=now,
            completed_at=now,
        )


class MockAgentPool:
    """Mock AgentPool"""

    def __init__(self, agents: dict[str, MockAgent] | None = None):
        self._agents = agents or {}

    def get_agent(self, name: str):
        return self._agents.get(name)

    def list_agents(self):
        return [{"name": a.name, "agent_type": a.agent_type} for a in self._agents.values()]


class TestGEOPipeline:
    """GEOPipeline unit tests"""

    @pytest.mark.asyncio
    async def test_sequential_pipeline(self):
        """Sequential steps should execute in order"""
        steps = [
            PipelineStep(name="step1", skill="skill_a"),
            PipelineStep(name="step2", skill="skill_b", depends_on=["step1"]),
        ]
        pool = MockAgentPool({
            "skill_a": MockAgent("skill_a", {"data": "result_a"}),
            "skill_b": MockAgent("skill_b", {"data": "result_b"}),
        })
        pipeline = GEOPipeline(name="test", steps=steps, agent_pool=pool)

        result = await pipeline.execute({"query": "test"})
        assert result.success
        assert len(result.steps) == 2
        assert result.steps[0].status == "success"
        assert result.steps[1].status == "success"

    @pytest.mark.asyncio
    async def test_parallel_steps(self):
        """Steps without dependencies should execute in parallel"""
        steps = [
            PipelineStep(name="step1", skill="skill_a"),
            PipelineStep(name="step2", skill="skill_b"),
        ]
        pool = MockAgentPool({
            "skill_a": MockAgent("skill_a", {"data": "a"}),
            "skill_b": MockAgent("skill_b", {"data": "b"}),
        })
        pipeline = GEOPipeline(name="test", steps=steps, agent_pool=pool)

        result = await pipeline.execute({"query": "test"})
        assert result.success
        assert len(result.steps) == 2

    @pytest.mark.asyncio
    async def test_dag_execution(self):
        """DAG with mixed parallel/sequential steps"""
        steps = [
            PipelineStep(name="detect", skill="skill_a"),
            PipelineStep(name="analyze_1", skill="skill_b", depends_on=["detect"]),
            PipelineStep(name="analyze_2", skill="skill_c", depends_on=["detect"]),
            PipelineStep(name="optimize", skill="skill_d", depends_on=["analyze_1", "analyze_2"]),
        ]
        pool = MockAgentPool({
            "skill_a": MockAgent("skill_a", {"citations": 5}),
            "skill_b": MockAgent("skill_b", {"competitor": "data"}),
            "skill_c": MockAgent("skill_c", {"trend": "up"}),
            "skill_d": MockAgent("skill_d", {"optimized": True}),
        })
        pipeline = GEOPipeline(name="test", steps=steps, agent_pool=pool)

        result = await pipeline.execute({"brand": "TestBrand"})
        assert result.success
        assert len(result.steps) == 4

        # Check execution groups
        groups = pipeline._build_execution_groups()
        assert len(groups) == 3  # [detect], [analyze_1, analyze_2], [optimize]
        assert "detect" in groups[0]
        assert set(groups[1]) == {"analyze_1", "analyze_2"}
        assert groups[2] == ["optimize"]

    @pytest.mark.asyncio
    async def test_step_failure(self):
        """Failed step should be recorded"""
        class FailingAgent:
            name = "skill_a"
            agent_type = "mock"
            async def execute(self, task):
                raise RuntimeError("Agent failed")

        steps = [PipelineStep(name="step1", skill="skill_a")]
        pool = MockAgentPool({"skill_a": FailingAgent()})
        pipeline = GEOPipeline(name="test", steps=steps, agent_pool=pool)

        result = await pipeline.execute({"query": "test"})
        assert not result.success
        assert result.steps[0].status == "failed"
        assert "Agent failed" in result.steps[0].error

    @pytest.mark.asyncio
    async def test_input_mapping(self):
        """Input mapping should resolve paths correctly"""
        steps = [
            PipelineStep(name="step1", skill="skill_a"),
            PipelineStep(
                name="step2",
                skill="skill_b",
                input_mapping={"brand": "$.input.brand"},
                depends_on=["step1"],
            ),
        ]
        pool = MockAgentPool({
            "skill_a": MockAgent("skill_a", {"data": "a"}),
            "skill_b": MockAgent("skill_b", {"data": "b"}),
        })
        pipeline = GEOPipeline(name="test", steps=steps, agent_pool=pool)

        result = await pipeline.execute({"brand": "TestBrand"})
        assert result.success

    @pytest.mark.asyncio
    async def test_from_config(self):
        """Pipeline should be created from YAML config"""
        config = {
            "name": "geo_test",
            "steps": [
                {"name": "detect", "skill": "citation_detector"},
                {"name": "analyze", "skill": "competitor_analyzer", "depends_on": ["detect"]},
            ],
        }
        pipeline = GEOPipeline.from_config(config)
        assert pipeline.name == "geo_test"
        assert len(pipeline._steps) == 2
        assert pipeline._steps[1].depends_on == ["detect"]

    @pytest.mark.asyncio
    async def test_execution_groups_topological_sort(self):
        """Execution groups should follow topological order"""
        steps = [
            PipelineStep(name="a", skill="s1"),
            PipelineStep(name="b", skill="s2", depends_on=["a"]),
            PipelineStep(name="c", skill="s3", depends_on=["a"]),
            PipelineStep(name="d", skill="s4", depends_on=["b", "c"]),
        ]
        pipeline = GEOPipeline(name="test", steps=steps)

        groups = pipeline._build_execution_groups()
        assert len(groups) == 3
        assert groups[0] == ["a"]
        assert set(groups[1]) == {"b", "c"}
        assert groups[2] == ["d"]

    @pytest.mark.asyncio
    async def test_resolve_mapping_path(self):
        """Mapping path resolution"""
        input_data = {"brand": "TestBrand", "platforms": ["chatgpt"]}
        step_outputs = {
            "detect": {"citations": 5, "records": []},
        }

        # $.input.brand
        result = GEOPipeline._resolve_mapping_path("$.input.brand", input_data, step_outputs)
        assert result == "TestBrand"

        # $.steps.detect.output.citations
        result = GEOPipeline._resolve_mapping_path("$.steps.detect.output.citations", input_data, step_outputs)
        assert result == 5

        # $.steps.detect (whole output)
        result = GEOPipeline._resolve_mapping_path("$.steps.detect", input_data, step_outputs)
        assert result == {"citations": 5, "records": []}

    @pytest.mark.asyncio
    async def test_final_output_includes_all_steps(self):
        """Final output should include all step results"""
        steps = [
            PipelineStep(name="step1", skill="skill_a"),
            PipelineStep(name="step2", skill="skill_b", depends_on=["step1"]),
        ]
        pool = MockAgentPool({
            "skill_a": MockAgent("skill_a", {"result": "a"}),
            "skill_b": MockAgent("skill_b", {"result": "b"}),
        })
        pipeline = GEOPipeline(name="test", steps=steps, agent_pool=pool)

        result = await pipeline.execute({"query": "test"})
        assert "step1" in result.final_output
        assert "step2" in result.final_output
        assert "input" in result.final_output

    @pytest.mark.asyncio
    async def test_empty_pipeline(self):
        """Empty pipeline should succeed with no steps"""
        pipeline = GEOPipeline(name="empty", steps=[])
        result = await pipeline.execute({"query": "test"})
        assert result.success
        assert len(result.steps) == 0