"""Tests for Pydantic models and step constructors."""

from areteos import (
    AgentStrategy,
    ApprovalRequestType,
    BranchCondition,
    Edge,
    InputMapping,
    ModelAlias,
    Run,
    RunStatus,
    Step,
    StepType,
)


class TestStepConstructors:
    def test_agent_step(self):
        step = Step.agent(
            "analyze",
            "Analyze Document",
            model=ModelAlias.CAPABLE,
            strategy=AgentStrategy.REACT,
            system_prompt="You are an analyst.",
        )
        assert step.id == "analyze"
        assert step.type == StepType.AGENT
        assert step.config["model"] == "capable"
        assert step.config["strategy"] == "react"
        assert step.config["system_prompt"] == "You are an analyst."
        assert step.timeout_ms == 300_000

    def test_python_step(self):
        step = Step.python(
            "transform",
            "Transform Data",
            code="print('hello')",
            timeout_ms=10_000,
        )
        assert step.type == StepType.TOOL
        assert step.config["tool"] == "python"
        assert step.config["script_content"] == "print('hello')"
        assert step.timeout_ms == 10_000

    def test_hitl_step(self):
        step = Step.hitl(
            "review",
            "Human Review",
            request_type=ApprovalRequestType.APPROVAL,
            timeout_minutes=60,
        )
        assert step.type == StepType.HITL
        assert step.config["request_type"] == "approval"
        assert step.config["timeout_minutes"] == 60

    def test_step_with_input_mappings(self):
        step = Step.agent(
            "summarize",
            "Summarize",
            input_mappings=[
                InputMapping(
                    source_step_id="extract",
                    source_field="text",
                    target_field="document",
                )
            ],
        )
        assert len(step.input_mappings) == 1
        assert step.input_mappings[0].source_step_id == "extract"

    def test_step_serialization(self):
        step = Step.python("test", "Test", code="print(1)")
        data = step.model_dump(exclude_none=True)
        assert data["id"] == "test"
        assert data["type"] == "tool"
        assert data["config"]["script_content"] == "print(1)"


class TestEdge:
    def test_simple_edge(self):
        edge = Edge(from_step_id="a", to_step_id="b")
        assert edge.from_step_id == "a"
        assert edge.condition is None

    def test_conditional_edge(self):
        edge = Edge(
            from_step_id="a",
            to_step_id="b",
            condition=BranchCondition(field="status", operator="equals", value="approved"),
        )
        assert edge.condition is not None
        assert edge.condition.field == "status"


class TestRunModel:
    def test_terminal_states(self):
        for status in (RunStatus.COMPLETED, RunStatus.FAILED, RunStatus.CANCELLED):
            run = Run(
                id="00000000-0000-0000-0000-000000000001",
                status=status,
            )
            assert run.is_terminal

    def test_non_terminal_states(self):
        for status in (RunStatus.PENDING, RunStatus.RUNNING, RunStatus.PAUSED):
            run = Run(
                id="00000000-0000-0000-0000-000000000001",
                status=status,
            )
            assert not run.is_terminal

    def test_from_api_response(self):
        run = Run.model_validate(
            {
                "id": "550e8400-e29b-41d4-a716-446655440000",
                "status": "running",
                "input": {"text": "hello"},
                "output": None,
                "error_message": None,
                "pause_requested": False,
            }
        )
        assert run.status == RunStatus.RUNNING
        assert run.input == {"text": "hello"}
