"""POST /ai/explain endpoint."""

from __future__ import annotations

import logging
from pathlib import Path
from typing import Annotated

from fastapi import APIRouter, Depends, HTTPException, Request

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._xml import extract_xml_tag
from codeflash_api.llm._models import EXPLANATIONS_MODEL
from codeflash_api.optimize._context import validate_trace_id

from .schemas import (
    ExplainErrorResponse,
    ExplainRequest,
    ExplainResponse,
)

log = logging.getLogger(__name__)

router = APIRouter()

_PROMPTS_DIR = Path(__file__).parent / "prompts"
SYSTEM_PROMPT_TEMPLATE = (_PROMPTS_DIR / "system_prompt.md").read_text()
USER_PROMPT_TEMPLATE = (_PROMPTS_DIR / "user_prompt.md").read_text()

_THROUGHPUT_CONTEXT = (
    "\nThroughput data (operations per second) is provided"
    " for both versions. When explaining, cover how the"
    " optimization affects processing rate. For async"
    " functions with both runtime and throughput data,"
    " explain both metrics.\n"
)
_CONCURRENCY_CONTEXT = (
    "\nConcurrency ratio data is provided -- it measures"
    " how well async code scales under concurrent"
    " execution (1.0x = no benefit, higher = better)."
    " Blocking code (e.g., `time.sleep()`) scores low;"
    " non-blocking code (e.g., `asyncio.sleep()`) scores"
    " high. When the ratio improves, the optimized code"
    " better yields to the event loop.\n"
)
_ACCEPTANCE_CRITERIA = (
    "\nThe optimization was accepted for improvements"
    " in **{acceptance_reason}**. Lead with that metric"
    " as the primary benefit. If other metrics regressed,"
    " frame them as a reasonable trade-off. Do not frame"
    " the optimization negatively.\n"
)


def _build_system_prompt(data: ExplainRequest) -> str:
    """
    Build the system prompt with optional context sections.
    """
    include_throughput = (
        data.original_throughput is not None
        and data.optimized_throughput is not None
    )
    include_concurrency = (
        data.original_concurrency_ratio is not None
        and data.optimized_concurrency_ratio is not None
    )

    throughput_context = _THROUGHPUT_CONTEXT if include_throughput else ""
    concurrency_context = _CONCURRENCY_CONTEXT if include_concurrency else ""
    acceptance_criteria = ""
    if data.acceptance_reason:
        acceptance_criteria = _ACCEPTANCE_CRITERIA.format(
            acceptance_reason=data.acceptance_reason,
        )

    return SYSTEM_PROMPT_TEMPLATE.format(
        throughput_context=throughput_context,
        concurrency_context=concurrency_context,
        acceptance_criteria=acceptance_criteria,
    )


def _build_user_prompt(data: ExplainRequest) -> str:
    """
    Build the user prompt with all data fields.
    """
    include_throughput = (
        data.original_throughput is not None
        and data.optimized_throughput is not None
    )
    include_concurrency = (
        data.original_concurrency_ratio is not None
        and data.optimized_concurrency_ratio is not None
    )

    throughput_fields = ""
    if include_throughput:
        throughput_fields = (
            f"<original_throughput>\n"
            f"{data.original_throughput}\n"
            f"</original_throughput>\n"
            f"<optimized_throughput>\n"
            f"{data.optimized_throughput}\n"
            f"</optimized_throughput>\n"
        )

    concurrency_fields = ""
    if include_concurrency:
        concurrency_fields = (
            f"<original_concurrency_ratio>\n"
            f"{data.original_concurrency_ratio}\n"
            f"</original_concurrency_ratio>\n"
            f"<optimized_concurrency_ratio>\n"
            f"{data.optimized_concurrency_ratio}\n"
            f"</optimized_concurrency_ratio>\n"
        )

    return USER_PROMPT_TEMPLATE.format(
        original_source_code=data.source_code,
        optimized_source_code=data.optimized_code,
        original_line_profiler_results=(
            data.original_profiler_results or "[No profiler results available]"
        ),
        optimized_line_profiler_results=(
            data.optimized_profiler_results
            or "[No profiler results available]"
        ),
        original_code_runtime=(data.original_runtime or "[Not available]"),
        optimized_code_runtime=(data.optimized_runtime or "[Not available]"),
        speedup=data.speedup or "[Not available]",
        annotated_tests=(data.annotated_tests or "[No tests available]"),
        read_only_dependency_code=(
            data.dependency_code or "[No read only code present]"
        ),
        original_explanation=(data.original_explanation or "[Not available]"),
        python_version=(data.python_version or "Not Available"),
        function_references=(data.function_references or "Not Available"),
        throughput_fields=throughput_fields,
        concurrency_fields=concurrency_fields,
    )


@router.post(
    "/ai/explain",
    response_model=ExplainResponse,
    responses={
        400: {"model": ExplainErrorResponse},
        422: {"model": ExplainErrorResponse},
        500: {"model": ExplainErrorResponse},
    },
)
async def explain(
    request: Request,
    data: ExplainRequest,
    user: Annotated[AuthenticatedUser, Depends(require_auth)],
    _rate: Annotated[None, Depends(check_rate_limit)],
    _usage: Annotated[None, Depends(track_usage)],
) -> ExplainResponse:
    """
    Generate an explanation of why optimized code is faster.
    """
    if not validate_trace_id(data.trace_id):
        raise HTTPException(
            status_code=400,
            detail=("Invalid trace ID. Please provide a valid UUIDv4."),
        )

    system_prompt = _build_system_prompt(data)
    user_prompt = _build_user_prompt(data)

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt},
    ]

    llm_client = request.app.state.llm_client
    try:
        output = await llm_client.call(
            llm=EXPLANATIONS_MODEL,
            messages=messages,
            call_type="explanation",
            trace_id=data.trace_id,
        )
    except Exception:
        log.exception(
            "explain LLM call failed: trace_id=%s, user=%s",
            data.trace_id,
            user.user_id,
        )
        raise HTTPException(
            status_code=422,
            detail=("Error generating explanation. Internal server error."),
        ) from None

    explanation = extract_xml_tag(output.content, "explain")
    if not explanation:
        explanation = output.content.strip()

    if not explanation:
        raise HTTPException(
            status_code=422,
            detail=("Failed to parse explanation from LLM response."),
        )

    return ExplainResponse(explanation=explanation)
