from __future__ import annotations

import enum
from typing import Any

from pydantic import BaseModel


class BehaviorDiffScope(str, enum.Enum):
    """
    The dimension on which a test diff was observed.
    """

    RETURN_VALUE = "return_value"
    STDOUT = "stdout"
    DID_PASS = "did_pass"  # noqa: S105
    TIMED_OUT = "timed_out"


class BehaviorDiff(BaseModel):
    """
    A single behavioural difference between original and optimised code.
    """

    scope: BehaviorDiffScope
    original_value: (
        bool | str | int | float | dict[str, Any] | list[Any] | None
    ) = None
    candidate_value: (
        bool | str | int | float | dict[str, Any] | list[Any] | None
    ) = None
    original_pass: bool
    candidate_pass: bool
    test_src_code: str | None = None
    candidate_pytest_error: str | None = None
    original_pytest_error: str | None = None


class CodeRepairRequest(BaseModel):
    """
    Request body for POST /ai/code_repair.
    """

    trace_id: str
    optimization_id: str
    original_source_code: str
    modified_source_code: str
    test_diffs: list[BehaviorDiff]
    language: str = "python"
    rerun_trace_id: str | None = None


class CodeRepairErrorResponse(BaseModel):
    """
    Error response from POST /ai/code_repair.
    """

    error: str
