"""Shared optimizer models used across language handlers."""

from __future__ import annotations

import enum
from typing import Self

from ninja import Schema
from pydantic import model_validator


class OptimizedCandidateSource(str, enum.Enum):
    OPTIMIZE = "OPTIMIZE"
    OPTIMIZE_LP = "OPTIMIZE_LP"
    REFINE = "REFINE"
    REPAIR = "REPAIR"
    ADAPTIVE = "ADAPTIVE"
    JIT_REWRITE = "JIT_REWRITE"


class OptimizeSchema(Schema):
    source_code: str
    dependency_code: str | None
    trace_id: str
    python_version: str | None = None  # Made optional for multi-language support
    language: str = "python"  # NEW: language identifier (python, javascript, typescript)
    language_version: str | None = None  # NEW: e.g., "ES2022", "Node 20", or Python version
    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
    model: str | None = None  # Deprecated: multi-model is now handled by get_model_distribution
    call_sequence: int | None = None  # Deprecated: call_sequence is now auto-assigned
    n_candidates: int = 5  # default value for backward compatibility
    is_numerical_code: bool | None = None
    rerun_trace_id: str | None = None

    @model_validator(mode="after")
    def resolve_python_version(self) -> Self:
        """Resolve python_version from language_version for backward compatibility."""
        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 OptimizeSchemaLP(Schema):
    source_code: str
    dependency_code: str | None
    line_profiler_results: str | None
    trace_id: str
    python_version: str | None = None  # Made optional for multi-language support
    language: str = "python"  # NEW: language identifier (python, javascript, typescript)
    language_version: str | None = None  # NEW: e.g., "ES2022", "Node 20", or Python version
    experiment_metadata: dict[str, str] | None = None
    codeflash_version: str | None = None
    n_candidates: int = 6  # default value for backward compatibility
    model: str | None = None
    call_sequence: int | None = None
    is_numerical_code: bool | None = None
    rerun_trace_id: str | None = None

    @model_validator(mode="after")
    def resolve_python_version(self) -> Self:
        """Resolve python_version from language_version for backward compatibility."""
        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
