"""POST /ai/adaptive_optimize endpoint."""

from __future__ import annotations

import asyncio
import logging
from pathlib import Path
from typing import Annotated

from fastapi import APIRouter, Depends, HTTPException, Request

from codeflash_api.adaptive.schemas import (
    AdaptiveOptimizeRequest,  # noqa: TC001
)
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.diff._base import DiffMethod
from codeflash_api.llm._models import ADAPTIVE_OPTIMIZE_MODEL
from codeflash_api.optimize._context import (
    MultiOptimizerContext,
    validate_trace_id,
)
from codeflash_api.optimize.schemas import (
    OptimizeErrorResponse,
    OptimizeResponse,
    OptimizeResponseItem,
)

log = logging.getLogger(__name__)

router = APIRouter()

_PROMPTS_DIR = Path(__file__).parent / "prompts"
SYSTEM_PROMPT = (_PROMPTS_DIR / "system_prompt.md").read_text()
USER_PROMPT = (_PROMPTS_DIR / "user_prompt.md").read_text()


def _build_attempts_str(
    data: AdaptiveOptimizeRequest,
) -> str:
    """
    Format previous optimization attempts for the user prompt.
    """
    return "".join(
        f"\n\n### Attempt {c.optimization_id}\n\n"
        f"### Source Code:\n{c.source_code}\n\n"
        f"### Explanation:\n{c.explanation}\n\n"
        f"**\n{c.speedup}\n**"
        for c in data.candidates
    )


@router.post(
    "/ai/adaptive_optimize",
    response_model=OptimizeResponse,
    responses={
        400: {"model": OptimizeErrorResponse},
        422: {"model": OptimizeErrorResponse},
        500: {"model": OptimizeErrorResponse},
    },
)
async def adaptive_optimize(
    request: Request,
    data: AdaptiveOptimizeRequest,
    user: Annotated[AuthenticatedUser, Depends(require_auth)],
    _rate: Annotated[None, Depends(check_rate_limit)],
    _usage: Annotated[None, Depends(track_usage)],
) -> OptimizeResponse:
    """
    Produce an adaptive optimization from previous candidates.
    """
    if not validate_trace_id(data.trace_id):
        raise HTTPException(
            status_code=400,
            detail=("Invalid trace ID. Please provide a valid UUIDv4."),
        )

    if not data.candidates:
        raise HTTPException(
            status_code=400,
            detail="At least one candidate is required.",
        )

    python_version_str = data.python_version or "3.12"

    system_prompt = SYSTEM_PROMPT.format(
        python_version_str=python_version_str,
    )
    attempts_str = _build_attempts_str(data)
    user_prompt = USER_PROMPT.format(
        ORIGINAL_CODE=data.original_source_code,
        OPTIMIZATION_ATTEMPTS=attempts_str,
    )

    ctx = MultiOptimizerContext(
        base_system_prompt=system_prompt,
        base_user_prompt=user_prompt,
        source_code=data.original_source_code,
        diff_method=DiffMethod.NO_DIFF,
    )

    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=ADAPTIVE_OPTIMIZE_MODEL,
            messages=messages,
            call_type="adaptive_optimize",
            trace_id=data.trace_id,
        )
    except Exception:
        log.exception(
            "adaptive_optimize LLM call failed: trace_id=%s, user=%s",
            data.trace_id,
            user.user_id,
        )
        raise HTTPException(
            status_code=422,
            detail=(
                "Could not generate adaptive optimization. Please try again."
            ),
        ) from None

    try:
        ctx.extract_code_and_explanation_from_llm_res(
            output.content,
        )
        new_opt = await asyncio.to_thread(
            ctx.parse_and_generate_candidate_schema,
        )
    except Exception:
        log.exception(
            "adaptive_optimize parsing failed: trace_id=%s",
            data.trace_id,
        )
        raise HTTPException(
            status_code=422,
            detail=(
                "Could not parse adaptive optimization"
                " result. Please try again."
            ),
        ) from None

    if new_opt is None or not ctx.is_valid_code():
        log.warning(
            "adaptive_optimize invalid code: trace_id=%s",
            data.trace_id,
        )
        raise HTTPException(
            status_code=422,
            detail=(
                "Could not generate a valid adaptive"
                " optimization. Please try again."
            ),
        )

    parent_id = data.candidates[-1].optimization_id
    result = OptimizeResponseItem(
        source_code=new_opt.source_code,
        explanation=new_opt.explanation,
        optimization_id=new_opt.optimization_id,
        parent_id=parent_id,
    )

    return OptimizeResponse(optimizations=[result])
