"""Request and response schemas for testgen endpoints."""

from __future__ import annotations

from typing import Self

from pydantic import BaseModel, model_validator

# ── Testgen ─────────────────────────────────────────────────────────


class TestGenRequest(BaseModel):
    """
    Request body for POST /ai/testgen.
    """

    __test__ = False

    source_code_being_tested: str
    function_name: str
    module_path: str
    test_module_path: str
    test_framework: str = "pytest"
    test_timeout: int = 300
    trace_id: str
    python_version: str | None = None
    language: str = "python"
    language_version: str | None = None
    codeflash_version: str | None = None
    test_index: int | None = None
    is_async: bool = False
    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.
        """
        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 TestGenResponse(BaseModel):
    """
    Successful response from POST /ai/testgen.
    """

    __test__ = False

    generated_tests: str
    raw_generated_tests: str | None = None


class TestGenErrorResponse(BaseModel):
    """
    Error response from testgen endpoints.
    """

    __test__ = False

    error: str
    trace_id: str | None = None


# ── Review ──────────────────────────────────────────────────────────


class TestSourceWithFailures(BaseModel):
    """
    A test source file with optional behavioral test failures.
    """

    __test__ = False

    test_source: str
    test_index: int
    failed_test_functions: list[str] = []
    failure_messages: dict[str, str] = {}


class CoverageFunctionDetail(BaseModel):
    """
    Per-function coverage detail.
    """

    name: str
    coverage: float
    executed_lines: list[int] = []
    unexecuted_lines: list[int] = []
    executed_branches: list[list[int]] = []
    unexecuted_branches: list[list[int]] = []


class CoverageDetails(BaseModel):
    """
    Coverage info for the function under test.
    """

    coverage_percentage: float
    threshold_percentage: float
    function_start_line: int = 1
    main_function: CoverageFunctionDetail
    dependent_function: CoverageFunctionDetail | None = None


class TestgenReviewRequest(BaseModel):
    """
    Request body for POST /ai/testgen_review.
    """

    __test__ = False

    tests: list[TestSourceWithFailures]
    function_source_code: str
    function_name: str
    trace_id: str
    language: str = "python"
    codeflash_version: str | None = None
    call_sequence: int | None = None
    coverage_summary: str = ""
    coverage_details: CoverageDetails | None = None


class FunctionVerdict(BaseModel):
    """
    Verdict for a single test function.
    """

    function_name: str
    verdict: str
    reason: str = ""


class TestReview(BaseModel):
    """
    Review result for one test source.
    """

    __test__ = False

    test_index: int
    functions: list[FunctionVerdict]


class TestgenReviewResponse(BaseModel):
    """
    Response from POST /ai/testgen_review.
    """

    __test__ = False

    reviews: list[TestReview]


# ── Repair ──────────────────────────────────────────────────────────


class FunctionToRepair(BaseModel):
    """
    A single function that needs repair with the reason.
    """

    function_name: str
    reason: str


class TestRepairRequest(BaseModel):
    """
    Request body for POST /ai/testgen_repair.
    """

    __test__ = False

    test_source: str
    functions_to_repair: list[FunctionToRepair]
    function_source_code: str
    function_name: str
    module_path: str
    test_module_path: str
    test_framework: str
    test_timeout: int
    trace_id: str
    python_version: str | None = None
    language_version: str | None = None
    language: str = "python"
    codeflash_version: str | None = None
    call_sequence: int | None = None
    coverage_details: CoverageDetails | None = None
    previous_repair_errors: dict[str, str] = {}
    module_source_code: str = ""

    @model_validator(mode="after")
    def resolve_python_version(self) -> Self:
        """
        Resolve python_version from language_version.
        """
        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 TestRepairResponse(BaseModel):
    """
    Response from POST /ai/testgen_repair.
    """

    __test__ = False

    generated_tests: str
