"""Unit tests for pipeline context and checkpointing."""

import pytest
from datetime import datetime, timezone
from aperion_flow.engine.context import (
    Checkpoint,
    CheckpointType,
    InMemoryCheckpointStore,
    PipelineContext,
    SecretRef,
    StageResult,
)
from aperion_flow.definitions.pipeline import PipelineStatus, StageStatus


class TestPipelineContext:
    """Tests for PipelineContext."""

    def test_context_creation(self) -> None:
        """Test basic context creation."""
        ctx = PipelineContext(
            pipeline_id="p1",
            pipeline_name="test_pipeline",
        )
        assert ctx.pipeline_id == "p1"
        assert ctx.pipeline_name == "test_pipeline"
        assert ctx.status == PipelineStatus.PENDING
        assert ctx.execution_id  # Should be generated

    def test_context_data_operations(self) -> None:
        """Test data get/set/update operations."""
        ctx = PipelineContext(pipeline_id="p1", pipeline_name="test")

        ctx.set("key1", "value1")
        assert ctx.get("key1") == "value1"
        assert ctx.get("missing", "default") == "default"

        ctx.update({"key2": "value2", "key3": "value3"})
        assert ctx.data == {"key1": "value1", "key2": "value2", "key3": "value3"}

    def test_context_status_transitions(self) -> None:
        """Test status transition methods."""
        ctx = PipelineContext(pipeline_id="p1", pipeline_name="test")

        ctx.mark_started()
        assert ctx.status == PipelineStatus.RUNNING
        assert ctx.started_at is not None

        ctx.mark_completed()
        assert ctx.status == PipelineStatus.COMPLETED
        assert ctx.completed_at is not None

    def test_context_failure_marking(self) -> None:
        """Test marking context as failed."""
        ctx = PipelineContext(pipeline_id="p1", pipeline_name="test")
        ctx.mark_started()
        ctx.mark_failed("Something went wrong", "RuntimeError", "stage1")

        assert ctx.status == PipelineStatus.FAILED
        assert ctx.error == "Something went wrong"
        assert ctx.error_type == "RuntimeError"
        assert ctx.error_stage == "stage1"

    def test_context_stage_recording(self) -> None:
        """Test recording stage execution."""
        ctx = PipelineContext(pipeline_id="p1", pipeline_name="test")

        ctx.record_stage_start("s1", "stage1")
        assert ctx.current_stage_id == "s1"
        assert "s1" in ctx.stage_results
        assert ctx.stage_results["s1"].status == StageStatus.RUNNING

        ctx.record_stage_complete("s1", {"output": "data"})
        assert ctx.current_stage_id is None
        assert ctx.stage_results["s1"].status == StageStatus.COMPLETED
        assert ctx.stage_results["s1"].data == {"output": "data"}

    def test_context_stage_failure_recording(self) -> None:
        """Test recording stage failure."""
        ctx = PipelineContext(pipeline_id="p1", pipeline_name="test")

        ctx.record_stage_start("s1", "stage1")
        ctx.record_stage_failed("s1", "Error occurred", "ValueError")

        assert ctx.stage_results["s1"].status == StageStatus.FAILED
        assert ctx.stage_results["s1"].error == "Error occurred"
        assert ctx.stage_results["s1"].error_type == "ValueError"

    def test_context_stage_skip_recording(self) -> None:
        """Test recording skipped stage."""
        ctx = PipelineContext(pipeline_id="p1", pipeline_name="test")

        ctx.record_stage_skipped("s1", "stage1", "Dependency failed")
        assert ctx.stage_results["s1"].status == StageStatus.SKIPPED
        assert ctx.stage_results["s1"].metadata["skip_reason"] == "Dependency failed"

    def test_completed_stage_ids(self) -> None:
        """Test getting completed stage IDs."""
        ctx = PipelineContext(pipeline_id="p1", pipeline_name="test")

        ctx.record_stage_start("s1", "stage1")
        ctx.record_stage_complete("s1")

        ctx.record_stage_start("s2", "stage2")
        ctx.record_stage_failed("s2", "error")

        ctx.record_stage_start("s3", "stage3")
        ctx.record_stage_complete("s3")

        assert set(ctx.completed_stage_ids) == {"s1", "s3"}
        assert ctx.failed_stage_ids == ["s2"]

    def test_context_duration(self) -> None:
        """Test duration calculation."""
        ctx = PipelineContext(pipeline_id="p1", pipeline_name="test")
        assert ctx.duration_seconds is None

        ctx.mark_started()
        ctx.mark_completed()
        assert ctx.duration_seconds is not None
        assert ctx.duration_seconds >= 0


