from __future__ import annotations

from pydantic import BaseModel


class RefinementRequestItem(BaseModel):
    """
    A single refinement candidate in the request list.
    """

    trace_id: str
    optimization_id: str
    original_source_code: str
    original_line_profiler_results: str = ""
    read_only_dependency_code: str = ""
    optimized_source_code: str
    optimized_line_profiler_results: str = ""
    optimized_explanation: str = ""
    original_code_runtime: str = ""
    optimized_code_runtime: str = ""
    speedup: str = ""
    python_version: str | None = None
    function_references: str | None = None
    call_sequence: int | None = None
    language: str = "python"
    language_version: str | None = None
    rerun_trace_id: str | None = None


class RefinementResponseItem(BaseModel):
    """
    A single successful refinement in the response.
    """

    source_code: str
    explanation: str
    optimization_id: str
    parent_id: str | None = None


class RefinementResponse(BaseModel):
    """
    Successful response from POST /ai/refinement.
    """

    refinements: list[RefinementResponseItem]


class RefinementErrorResponse(BaseModel):
    """
    Error response from POST /ai/refinement.
    """

    error: str
