"""Unit tests for pipeline and stage models."""

import pytest
from aperion_flow.definitions.pipeline import (
    LLMAssistanceConfig,
    LLMAssistanceMode,
    Pipeline,
    PipelineBuilder,
    PipelineStage,
    PipelineStatus,
    StageStatus,
)


class TestPipelineStage:
    """Tests for PipelineStage model."""

    def test_stage_creation(self) -> None:
        """Test basic stage creation."""
        stage = PipelineStage(
            name="test_stage",
            handler="module.handler",
        )
        assert stage.name == "test_stage"
        assert stage.handler == "module.handler"
        assert stage.timeout == 30.0
        assert stage.retry_count == 0
        assert stage.depends_on == []

    def test_stage_with_dependencies(self) -> None:
        """Test stage with dependencies."""
        stage = PipelineStage(
            name="dependent_stage",
            handler="module.handler",
            depends_on=["stage1", "stage2"],
        )
        assert stage.depends_on == ["stage1", "stage2"]

    def test_stage_with_llm_assistance(self) -> None:
        """Test stage with LLM assistance configured."""
        config = LLMAssistanceConfig(mode=LLMAssistanceMode.PRE_VALIDATE)
        stage = PipelineStage(
            name="llm_stage",
            handler="module.handler",
            llm_assistance=[config],
        )
        assert stage.has_llm_mode(LLMAssistanceMode.PRE_VALIDATE)
        assert not stage.has_llm_mode(LLMAssistanceMode.ANALYZE_OUTPUT)

    def test_get_llm_config(self) -> None:
        """Test getting LLM config for specific mode."""
        config = LLMAssistanceConfig(
            mode=LLMAssistanceMode.ERROR_DIAGNOSIS,
            timeout=15.0,
        )
        stage = PipelineStage(
            name="llm_stage",
            handler="module.handler",
            llm_assistance=[config],
        )
        retrieved = stage.get_llm_config(LLMAssistanceMode.ERROR_DIAGNOSIS)
        assert retrieved is not None
        assert retrieved.timeout == 15.0

        missing = stage.get_llm_config(LLMAssistanceMode.ENRICH_CONTEXT)
        assert missing is None

    def test_stage_is_immutable(self) -> None:
        """Test that stages are frozen/immutable."""
        stage = PipelineStage(name="test", handler="module.handler")
        with pytest.raises(Exception):  # Pydantic ValidationError
            stage.name = "modified"


class TestPipeline:
    """Tests for Pipeline model."""

    def test_simple_pipeline(self) -> None:
        """Test creating a simple pipeline."""
        stages = [
            PipelineStage(id="s1", name="stage1", handler="mod.h1"),
            PipelineStage(id="s2", name="stage2", handler="mod.h2"),
        ]
        pipeline = Pipeline(name="test_pipeline", stages=stages)
        assert pipeline.name == "test_pipeline"
        assert len(pipeline.stages) == 2

    def test_pipeline_with_dependencies(self) -> None:
        """Test pipeline with stage dependencies."""
        stages = [
            PipelineStage(id="s1", name="fetch", handler="mod.fetch"),
            PipelineStage(id="s2", name="process", handler="mod.process", depends_on=["s1"]),
            PipelineStage(id="s3", name="store", handler="mod.store", depends_on=["s2"]),
        ]
        pipeline = Pipeline(name="data_pipeline", stages=stages)
        assert len(pipeline.stages) == 3

    def test_pipeline_rejects_unknown_dependency(self) -> None:
        """Test that unknown dependencies are rejected."""
        stages = [
            PipelineStage(id="s1", name="stage1", handler="mod.h1"),
            PipelineStage(id="s2", name="stage2", handler="mod.h2", depends_on=["unknown"]),
        ]
        with pytest.raises(ValueError, match="unknown stage ID"):
            Pipeline(name="invalid", stages=stages)

    def test_pipeline_rejects_cycle(self) -> None:
        """Test that cycles are rejected."""
        stages = [
            PipelineStage(id="s1", name="stage1", handler="mod.h1", depends_on=["s2"]),
            PipelineStage(id="s2", name="stage2", handler="mod.h2", depends_on=["s1"]),
        ]
        with pytest.raises(ValueError, match="cycle"):
            Pipeline(name="cyclic", stages=stages)

    def test_get_stage_by_id(self) -> None:
        """Test getting stage by ID."""
        stages = [
            PipelineStage(id="s1", name="stage1", handler="mod.h1"),
            PipelineStage(id="s2", name="stage2", handler="mod.h2"),
        ]
        pipeline = Pipeline(name="test", stages=stages)
        stage = pipeline.get_stage("s1")
        assert stage is not None
        assert stage.name == "stage1"

        missing = pipeline.get_stage("nonexistent")
        assert missing is None

    def test_get_root_stages(self) -> None:
        """Test getting stages with no dependencies."""
        stages = [
            PipelineStage(id="s1", name="root1", handler="mod.h1"),
            PipelineStage(id="s2", name="root2", handler="mod.h2"),
            PipelineStage(id="s3", name="child", handler="mod.h3", depends_on=["s1"]),
        ]
        pipeline = Pipeline(name="test", stages=stages)
        roots = pipeline.get_root_stages()
        assert len(roots) == 2
        root_names = {s.name for s in roots}
        assert root_names == {"root1", "root2"}

    def test_topological_order(self) -> None:
        """Test topological sorting of stages."""
        stages = [
            PipelineStage(id="s3", name="third", handler="mod.h3", depends_on=["s2"]),
            PipelineStage(id="s1", name="first", handler="mod.h1"),
            PipelineStage(id="s2", name="second", handler="mod.h2", depends_on=["s1"]),
        ]
        pipeline = Pipeline(name="test", stages=stages)
        ordered = pipeline.topological_order()

        names = [s.name for s in ordered]
        assert names.index("first") < names.index("second")
        assert names.index("second") < names.index("third")

    def test_get_dependents(self) -> None:
        """Test getting stages that depend on a given stage."""
        stages = [
            PipelineStage(id="s1", name="base", handler="mod.h1"),
            PipelineStage(id="s2", name="dep1", handler="mod.h2", depends_on=["s1"]),
            PipelineStage(id="s3", name="dep2", handler="mod.h3", depends_on=["s1"]),
        ]
        pipeline = Pipeline(name="test", stages=stages)
        deps = pipeline.get_dependents("s1")
        assert len(deps) == 2
        dep_names = {s.name for s in deps}
        assert dep_names == {"dep1", "dep2"}


