from __future__ import annotations

import json
import logging
from pathlib import Path
from typing import TYPE_CHECKING, Annotated

from fastapi import APIRouter, Depends, HTTPException, Request
from packaging import version

from codeflash_api.auth._deps import (
    check_rate_limit,
    require_auth,
    track_usage,
)
from codeflash_api.auth.models import AuthenticatedUser  # noqa: TC001
from codeflash_api.languages.python._markdown import (
    extract_code_block_with_context,
    wrap_code_in_markdown,
)
from codeflash_api.llm._client import LLMOutputUnparseableError
from codeflash_api.llm._models import OPTIMIZATION_REVIEW_MODEL
from codeflash_api.optimize._context import validate_trace_id
from codeflash_api.review.schemas import (
    OptimizationReviewErrorResponse,
    OptimizationReviewRequest,
    OptimizationReviewResponse,
    ReviewLevel,
)

if TYPE_CHECKING:
    from codeflash_api.llm._client import LLMClient
    from codeflash_api.llm._models import LLM

log = logging.getLogger(__name__)

router = APIRouter()

_PROMPTS_DIR = Path(__file__).parent / "prompts"
SYSTEM_PROMPT = (
    _PROMPTS_DIR / "optimization_review_system_prompt.md"
).read_text()
USER_PROMPT_TEMPLATE = (_PROMPTS_DIR / "user_prompt_template.md").read_text()

_MAX_RETRIES = 2


def _build_review_messages(
    data: OptimizationReviewRequest,
) -> list[dict[str, str]]:
    """
    Build system and user messages for the review LLM call.
    """
    if (
        version.parse(data.codeflash_version) <= version.parse("0.18.2")
        and data.generated_tests
    ):
        data.generated_tests = wrap_code_in_markdown(data.generated_tests)

    system_message = {"role": "system", "content": SYSTEM_PROMPT}

    if data.original_code and data.optimized_code:
        if data.language == "python":
            language_version = data.python_version or "Not specified"
        else:
            language_version = "Not specified"

        user_prompt = USER_PROMPT_TEMPLATE.format(
            original_runtime=data.original_runtime,
            optimized_runtime=data.optimized_runtime,
            speedup=data.speedup,
            loop_count=data.loop_count,
            coverage_message=data.coverage_message,
            language=data.language,
            language_version=language_version,
            original_code=data.original_code,
            optimized_code=data.optimized_code,
            explanation=data.explanation,
            generated_tests=data.generated_tests or "Not Available",
            existing_tests=data.existing_tests or "Not Available",
            replay_tests=data.replay_tests or "Not Available",
            benchmark_details=(data.benchmark_details or "Not Available"),
            calling_fn_details=(data.calling_fn_details or "Not Available"),
            acceptance_reason=data.acceptance_reason or "runtime",
            original_concurrency_ratio=(
                data.original_concurrency_ratio or "N/A"
            ),
            optimized_concurrency_ratio=(
                data.optimized_concurrency_ratio or "N/A"
            ),
            original_throughput=(data.original_throughput or "N/A"),
            optimized_throughput=(data.optimized_throughput or "N/A"),
        )
    elif data.code_diff:
        language_version = (
            data.python_version or "Not specified"
            if data.language == "python"
            else "Not specified"
        )
        user_prompt = f"""
<optimization_review_request>

<performance_metrics>
<original_runtime>{data.original_runtime}</original_runtime>
<optimized_runtime>{data.optimized_runtime}</optimized_runtime>
<speedup_percentage>{data.speedup}</speedup_percentage>
<loop_count>{data.loop_count}</loop_count>
<test_coverage>{data.coverage_message}</test_coverage>
<language>{data.language}</language>
<language_version>{language_version}</language_version>
</performance_metrics>

<code_changes>
{data.code_diff}
</code_changes>

<explanation>
{data.explanation}
</explanation>

<validation_data>
<generated_tests>
{data.generated_tests or "Not Available"}
</generated_tests>

<existing_tests>
{data.existing_tests or "Not Available"}
</existing_tests>

<replay_tests>
{data.replay_tests or "Not Available"}
</replay_tests>

<benchmark_details>
{data.benchmark_details or "Not Available"}
</benchmark_details>
</validation_data>

<codebase_context>
<calling_functions>
{data.calling_fn_details or "Not Available"}
</calling_functions>
</codebase_context>

</optimization_review_request>

Please analyze this optimization following the structured \
assessment process and provide your rating.
"""
    else:
        msg = (
            "Must provide either (original_code + optimized_code) or code_diff"
        )
        raise ValueError(msg)

    user_message = {"role": "user", "content": user_prompt}
    return [system_message, user_message]


