from __future__ import annotations

from typing import Self

from pydantic import BaseModel, model_validator


class JitRewriteRequest(BaseModel):
    """
    Request body for POST /ai/rewrite_jit.
    """

    source_code: str
    trace_id: str
    dependency_code: str | None = None
    python_version: str | None = None
    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
    n_candidates: int = 1
    language: str = "python"

    @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
