"""Tests for experiment type definitions."""

from datetime import datetime, timezone

from lumenova_beacon.experiments.types import (
    EnrichedRecord,
    ExperimentRun,
    ExperimentRunStats,
    ExperimentRunStatus,
    ExperimentStatus,
    ExperimentStep,
    ExperimentStepType,
    StepRunSummary,
)


class TestExperimentStatus:
    def test_all_values(self):
        assert ExperimentStatus.DRAFT == 1
        assert ExperimentStatus.QUEUED == 2
        assert ExperimentStatus.RUNNING == 3
        assert ExperimentStatus.COMPLETED == 4
        assert ExperimentStatus.FAILED == 5
        assert ExperimentStatus.STOPPED == 6
        assert ExperimentStatus.COMPLETED_WITH_ERRORS == 7
        assert ExperimentStatus.WAITING_FOR_EXTERNAL == 8

    def test_has_eight_members(self):
        assert len(ExperimentStatus) == 8

    def test_construct_from_int(self):
        assert ExperimentStatus(7) == ExperimentStatus.COMPLETED_WITH_ERRORS
        assert ExperimentStatus(8) == ExperimentStatus.WAITING_FOR_EXTERNAL


class TestExperimentStepType:
    def test_all_values(self):
        assert ExperimentStepType.PROMPT == "prompt"
        assert ExperimentStepType.LLM == "llm"
        assert ExperimentStepType.EVALUATION == "evaluation"
        assert ExperimentStepType.EXTERNAL_AGENT == "external_agent"

    def test_has_four_members(self):
        assert len(ExperimentStepType) == 4

    def test_is_string(self):
        assert isinstance(ExperimentStepType.PROMPT, str)


class TestExperimentRunStatus:
    def test_all_values(self):
        assert ExperimentRunStatus.PENDING == "pending"
        assert ExperimentRunStatus.RUNNING == "running"
        assert ExperimentRunStatus.COMPLETED == "completed"
        assert ExperimentRunStatus.FAILED == "failed"

    def test_has_four_members(self):
        assert len(ExperimentRunStatus) == 4

    def test_is_string(self):
        assert isinstance(ExperimentRunStatus.PENDING, str)


class TestExperimentStep:
    def test_to_dict(self):
        step = ExperimentStep(
            step_type=ExperimentStepType.LLM,
            output_column="llm_output",
            config={"model_config_id": "abc", "variable_mappings": {"input": "col1"}},
        )
        result = step.to_dict()
        assert result == {
            "step_type": "llm",
            "output_column": "llm_output",
            "config": {
                "model_config_id": "abc",
                "variable_mappings": {"input": "col1"},
            },
        }

    def test_to_dict_with_string_step_type(self):
        step = ExperimentStep(
            step_type="external_agent",
            output_column="output",
            config={},
        )
        result = step.to_dict()
        assert result["step_type"] == "external_agent"

    def test_from_dict_full(self):
        data = {
            "id": "step-123",
            "experiment_id": "exp-456",
            "position": 2,
            "step_type": "prompt",
            "output_column": "rendered",
            "config": {"prompt_content": "Hello {{name}}"},
            "created_at": "2025-01-15T10:00:00Z",
            "updated_at": "2025-01-15T11:00:00Z",
        }
        step = ExperimentStep.from_dict(data)
        assert step.id == "step-123"
        assert step.experiment_id == "exp-456"
        assert step.position == 2
        assert step.step_type == "prompt"
        assert step.output_column == "rendered"
        assert step.config == {"prompt_content": "Hello {{name}}"}
        assert step.created_at == datetime(2025, 1, 15, 10, 0, tzinfo=timezone.utc)
        assert step.updated_at == datetime(2025, 1, 15, 11, 0, tzinfo=timezone.utc)

    def test_from_dict_minimal(self):
        data = {
            "step_type": "llm",
            "output_column": "output",
        }
        step = ExperimentStep.from_dict(data)
        assert step.step_type == "llm"
        assert step.output_column == "output"
        assert step.config == {}
        assert step.id is None
        assert step.position is None

    def test_roundtrip(self):
        original = ExperimentStep(
            step_type=ExperimentStepType.EVALUATION,
            output_column="eval_score",
            config={"evaluator_id": "eval-1"},
        )
        serialized = original.to_dict()
        restored = ExperimentStep.from_dict(serialized)
        assert restored.step_type == "evaluation"
        assert restored.output_column == "eval_score"
        assert restored.config == {"evaluator_id": "eval-1"}