def _format_scores_summary(
    review_data: dict[str, object],
) -> str:
    """
    Build the dimension-scores markdown summary from parsed review JSON.
    """
    performance_score = review_data.get("performance_score")
    code_quality_score = review_data.get("code_quality_score")
    hot_path_score = review_data.get("hot_path_score")

    if not any([performance_score, code_quality_score, hot_path_score]):
        return ""

    parts = ["\n\n---\n**Dimension Scores:**"]
    if performance_score:
        parts.append(f"\n- **Performance: {performance_score}/10**")
        reasoning = review_data.get("performance_reasoning", "")
        if reasoning:
            parts.append(f"\n  {reasoning}")
    if code_quality_score:
        parts.append(f"\n- **Code Quality: {code_quality_score}/10**")
        reasoning = review_data.get("code_quality_reasoning", "")
        if reasoning:
            parts.append(f"\n  {reasoning}")
    if hot_path_score:
        parts.append(f"\n- **Hot Path: {hot_path_score}/10**")
        reasoning = review_data.get("hot_path_reasoning", "")
        if reasoning:
            parts.append(f"\n  {reasoning}")
    return "".join(parts)


def _parse_review_response(
    review_text: str,
    cost: float,
) -> OptimizationReviewResponse:
    """
    Extract the JSON rating block from LLM output and build the response.
    """
    result = extract_code_block_with_context(review_text, language="json")
    if result is None:
        msg = "No JSON block found in review response"
        raise LLMOutputUnparseableError(msg, cost=cost)

    before_text, json_content, after_text = result
    try:
        review_data = json.loads(json_content.lower().strip())
    except (json.JSONDecodeError, ValueError) as exc:
        msg = "Invalid JSON in review response"
        raise LLMOutputUnparseableError(msg, cost=cost) from exc

    try:
        review_level = ReviewLevel(review_data["rating"])
    except (KeyError, ValueError) as exc:
        msg = "Missing or invalid 'rating' in review JSON"
        raise LLMOutputUnparseableError(msg, cost=cost) from exc

    review_explanation = (
        before_text + " " + after_text + _format_scores_summary(review_data)
    )

    return OptimizationReviewResponse(
        review=review_level,
        review_explanation=review_explanation,
        performance_score=review_data.get("performance_score"),
        code_quality_score=review_data.get("code_quality_score"),
        hot_path_score=review_data.get("hot_path_score"),
    )


async def get_optimization_review(
    llm_client: LLMClient,
    data: OptimizationReviewRequest,
    *,
    model: LLM = OPTIMIZATION_REVIEW_MODEL,
) -> tuple[OptimizationReviewResponse, float]:
    """
    Call the LLM and parse the optimization review.

    Retries up to *_MAX_RETRIES* times on unparseable output.
    Returns *(response, cost)*.
    """
    messages = _build_review_messages(data)
    total_cost = 0.0
    last_error: LLMOutputUnparseableError | None = None

    for _attempt in range(_MAX_RETRIES):
        call_cost = 0.0
        try:
            llm_response = await llm_client.call(
                llm=model,
                messages=messages,
                call_type="optimization_review",
                trace_id=data.trace_id,
            )
            call_cost = llm_response.cost
            total_cost += call_cost
            review_text = llm_response.content.strip()
            return (
                _parse_review_response(review_text, call_cost),
                total_cost,
            )
        except LLMOutputUnparseableError as exc:
            if call_cost == 0.0:
                total_cost += exc.cost
            last_error = exc
            log.warning(
                "Unparseable review response (attempt %d): trace_id=%s: %s",
                _attempt + 1,
                data.trace_id,
                exc,
            )

    raise last_error  # type: ignore[misc]


@router.post(
    "/ai/optimization_review",
    response_model=OptimizationReviewResponse,
    responses={
        400: {"model": OptimizationReviewErrorResponse},
        422: {"model": OptimizationReviewErrorResponse},
        500: {"model": OptimizationReviewErrorResponse},
    },
)
async def optimization_review(
    request: Request,
    data: OptimizationReviewRequest,
    user: Annotated[AuthenticatedUser, Depends(require_auth)],
    _rate: Annotated[None, Depends(check_rate_limit)],
    _usage: Annotated[None, Depends(track_usage)],
) -> OptimizationReviewResponse:
    """
    Review an optimization and return a quality rating.
    """
    if not validate_trace_id(data.trace_id):
        raise HTTPException(
            status_code=400,
            detail=("Invalid trace ID. Please provide a valid UUIDv4."),
        )

    llm_client: LLMClient = request.app.state.llm_client

    try:
        review_response, _cost = await get_optimization_review(
            llm_client, data
        )
    except LLMOutputUnparseableError as exc:
        log.warning(
            "Review unparseable after retries: trace_id=%s: %s",
            data.trace_id,
            exc,
        )
        raise HTTPException(
            status_code=422,
            detail="Invalid response",
        ) from exc
    except ValueError as exc:
        raise HTTPException(
            status_code=400,
            detail=str(exc),
        ) from exc
    except Exception as exc:
        log.exception(
            "Error in optimization_review: trace_id=%s",
            data.trace_id,
        )
        raise HTTPException(
            status_code=500,
            detail="Internal server error",
        ) from exc

    return review_response