class TestCheckpoint:
    """Tests for Checkpoint model."""

    def test_checkpoint_creation(self) -> None:
        """Test basic checkpoint creation."""
        checkpoint = Checkpoint(
            pipeline_id="p1",
            execution_id="e1",
            checkpoint_type=CheckpointType.PIPELINE_START,
            status=PipelineStatus.RUNNING,
        )
        assert checkpoint.pipeline_id == "p1"
        assert checkpoint.checkpoint_type == CheckpointType.PIPELINE_START

    def test_checkpoint_serialization(self) -> None:
        """Test checkpoint JSON serialization."""
        checkpoint = Checkpoint(
            pipeline_id="p1",
            execution_id="e1",
            checkpoint_type=CheckpointType.STAGE_COMPLETE,
            stage_id="s1",
            status=StageStatus.COMPLETED,
            data_snapshot={"key": "value"},
            completed_stages=["s1"],
        )

        json_str = checkpoint.to_json()
        restored = Checkpoint.from_json(json_str)

        assert restored.pipeline_id == checkpoint.pipeline_id
        assert restored.stage_id == checkpoint.stage_id
        assert restored.data_snapshot == checkpoint.data_snapshot

    def test_checkpoint_types(self) -> None:
        """Test all checkpoint types exist."""
        assert CheckpointType.PIPELINE_START
        assert CheckpointType.STAGE_START
        assert CheckpointType.STAGE_COMPLETE
        assert CheckpointType.STAGE_FAILED
        assert CheckpointType.PIPELINE_COMPLETE
        assert CheckpointType.PIPELINE_FAILED
        assert CheckpointType.RECOVERY


class TestContextCheckpointing:
    """Tests for context checkpoint creation and restoration."""

    def test_create_checkpoint(self) -> None:
        """Test creating checkpoint from context."""
        ctx = PipelineContext(pipeline_id="p1", pipeline_name="test")
        ctx.set("key", "value")
        ctx.record_stage_start("s1", "stage1")
        ctx.record_stage_complete("s1")
        ctx.mark_started()

        checkpoint = ctx.create_checkpoint(CheckpointType.STAGE_COMPLETE)

        assert checkpoint.pipeline_id == "p1"
        assert checkpoint.execution_id == ctx.execution_id
        assert checkpoint.data_snapshot == {"key": "value"}
        assert "s1" in checkpoint.completed_stages

    def test_restore_from_checkpoint(self) -> None:
        """Test restoring context from checkpoint."""
        checkpoint = Checkpoint(
            pipeline_id="p1",
            execution_id="e1",
            checkpoint_type=CheckpointType.STAGE_COMPLETE,
            status=PipelineStatus.RUNNING,
            data_snapshot={"restored": "data"},
            completed_stages=["s1", "s2"],
            failed_stages=["s3"],
        )

        ctx = PipelineContext(pipeline_id="p1", pipeline_name="test", execution_id="e1")
        ctx.restore_from_checkpoint(checkpoint)

        assert ctx.data == {"restored": "data"}
        assert "s1" in ctx.stage_results
        assert ctx.stage_results["s1"].status == StageStatus.COMPLETED
        assert ctx.stage_results["s3"].status == StageStatus.FAILED


