from __future__ import annotations

from enum import Enum

from pydantic import BaseModel


class ReviewLevel(str, Enum):
    """
    Quality rating for an optimization review.
    """

    HIGH = "high"
    MEDIUM = "medium"
    LOW = "low"


class OptimizationReviewRequest(BaseModel):
    """
    Request body for the optimization review endpoint.
    """

    trace_id: str
    code_diff: str | None = None
    original_code: str | None = None
    optimized_code: str | None = None
    original_runtime: str
    optimized_runtime: str
    speedup: str
    existing_tests: str | None = None
    generated_tests: str | None = None
    replay_tests: str | None = None
    benchmark_details: str | None = None
    coverage_message: str
    loop_count: int
    explanation: str
    calling_fn_details: str
    language: str = "python"
    python_version: str | None = None
    codeflash_version: str = "0.18.2"
    call_sequence: int | None = None
    acceptance_reason: str | None = None
    original_concurrency_ratio: float | None = None
    optimized_concurrency_ratio: float | None = None
    original_throughput: int | None = None
    optimized_throughput: int | None = None


class OptimizationReviewResponse(BaseModel):
    """
    Successful response from the optimization review endpoint.
    """

    review: ReviewLevel
    review_explanation: str
    performance_score: int | None = None
    code_quality_score: int | None = None
    hot_path_score: int | None = None


class OptimizationReviewErrorResponse(BaseModel):
    """
    Error response from the optimization review endpoint.
    """

    error: str
