from __future__ import annotations

import enum
from typing import Self

from pydantic import BaseModel, model_validator


class OptimizedCandidateSource(str, enum.Enum):
    """
    Source of an optimized candidate.
    """

    OPTIMIZE = "OPTIMIZE"
    OPTIMIZE_LP = "OPTIMIZE_LP"
    REFINE = "REFINE"
    REPAIR = "REPAIR"
    ADAPTIVE = "ADAPTIVE"
    JIT_REWRITE = "JIT_REWRITE"


class OptimizeRequest(BaseModel):
    """
    Request body for POST /ai/optimize.
    """

    source_code: str
    dependency_code: str | None = None
    trace_id: str
    python_version: str | None = None
    language: str = "python"
    language_version: str | None = None
    experiment_metadata: dict[str, str] | None = None
    codeflash_version: str | None = None
    current_username: str | None = None
    repo_owner: str | None = None
    repo_name: str | None = None
    is_async: bool | None = False
    n_candidates: int = 5
    is_numerical_code: bool | None = None
    rerun_trace_id: str | None = None
    baseline_runtime_ns: int | None = None
    loop_count: int | None = None
    line_profiler_results: str | None = None
    test_input_examples: str | None = None

    @model_validator(mode="after")
    def resolve_python_version(self) -> Self:
        """
        Resolve python_version from language_version for backward compat.
        """
        if (
            self.python_version is None
            and self.language_version is not None
            and self.language == "python"
        ):
            self.python_version = self.language_version
        return self


class OptimizeResponseItem(BaseModel):
    """
    A single optimization candidate in the response.
    """

    source_code: str
    explanation: str
    optimization_id: str
    parent_id: str | None = None
    optimization_event_id: str | None = None


class OptimizeResponse(BaseModel):
    """
    Successful response from POST /ai/optimize.
    """

    optimizations: list[OptimizeResponseItem]


class OptimizeErrorResponse(BaseModel):
    """
    Error response from POST /ai/optimize.
    """

    error: str
