"""POST /ai/testgen endpoint."""

from __future__ import annotations

import logging
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.optimize._context import validate_trace_id
from codeflash_api.testgen._generate import generate_tests
from codeflash_api.testgen._validate import CodeValidationError

from .schemas import (
    TestGenErrorResponse,
    TestGenRequest,
    TestGenResponse,
)

log = logging.getLogger(__name__)

router = APIRouter()


@router.post(
    "/ai/testgen",
    response_model=TestGenResponse,
    responses={
        400: {"model": TestGenErrorResponse},
        422: {"model": TestGenErrorResponse},
        500: {"model": TestGenErrorResponse},
    },
)
async def testgen(
    request: Request,
    data: TestGenRequest,
    user: Annotated[AuthenticatedUser, Depends(require_auth)],
    _rate: Annotated[None, Depends(check_rate_limit)],
    _usage: Annotated[None, Depends(track_usage)],
) -> TestGenResponse:
    """
    Generate tests for the given source code.
    """
    if not validate_trace_id(data.trace_id):
        raise HTTPException(
            status_code=400,
            detail="Invalid trace ID. Please provide a valid UUIDv4.",
        )

    if data.language != "python":
        raise HTTPException(
            status_code=400,
            detail="Only Python is currently supported.",
        )

    if data.test_framework not in ("unittest", "pytest"):
        raise HTTPException(
            status_code=400,
            detail=(
                "Invalid test framework. We only support unittest and pytest."
            ),
        )

    python_version = data.python_version or data.language_version
    if not python_version:
        raise HTTPException(
            status_code=400,
            detail="Python version is required.",
        )

    llm_client = request.app.state.llm_client
    test_index = data.test_index if data.test_index is not None else 0

    try:
        generated_tests = await generate_tests(
            llm_client,
            source_code=data.source_code_being_tested,
            function_name=data.function_name,
            test_framework=data.test_framework,
            python_version_str=python_version,
            is_async=data.is_async,
            trace_id=data.trace_id,
            user_id=user.user_id,
            test_index=test_index,
        )

        return TestGenResponse(
            generated_tests=generated_tests,
        )

    except CodeValidationError as e:
        log.exception("Test generation failed: trace_id=%s", data.trace_id)
        raise HTTPException(
            status_code=422,
            detail=str(e),
        ) from None

    except Exception:
        log.exception("Test generation failed: trace_id=%s", data.trace_id)
        raise HTTPException(
            status_code=500,
            detail="Error generating tests. Internal server error.",
        ) from None