class TestPipelineBuilder:
    """Tests for PipelineBuilder."""

    def test_builder_basic(self) -> None:
        """Test basic pipeline building."""
        pipeline = (
            PipelineBuilder("test_pipeline", "A test pipeline")
            .add_stage("fetch", "mod.fetch")
            .add_stage("process", "mod.process")
            .build()
        )
        assert pipeline.name == "test_pipeline"
        assert len(pipeline.stages) == 2

    def test_builder_with_options(self) -> None:
        """Test builder with various options."""
        pipeline = (
            PipelineBuilder("complex")
            .add_stage("s1", "mod.h1", timeout=60.0, retry_count=3)
            .with_max_parallel(2)
            .with_timeout(600.0)
            .with_fail_fast(False)
            .with_category("testing")
            .with_tags("ci", "critical")
            .build()
        )
        assert pipeline.max_parallel == 2
        assert pipeline.timeout == 600.0
        assert pipeline.fail_fast is False
        assert pipeline.category == "testing"
        assert "ci" in pipeline.tags

    def test_builder_with_llm_modes(self) -> None:
        """Test builder with LLM assistance modes."""
        pipeline = (
            PipelineBuilder("llm_pipeline")
            .add_stage(
                "analyze",
                "mod.analyze",
                llm_modes=[LLMAssistanceMode.PRE_VALIDATE, LLMAssistanceMode.ANALYZE_OUTPUT],
            )
            .build()
        )
        stage = pipeline.stages[0]
        assert len(stage.llm_assistance) == 2
        assert stage.has_llm_mode(LLMAssistanceMode.PRE_VALIDATE)


class TestLLMAssistanceConfig:
    """Tests for LLMAssistanceConfig."""

    def test_config_defaults(self) -> None:
        """Test default configuration values."""
        config = LLMAssistanceConfig(mode=LLMAssistanceMode.PRE_VALIDATE)
        assert config.enabled is True
        assert config.timeout == 10.0
        assert config.cache_results is True
        assert config.fallback_on_error is True

    def test_config_custom_values(self) -> None:
        """Test custom configuration values."""
        config = LLMAssistanceConfig(
            mode=LLMAssistanceMode.ERROR_DIAGNOSIS,
            enabled=True,
            timeout=30.0,
            cache_results=False,
            fallback_on_error=False,
            prompt_template="Custom prompt: {error}",
        )
        assert config.timeout == 30.0
        assert config.cache_results is False
        assert config.prompt_template == "Custom prompt: {error}"

    def test_config_is_immutable(self) -> None:
        """Test that config is frozen."""
        config = LLMAssistanceConfig(mode=LLMAssistanceMode.ANALYZE_OUTPUT)
        with pytest.raises(Exception):
            config.enabled = False


class TestEnums:
    """Tests for status enums."""

    def test_pipeline_status_values(self) -> None:
        """Test pipeline status enum values."""
        assert PipelineStatus.PENDING.value == "pending"
        assert PipelineStatus.RUNNING.value == "running"
        assert PipelineStatus.COMPLETED.value == "completed"
        assert PipelineStatus.FAILED.value == "failed"

    def test_stage_status_values(self) -> None:
        """Test stage status enum values."""
        assert StageStatus.PENDING.value == "pending"
        assert StageStatus.RUNNING.value == "running"
        assert StageStatus.COMPLETED.value == "completed"
        assert StageStatus.SKIPPED.value == "skipped"

    def test_llm_mode_values(self) -> None:
        """Test LLM assistance mode enum values."""
        assert LLMAssistanceMode.PRE_VALIDATE.value == "pre_validate"
        assert LLMAssistanceMode.ENRICH_CONTEXT.value == "enrich_context"
        assert LLMAssistanceMode.ANALYZE_OUTPUT.value == "analyze_output"
        assert LLMAssistanceMode.ERROR_DIAGNOSIS.value == "error_diagnosis"
