"""Tests for Pipeline system"""

import pytest

from agentkit.orchestrator.pipeline_schema import Pipeline, PipelineStage
from agentkit.orchestrator.pipeline_engine import PipelineEngine
from agentkit.orchestrator.pipeline_loader import PipelineLoader


def test_pipeline_schema():
    stage = PipelineStage(
        name="step1",
        agent="agent_a",
        action="process",
        depends_on=[],
        inputs={"data": "${input}"},
    )
    pipeline = Pipeline(
        name="test_pipeline",
        version="1.0.0",
        description="Test",
        stages=[stage],
        variables={"input": "hello"},
    )
    assert len(pipeline.stages) == 1
    assert pipeline.variables["input"] == "hello"


def test_topological_group():
    stages = [
        PipelineStage(name="a", agent="agent1", action="do_a", depends_on=[]),
        PipelineStage(name="b", agent="agent2", action="do_b", depends_on=["a"]),
        PipelineStage(name="c", agent="agent3", action="do_c", depends_on=["a"]),
        PipelineStage(name="d", agent="agent4", action="do_d", depends_on=["b", "c"]),
    ]

    groups = PipelineEngine._topological_group(stages)
    assert len(groups) == 3  # [a], [b, c], [d]
    assert groups[0][0].name == "a"
    assert {s.name for s in groups[1]} == {"b", "c"}
    assert groups[2][0].name == "d"


def test_topological_group_circular():
    stages = [
        PipelineStage(name="a", agent="agent1", action="do_a", depends_on=["b"]),
        PipelineStage(name="b", agent="agent2", action="do_b", depends_on=["a"]),
    ]

    with pytest.raises(ValueError, match="Circular dependency"):
        PipelineEngine._topological_group(stages)


@pytest.mark.asyncio
async def test_pipeline_dry_run():
    pipeline = Pipeline(
        name="test",
        version="1.0.0",
        description="Test",
        stages=[
            PipelineStage(name="step1", agent="agent1", action="process", inputs={"x": "1"}),
            PipelineStage(name="step2", agent="agent2", action="analyze", depends_on=["step1"], inputs={"y": "2"}),
        ],
    )

    engine = PipelineEngine(dispatcher=None)  # dry-run mode
    result = await engine.execute(pipeline)
    assert result.status.value == "completed"


def test_yaml_loader():
    yaml_content = """
name: test_pipeline
version: "1.0"
description: Test pipeline
stages:
  - name: step1
    agent: agent1
    action: process
    inputs:
      data: hello
  - name: step2
    agent: agent2
    action: analyze
    depends_on: [step1]
    inputs:
      result: ${{step1_output}}
variables:
  input: world
"""
    loader = PipelineLoader()
    pipeline = loader.load_from_yaml(yaml_content, "test")
    assert pipeline.name == "test_pipeline"
    assert len(pipeline.stages) == 2
    assert pipeline.stages[1].depends_on == ["step1"]


def test_validate_dag():
    stages = [
        PipelineStage(name="a", agent="a1", action="do", depends_on=[]),
        PipelineStage(name="b", agent="a2", action="do", depends_on=["a"]),
    ]
    assert PipelineLoader.validate_dag(stages) is True

    circular = [
        PipelineStage(name="a", agent="a1", action="do", depends_on=["b"]),
        PipelineStage(name="b", agent="a2", action="do", depends_on=["a"]),
    ]
    assert PipelineLoader.validate_dag(circular) is False