class TestExperimentRun:
    def test_from_dict_full(self):
        data = {
            "id": "run-001",
            "experiment_id": "exp-001",
            "record_id": "rec-001",
            "step_id": "step-001",
            "configuration_id": None,
            "status": "completed",
            "result": {"output": "hello", "tokens": 10},
            "error": None,
            "execution_time_ms": 150,
            "started_at": "2025-01-15T10:00:00Z",
            "completed_at": "2025-01-15T10:00:01Z",
            "created_at": "2025-01-15T09:59:00Z",
            "updated_at": "2025-01-15T10:00:01Z",
        }
        run = ExperimentRun.from_dict(data)
        assert run.id == "run-001"
        assert run.experiment_id == "exp-001"
        assert run.record_id == "rec-001"
        assert run.step_id == "step-001"
        assert run.configuration_id is None
        assert run.status == ExperimentRunStatus.COMPLETED
        assert run.result == {"output": "hello", "tokens": 10}
        assert run.error is None
        assert run.execution_time_ms == 150
        assert run.started_at is not None
        assert run.completed_at is not None

    def test_from_dict_minimal(self):
        data = {
            "id": "run-002",
            "experiment_id": "exp-001",
            "record_id": "rec-002",
            "status": "pending",
        }
        run = ExperimentRun.from_dict(data)
        assert run.id == "run-002"
        assert run.status == ExperimentRunStatus.PENDING
        assert run.step_id is None
        assert run.result is None
        assert run.execution_time_ms is None

    def test_from_dict_failed(self):
        data = {
            "id": "run-003",
            "experiment_id": "exp-001",
            "record_id": "rec-003",
            "status": "failed",
            "error": "LLM timeout",
            "execution_time_ms": 30000,
        }
        run = ExperimentRun.from_dict(data)
        assert run.status == ExperimentRunStatus.FAILED
        assert run.error == "LLM timeout"

    def test_from_dict_with_configuration_label(self):
        data = {
            "id": "run-004",
            "experiment_id": "exp-001",
            "record_id": "rec-004",
            "status": "completed",
            "configuration_label": "B",
        }
        run = ExperimentRun.from_dict(data)
        assert run.configuration_label == "B"

    def test_from_dict_without_configuration_label(self):
        data = {
            "id": "run-005",
            "experiment_id": "exp-001",
            "record_id": "rec-005",
            "status": "pending",
        }
        run = ExperimentRun.from_dict(data)
        assert run.configuration_label is None


class TestExperimentRunStats:
    def test_from_dict_full(self):
        data = {
            "experiment_id": "exp-001",
            "total_runs": 300,
            "pending_runs": 50,
            "running_runs": 10,
            "completed_runs": 230,
            "failed_runs": 10,
            "avg_execution_time_ms": 250.5,
            "min_execution_time_ms": 50,
            "max_execution_time_ms": 5000,
        }
        stats = ExperimentRunStats.from_dict(data)
        assert stats.experiment_id == "exp-001"
        assert stats.total_runs == 300
        assert stats.pending_runs == 50
        assert stats.running_runs == 10
        assert stats.completed_runs == 230
        assert stats.failed_runs == 10
        assert stats.avg_execution_time_ms == 250.5
        assert stats.min_execution_time_ms == 50
        assert stats.max_execution_time_ms == 5000

    def test_from_dict_no_timing(self):
        data = {
            "experiment_id": "exp-002",
            "total_runs": 0,
            "pending_runs": 0,
            "running_runs": 0,
            "completed_runs": 0,
            "failed_runs": 0,
        }
        stats = ExperimentRunStats.from_dict(data)
        assert stats.avg_execution_time_ms is None
        assert stats.min_execution_time_ms is None
        assert stats.max_execution_time_ms is None


class TestStepRunSummary:
    def test_from_dict_full(self):
        data = {
            "step_id": "step-001",
            "step_position": 0,
            "step_type": "llm",
            "output_column": "llm_output",
            "status": "completed",
            "output": "Generated text",
            "error_message": None,
            "execution_time_ms": 200,
            "execution_metadata": {"model": "gpt-4", "tokens": 50},
        }
        summary = StepRunSummary.from_dict(data)
        assert summary.step_id == "step-001"
        assert summary.step_position == 0
        assert summary.step_type == "llm"
        assert summary.output_column == "llm_output"
        assert summary.status == ExperimentRunStatus.COMPLETED
        assert summary.output == "Generated text"
        assert summary.error_message is None
        assert summary.execution_time_ms == 200
        assert summary.execution_metadata == {"model": "gpt-4", "tokens": 50}

    def test_from_dict_failed(self):
        data = {
            "step_id": "step-002",
            "step_position": 1,
            "step_type": "evaluation",
            "output_column": "eval_score",
            "status": "failed",
            "error_message": "Evaluator timeout",
        }
        summary = StepRunSummary.from_dict(data)
        assert summary.status == ExperimentRunStatus.FAILED
        assert summary.error_message == "Evaluator timeout"
        assert summary.output is None


class TestEnrichedRecord:
    def test_from_dict_full(self):
        data = {
            "record_id": "rec-001",
            "data": {"input": "hello", "category": "greeting"},
            "overall_status": "success",
            "step_runs": [
                {
                    "step_id": "step-001",
                    "step_position": 0,
                    "step_type": "prompt",
                    "output_column": "rendered",
                    "status": "completed",
                    "output": "Rendered: hello",
                },
                {
                    "step_id": "step-002",
                    "step_position": 1,
                    "step_type": "llm",
                    "output_column": "response",
                    "status": "completed",
                    "output": "Hi there!",
                    "execution_time_ms": 150,
                },
            ],
        }
        record = EnrichedRecord.from_dict(data)
        assert record.record_id == "rec-001"
        assert record.data == {"input": "hello", "category": "greeting"}
        assert record.overall_status == "success"
        assert len(record.step_runs) == 2
        assert record.step_runs[0].step_type == "prompt"
        assert record.step_runs[1].output == "Hi there!"

    def test_from_dict_no_step_runs(self):
        data = {
            "record_id": "rec-002",
            "data": {"input": "test"},
            "overall_status": "pending",
        }
        record = EnrichedRecord.from_dict(data)
        assert record.step_runs == []

    def test_from_dict_empty_data(self):
        data = {
            "record_id": "rec-003",
            "overall_status": "pending",
        }
        record = EnrichedRecord.from_dict(data)
        assert record.data == {}
        assert record.step_runs == []