class TestInMemoryCheckpointStore:
    """Tests for InMemoryCheckpointStore."""

    @pytest.fixture
    def store(self) -> InMemoryCheckpointStore:
        return InMemoryCheckpointStore()

    @pytest.mark.asyncio
    async def test_save_and_load(self, store: InMemoryCheckpointStore) -> None:
        """Test saving and loading checkpoints."""
        checkpoint = Checkpoint(
            pipeline_id="p1",
            execution_id="e1",
            checkpoint_type=CheckpointType.PIPELINE_START,
            status=PipelineStatus.RUNNING,
        )

        await store.save(checkpoint)
        loaded = await store.load("e1")

        assert loaded is not None
        assert loaded.pipeline_id == "p1"

    @pytest.mark.asyncio
    async def test_load_returns_latest(self, store: InMemoryCheckpointStore) -> None:
        """Test that load returns the latest checkpoint."""
        cp1 = Checkpoint(
            pipeline_id="p1",
            execution_id="e1",
            checkpoint_type=CheckpointType.PIPELINE_START,
            status=PipelineStatus.RUNNING,
        )
        cp2 = Checkpoint(
            pipeline_id="p1",
            execution_id="e1",
            checkpoint_type=CheckpointType.STAGE_COMPLETE,
            status=PipelineStatus.RUNNING,
        )

        await store.save(cp1)
        await store.save(cp2)
        loaded = await store.load("e1")

        assert loaded is not None
        assert loaded.checkpoint_type == CheckpointType.STAGE_COMPLETE

    @pytest.mark.asyncio
    async def test_list_checkpoints(self, store: InMemoryCheckpointStore) -> None:
        """Test listing all checkpoints for an execution."""
        for i in range(3):
            cp = Checkpoint(
                pipeline_id="p1",
                execution_id="e1",
                checkpoint_type=CheckpointType.STAGE_COMPLETE,
                stage_id=f"s{i}",
                status=PipelineStatus.RUNNING,
            )
            await store.save(cp)

        checkpoints = await store.list_checkpoints("e1")
        assert len(checkpoints) == 3

    @pytest.mark.asyncio
    async def test_delete_checkpoints(self, store: InMemoryCheckpointStore) -> None:
        """Test deleting checkpoints."""
        cp = Checkpoint(
            pipeline_id="p1",
            execution_id="e1",
            checkpoint_type=CheckpointType.PIPELINE_START,
            status=PipelineStatus.RUNNING,
        )
        await store.save(cp)
        await store.delete("e1")

        loaded = await store.load("e1")
        assert loaded is None

    @pytest.mark.asyncio
    async def test_load_nonexistent(self, store: InMemoryCheckpointStore) -> None:
        """Test loading from nonexistent execution."""
        loaded = await store.load("nonexistent")
        assert loaded is None


class TestSecretRef:
    """Tests for SecretRef."""

    def test_secret_ref_creation(self) -> None:
        """Test creating a secret reference."""
        ref = SecretRef(name="API_KEY")
        assert ref.name == "API_KEY"
        assert ref.source == "env"
        assert ref.required is True

    def test_secret_ref_with_vault(self) -> None:
        """Test secret reference with vault source."""
        ref = SecretRef(
            name="db_password",
            source="vault",
            required=True,
            default_ref="DB_PASSWORD_DEV",
        )
        assert ref.source == "vault"
        assert ref.default_ref == "DB_PASSWORD_DEV"

    def test_secret_ref_repr(self) -> None:
        """Test secret reference string representation."""
        ref = SecretRef(name="SECRET", source="env")
        assert "SecretRef(env:SECRET)" in repr(ref)

    def test_secret_ref_is_immutable(self) -> None:
        """Test that secret ref is frozen."""
        ref = SecretRef(name="SECRET")
        with pytest.raises(Exception):
            ref.name = "CHANGED"


class TestStageResult:
    """Tests for StageResult."""

    def test_stage_result_creation(self) -> None:
        """Test basic stage result creation."""
        result = StageResult(
            stage_id="s1",
            stage_name="stage1",
            status=StageStatus.COMPLETED,
            data={"output": "value"},
        )
        assert result.stage_id == "s1"
        assert result.status == StageStatus.COMPLETED

    def test_stage_result_duration(self) -> None:
        """Test duration calculation."""
        start = datetime.now(timezone.utc)
        # Simulate 1 second later
        end = datetime.now(timezone.utc)

        result = StageResult(
            stage_id="s1",
            stage_name="stage1",
            status=StageStatus.COMPLETED,
            started_at=start,
            completed_at=end,
        )

        assert result.duration_seconds is not None
        assert result.duration_seconds >= 0

    def test_stage_result_no_duration_without_times(self) -> None:
        """Test duration is None when times not set."""
        result = StageResult(
            stage_id="s1",
            stage_name="stage1",
            status=StageStatus.PENDING,
        )
        assert result.duration_seconds is None
