from __future__ import annotations

import asyncio
import logging
import uuid
from pathlib import Path
from typing import TYPE_CHECKING, Annotated, Any

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._client import LLMOutputUnparseableError
from codeflash_api.llm._models import REFINEMENT_MODEL
from codeflash_api.optimize._context import validate_trace_id
from codeflash_api.refinement._context import (
    BaseRefinerContext,
    RefinementContextData,
)
from codeflash_api.refinement.schemas import (
    RefinementErrorResponse,
    RefinementRequestItem,
    RefinementResponse,
    RefinementResponseItem,
)

if TYPE_CHECKING:
    from collections.abc import Coroutine

    from codeflash_api.llm._client import LLMClient

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()


async def _refine_single(
    llm_client: LLMClient,
    item: RefinementRequestItem,
    ctx: BaseRefinerContext,
    trace_id: str,
    call_sequence: int,
) -> RefinementResponseItem | None:
    """
    Run one refinement: call LLM, apply patches, validate result.
    """
    system_prompt = ctx.get_system_prompt()
    user_prompt = ctx.get_user_prompt()

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

    try:
        output = await llm_client.call(
            llm=REFINEMENT_MODEL,
            messages=messages,
            call_type="refinement",
            trace_id=trace_id,
        )
    except (LLMOutputUnparseableError, RuntimeError, ValueError, OSError):
        log.warning(
            "LLM call failed for optimization_id=%s"
            " trace_id=%s call_sequence=%d",
            item.optimization_id,
            trace_id,
            call_sequence,
            exc_info=True,
        )
        return None

    explanation_tag = extract_xml_tag(output.content, "explanation")
    refined_explanation = (
        explanation_tag.strip() if explanation_tag else output.content
    )

    refined_code = ""
    try:
        refined_code = ctx.apply_patches_to_optimized_code(output.content)
    except (ValueError, KeyError, IndexError):
        log.warning(
            "Patch application failed for optimization_id=%s",
            item.optimization_id,
            exc_info=True,
        )
        refined_code = ""

    if not ctx.is_valid_refinement(refined_code):
        return None

    return RefinementResponseItem(
        source_code=refined_code,
        explanation=refined_explanation,
        optimization_id=str(uuid.uuid4()),
        parent_id=item.optimization_id,
    )


@router.post(
    "/ai/refinement",
    response_model=RefinementResponse,
    responses={
        400: {"model": RefinementErrorResponse},
        422: {"model": RefinementErrorResponse},
        500: {"model": RefinementErrorResponse},
    },
)
async def refine(
    request: Request,
    data: list[RefinementRequestItem],
    user: Annotated[AuthenticatedUser, Depends(require_auth)],
    _rate: Annotated[None, Depends(check_rate_limit)],
    _usage: Annotated[None, Depends(track_usage)],
) -> RefinementResponse:
    """
    Refine optimization candidates for code quality.
    """
    if not data:
        raise HTTPException(
            status_code=400,
            detail="Request list cannot be empty.",
        )

    trace_id = data[0].trace_id
    if not validate_trace_id(trace_id):
        raise HTTPException(
            status_code=400,
            detail=("Invalid trace ID. Please provide a valid UUIDv4."),
        )

    llm_client: LLMClient = request.app.state.llm_client

    coros: list[Coroutine[Any, Any, RefinementResponseItem | None]] = []
    for i, item in enumerate(data):
        ctx_data = RefinementContextData(
            original_source_code=item.original_source_code,
            original_line_profiler_results=(
                item.original_line_profiler_results or None
            ),
            original_code_runtime=item.original_code_runtime,
            read_only_dependency_code=(item.read_only_dependency_code or None),
            optimized_source_code=item.optimized_source_code,
            optimized_line_profiler_results=(
                item.optimized_line_profiler_results or None
            ),
            optimized_code_runtime=item.optimized_code_runtime,
            speedup=item.speedup,
            optimized_explanation=item.optimized_explanation,
            python_version=item.python_version,
            function_references=item.function_references,
            language=item.language,
            language_version=item.language_version,
        )
        ctx = BaseRefinerContext.get_dynamic_context(
            ctx_data=ctx_data,
            base_system_prompt=SYSTEM_PROMPT,
            base_user_prompt=USER_PROMPT,
        )
        sequence = (
            item.call_sequence if item.call_sequence is not None else i + 1
        )
        coros.append(
            _refine_single(
                llm_client,
                item,
                ctx,
                trace_id,
                sequence,
            )
        )

    results = await asyncio.gather(*coros, return_exceptions=True)

    filtered: list[RefinementResponseItem] = []
    seen_codes: set[str] = set()

    for result in results:
        if isinstance(result, BaseException):
            log.warning(
                "Refinement task raised: %s",
                result,
            )
            continue
        if result is None:
            continue

        try:
            ctx.validate_code_syntax(result.source_code)
        except (SyntaxError, ValueError):
            log.warning(
                "Syntax validation failed for refinement %s",
                result.optimization_id,
                exc_info=True,
            )
            continue

        code_key = result.source_code.strip()
        if code_key and code_key not in seen_codes:
            seen_codes.add(code_key)
            filtered.append(result)

    if not filtered:
        raise HTTPException(
            status_code=422,
            detail=("Could not generate any refinements. Please try again."),
        )

    return RefinementResponse(refinements=filtered)
