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.diff._base import DiffMethod
from codeflash_api.optimize._context import (
    BaseOptimizerContext,
    parse_python_version,
    validate_trace_id,
)
from codeflash_api.optimize._pipeline import optimize_python_code
from codeflash_api.optimize.schemas import (
    OptimizeErrorResponse,
    OptimizeRequest,
    OptimizeResponse,
)

log = logging.getLogger(__name__)

router = APIRouter()

_PROMPTS_DIR = (
    Path(__file__).parent.parent / "languages" / "python" / "prompts"
)
SYSTEM_PROMPT = (_PROMPTS_DIR / "system_prompt.md").read_text()
USER_PROMPT = (_PROMPTS_DIR / "user_prompt.md").read_text()
ASYNC_SYSTEM_PROMPT = (_PROMPTS_DIR / "async_system_prompt.md").read_text()
ASYNC_USER_PROMPT = (_PROMPTS_DIR / "async_user_prompt.md").read_text()
JIT_INSTRUCTIONS = (_PROMPTS_DIR / "jit_instructions.md").read_text()


def _validate_request(
    data: OptimizeRequest,
    ctx: BaseOptimizerContext,
) -> tuple[int, int, int]:
    """
    Validate request fields and parse source code.
    """
    if not data.source_code:
        raise HTTPException(
            status_code=400,
            detail="Source code cannot be empty.",
        )
    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.python_version:
        raise HTTPException(
            status_code=400,
            detail="Python version is required.",
        )
    try:
        python_version = parse_python_version(data.python_version)
    except ValueError as e:
        raise HTTPException(
            status_code=400,
            detail=(
                "Invalid Python version, it should look like"
                " 3.x.x. We only support Python 3.9 and above."
            ),
        ) from e

    try:
        ctx.validate_and_parse_source_code(
            data.source_code,
            feature_version=python_version[:2],
        )
    except SyntaxError as e:
        raise HTTPException(
            status_code=400,
            detail=(
                "Invalid source code. It is not valid Python"
                " code. Please check syntax of your code."
            ),
        ) from e

    return python_version


@router.post(
    "/ai/optimize",
    response_model=OptimizeResponse,
    responses={
        400: {"model": OptimizeErrorResponse},
        422: {"model": OptimizeErrorResponse},
        500: {"model": OptimizeErrorResponse},
    },
)
async def optimize(
    request: Request,
    data: OptimizeRequest,
    user: Annotated[AuthenticatedUser, Depends(require_auth)],
    _rate: Annotated[None, Depends(check_rate_limit)],
    _usage: Annotated[None, Depends(track_usage)],
) -> OptimizeResponse:
    """
    Optimize Python code for performance using LLMs.
    """
    system_prompt = ASYNC_SYSTEM_PROMPT if data.is_async else SYSTEM_PROMPT
    user_prompt = ASYNC_USER_PROMPT if data.is_async else USER_PROMPT
    if data.is_numerical_code:
        system_prompt += f"\n{JIT_INSTRUCTIONS}\n"

    ctx = BaseOptimizerContext.get_dynamic_context(
        system_prompt,
        user_prompt,
        data.source_code,
        DiffMethod.NO_DIFF,
    )

    try:
        python_version = _validate_request(data, ctx)
    except HTTPException:
        raise
    except Exception as e:
        log.exception("Validation error: trace_id=%s", data.trace_id)
        raise HTTPException(
            status_code=400,
            detail="Request validation failed.",
        ) from e

    llm_client = request.app.state.llm_client

    (
        optimization_results,
        _llm_cost,
        _code_and_explanations,
        _optimization_models,
    ) = await optimize_python_code(
        llm_client,
        user_id=user.user_id,
        ctx=ctx,
        trace_id=data.trace_id,
        original_source_code=data.source_code,
        dependency_code=data.dependency_code,
        python_version=python_version,
        n_candidates=data.n_candidates,
        baseline_runtime_ns=data.baseline_runtime_ns,
        loop_count=data.loop_count,
        line_profiler_results=data.line_profiler_results,
        test_input_examples=data.test_input_examples,
    )

    if len(optimization_results) == 0:
        log.warning(
            "No optimizations found: trace_id=%s, repo=%s/%s, n_candidates=%d",
            data.trace_id,
            data.repo_owner,
            data.repo_name,
            data.n_candidates,
        )
        raise HTTPException(
            status_code=422,
            detail=("Could not generate any optimizations. Please try again."),
        )

    return OptimizeResponse(optimizations=optimization_results)
