from __future__ import annotations

import logging
import uuid
from pathlib import Path
from typing import TYPE_CHECKING, 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.llm._client import LLMOutputUnparseableError
from codeflash_api.llm._models import CODE_REPAIR_MODEL
from codeflash_api.optimize._context import validate_trace_id
from codeflash_api.optimize.schemas import OptimizeResponseItem
from codeflash_api.repair._context import (
    apply_patches_to_optimized_code,
    build_user_prompt,
    is_valid_repair,
)
from codeflash_api.repair.schemas import (
    CodeRepairErrorResponse,
    CodeRepairRequest,
)

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 / "system_prompt.md").read_text()
USER_PROMPT = (_PROMPTS_DIR / "user_prompt.md").read_text()

MAX_REPAIR_ATTEMPTS = 2


async def _call_repair_llm(  # noqa: PLR0913
    llm_client: LLMClient,
    *,
    user_id: str,
    system_prompt: str,
    user_prompt: str,
    model: LLM,
    trace_id: str,
) -> str:
    """
    Call the LLM for code repair and return the raw response content.
    """
    messages: list[dict[str, str]] = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt},
    ]
    response = await llm_client.call(
        model,
        messages,
        call_type="code_repair",
        trace_id=trace_id,
    )
    return response.content


@router.post(
    "/ai/code_repair",
    response_model=OptimizeResponseItem,
    responses={
        400: {"model": CodeRepairErrorResponse},
        422: {"model": CodeRepairErrorResponse},
        500: {"model": CodeRepairErrorResponse},
    },
)
async def code_repair(
    request: Request,
    data: CodeRepairRequest,
    user: Annotated[AuthenticatedUser, Depends(require_auth)],
    _rate: Annotated[None, Depends(check_rate_limit)],
    _usage: Annotated[None, Depends(track_usage)],
) -> OptimizeResponseItem:
    """
    Repair optimised code to match the original's behaviour.
    """
    if not validate_trace_id(data.trace_id):
        raise HTTPException(
            status_code=400,
            detail=("Invalid trace ID. Please provide a valid UUIDv4."),
        )

    user_prompt = build_user_prompt(
        USER_PROMPT,
        original_source_code=data.original_source_code,
        modified_source_code=data.modified_source_code,
        test_diffs=data.test_diffs,
        language=data.language,
    )

    llm_client: LLMClient = request.app.state.llm_client
    last_error: Exception | None = None

    for attempt in range(MAX_REPAIR_ATTEMPTS):
        try:
            llm_response = await _call_repair_llm(
                llm_client,
                user_id=user.user_id,
                system_prompt=SYSTEM_PROMPT,
                user_prompt=user_prompt,
                model=CODE_REPAIR_MODEL,
                trace_id=data.trace_id,
            )
        except (
            LLMOutputUnparseableError,
            ValueError,
            RuntimeError,
        ) as exc:
            log.warning(
                "LLM call failed on attempt %d: trace_id=%s",
                attempt + 1,
                data.trace_id,
                exc_info=True,
            )
            last_error = exc
            continue

        try:
            repaired_code = apply_patches_to_optimized_code(
                llm_response,
                data.modified_source_code,
                data.language,
            )
        except ValueError as exc:
            log.warning(
                "Patch application failed on attempt %d: trace_id=%s",
                attempt + 1,
                data.trace_id,
                exc_info=True,
            )
            last_error = exc
            continue

        if not is_valid_repair(
            repaired_code,
            data.modified_source_code,
            data.language,
        ):
            log.warning(
                "Repaired code failed validation on attempt %d: trace_id=%s",
                attempt + 1,
                data.trace_id,
            )
            last_error = LLMOutputUnparseableError(
                "Repaired code is not valid"
            )
            continue

        new_optimization_id = str(uuid.uuid4())
        return OptimizeResponseItem(
            source_code=repaired_code,
            optimization_id=new_optimization_id,
            parent_id=data.optimization_id,
            explanation=llm_response,
        )

    detail = str(last_error) if last_error else "Repair failed"
    raise HTTPException(status_code=422, detail=detail)
