"""Request/response schemas for the testgen review and repair endpoints."""

from typing import Self

from ninja import Schema
from pydantic import model_validator

from aiservice.models.functions_to_optimize import FunctionToOptimize

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


class TestSourceWithFailures(Schema):
    test_source: str
    test_index: int
    failed_test_functions: list[str] = []  # functions that failed behavioral tests
    failure_messages: dict[str, str] = {}  # function_name → error message from test run


class CoverageFunctionDetail(Schema):
    name: str
    coverage: float
    executed_lines: list[int] = []
    unexecuted_lines: list[int] = []
    executed_branches: list[list[int]] = []
    unexecuted_branches: list[list[int]] = []


class CoverageDetails(Schema):
    coverage_percentage: float
    threshold_percentage: float
    function_start_line: int = 1
    main_function: CoverageFunctionDetail
    dependent_function: CoverageFunctionDetail | None = None


class TestgenReviewSchema(Schema):
    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(Schema):
    function_name: str
    verdict: str  # "pass" or "repair"
    reason: str = ""


class TestReview(Schema):
    test_index: int
    functions: list[FunctionVerdict]


class TestgenReviewResponseSchema(Schema):
    reviews: list[TestReview]


class TestgenReviewErrorSchema(Schema):
    error: str


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


class FunctionToRepair(Schema):
    function_name: str
    reason: str


class TestRepairSchema(Schema):
    test_source: str
    functions_to_repair: list[FunctionToRepair]
    function_source_code: str
    function_to_optimize: FunctionToOptimize
    helper_function_names: list[str] | None = None
    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 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 TestRepairResponseSchema(Schema):
    generated_tests: str
    instrumented_behavior_tests: str
    instrumented_perf_tests: str


class TestRepairErrorSchema(Schema):
    error: str
