"""Type definitions for experiment operations."""

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum, IntEnum
from typing import Any

from lumenova_beacon.utils.datetime import parse_iso_datetime


class ExperimentStatus(IntEnum):
    """Experiment status enum matching API values."""

    DRAFT = 1
    QUEUED = 2
    RUNNING = 3
    COMPLETED = 4
    FAILED = 5
    STOPPED = 6
    COMPLETED_WITH_ERRORS = 7
    WAITING_FOR_EXTERNAL = 8


class ExperimentStepType(str, Enum):
    """Step type within an experiment pipeline."""

    PROMPT = "prompt"
    LLM = "llm"
    EVALUATION = "evaluation"
    EXTERNAL_AGENT = "external_agent"


class ExperimentRunStatus(str, Enum):
    """Status of an individual experiment run."""

    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"


@dataclass
class ExperimentStep:
    """A step in an experiment pipeline.

    Steps define the operations executed sequentially per record.
    Each step has a type, an output column for results, and type-specific
    configuration including variable mappings.

    Examples:
        Create a prompt step:
            >>> step = ExperimentStep(
            ...     step_type=ExperimentStepType.PROMPT,
            ...     output_column="rendered_prompt",
            ...     config={
            ...         "prompt_content": "Summarize: {{input}}",
            ...         "variable_mappings": {"input": "text_column"},
            ...     },
            ... )

        Create an LLM step:
            >>> step = ExperimentStep(
            ...     step_type=ExperimentStepType.LLM,
            ...     output_column="llm_output",
            ...     config={
            ...         "model_config_id": "config-uuid",
            ...         "model_parameters": {"temperature": 0.7},
            ...         "variable_mappings": {"input": "rendered_prompt"},
            ...     },
            ... )

        Create an external agent step:
            >>> step = ExperimentStep(
            ...     step_type=ExperimentStepType.EXTERNAL_AGENT,
            ...     output_column="agent_output",
            ...     config={
            ...         "agent_name": "My Agent",
            ...         "variable_mappings": {"input": "text_column"},
            ...     },
            ... )
    """

    step_type: ExperimentStepType | str
    """Type of step: prompt, llm, evaluation, or external_agent."""

    output_column: str
    """Column name in the dataset where this step writes its output."""

    config: dict[str, Any] = field(default_factory=dict)
    """Type-specific configuration including variable_mappings."""

    id: str | None = None
    """Step UUID (populated from server response)."""

    experiment_id: str | None = None
    """Parent experiment UUID (populated from server response)."""

    position: int | None = None
    """Position in the pipeline (0-indexed, auto-assigned by server)."""

    created_at: datetime | None = None
    """Creation timestamp (populated from server response)."""

    updated_at: datetime | None = None
    """Last update timestamp (populated from server response)."""

    def to_dict(self) -> dict[str, Any]:
        """Convert to API request format for creating a step.

        Returns:
            Dictionary with step_type, output_column, and config.
        """
        step_type = (
            self.step_type.value
            if isinstance(self.step_type, ExperimentStepType)
            else self.step_type
        )
        return {
            "step_type": step_type,
            "output_column": self.output_column,
            "config": self.config,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> ExperimentStep:
        """Create an ExperimentStep from an API response dict."""
        return cls(
            step_type=data["step_type"],
            output_column=data["output_column"],
            config=data.get("config", {}),
            id=data.get("id"),
            experiment_id=data.get("experiment_id"),
            position=data.get("position"),
            created_at=parse_iso_datetime(data.get("created_at")),
            updated_at=parse_iso_datetime(data.get("updated_at")),
        )


@dataclass
class ExperimentRun:
    """An individual experiment run execution record.

    Represents a single (step x record) execution. For a 3-step pipeline
    with 100 records, there will be 300 runs total.
    """

    id: str
    """Run UUID."""

    experiment_id: str
    """Parent experiment UUID."""

    record_id: str
    """Dataset record UUID this run was executed against."""

    status: ExperimentRunStatus
    """Current status of the run."""

    step_id: str | None = None
    """Step UUID (for pipeline experiments)."""

    configuration_id: str | None = None
    """Configuration UUID (for legacy experiments)."""

    result: dict[str, Any] | None = None
    """Execution result (output + metadata)."""

    error: str | None = None
    """Error message if the run failed."""

    execution_time_ms: int | None = None
    """Execution time in milliseconds."""

    started_at: datetime | None = None
    """When the run started executing."""

    completed_at: datetime | None = None
    """When the run finished."""

    created_at: datetime | None = None
    """When the run record was created."""

    updated_at: datetime | None = None
    """When the run record was last updated."""

    configuration_label: str | None = None
    """Configuration label (A, B, C, etc.) from list responses."""

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> ExperimentRun:
        """Create an ExperimentRun from an API response dict."""
        return cls(
            id=data["id"],
            experiment_id=data["experiment_id"],
            record_id=data["record_id"],
            status=ExperimentRunStatus(data["status"]),
            step_id=data.get("step_id"),
            configuration_id=data.get("configuration_id"),
            result=data.get("result"),
            error=data.get("error"),
            execution_time_ms=data.get("execution_time_ms"),
            started_at=parse_iso_datetime(data.get("started_at")),
            completed_at=parse_iso_datetime(data.get("completed_at")),
            created_at=parse_iso_datetime(data.get("created_at")),
            updated_at=parse_iso_datetime(data.get("updated_at")),
            configuration_label=data.get("configuration_label"),
        )


@dataclass
class ExperimentRunStats:
    """Aggregated run statistics for an experiment."""

    experiment_id: str
    """Experiment UUID."""

    total_runs: int
    """Total number of runs."""

    pending_runs: int
    """Number of runs pending execution."""

    running_runs: int
    """Number of runs currently executing."""

    completed_runs: int
    """Number of successfully completed runs."""

    failed_runs: int
    """Number of failed runs."""

    avg_execution_time_ms: float | None = None
    """Average execution time across completed runs."""

    min_execution_time_ms: int | None = None
    """Minimum execution time across completed runs."""

    max_execution_time_ms: int | None = None
    """Maximum execution time across completed runs."""

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> ExperimentRunStats:
        """Create an ExperimentRunStats from an API response dict."""
        return cls(
            experiment_id=data["experiment_id"],
            total_runs=data["total_runs"],
            pending_runs=data["pending_runs"],
            running_runs=data["running_runs"],
            completed_runs=data["completed_runs"],
            failed_runs=data["failed_runs"],
            avg_execution_time_ms=data.get("avg_execution_time_ms"),
            min_execution_time_ms=data.get("min_execution_time_ms"),
            max_execution_time_ms=data.get("max_execution_time_ms"),
        )


@dataclass
class StepRunSummary:
    """Per-step execution summary within an enriched record."""

    step_id: str
    """Step UUID."""

    step_position: int
    """Position of the step in the pipeline (0-indexed)."""

    step_type: str
    """Type of the step (prompt, llm, evaluation, external_agent)."""

    output_column: str
    """Column name where this step writes its output."""

    status: ExperimentRunStatus
    """Execution status for this step on this record."""

    output: Any | None = None
    """Step output value (null if not yet executed or failed)."""

    error_message: str | None = None
    """Error details if the step failed."""

    execution_time_ms: int | None = None
    """Execution time in milliseconds."""

    execution_metadata: dict[str, Any] | None = None
    """Additional execution metadata (model, tokens, etc.)."""

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> StepRunSummary:
        """Create a StepRunSummary from an API response dict."""
        return cls(
            step_id=data["step_id"],
            step_position=data["step_position"],
            step_type=data["step_type"],
            output_column=data["output_column"],
            status=ExperimentRunStatus(data["status"]),
            output=data.get("output"),
            error_message=data.get("error_message"),
            execution_time_ms=data.get("execution_time_ms"),
            execution_metadata=data.get("execution_metadata"),
        )


@dataclass
class EnrichedRecord:
    """A dataset record enriched with per-step execution data."""

    record_id: str
    """Dataset record UUID."""

    data: dict[str, Any]
    """Original record data."""

    overall_status: str
    """Derived status: pending, running, success, or failed."""

    step_runs: list[StepRunSummary] = field(default_factory=list)
    """Per-step execution details for this record."""

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> EnrichedRecord:
        """Create an EnrichedRecord from an API response dict."""
        return cls(
            record_id=data["record_id"],
            data=data.get("data", {}),
            overall_status=data["overall_status"],
            step_runs=[
                StepRunSummary.from_dict(sr) for sr in data.get("step_runs", [])
            ],
        )
