"""POST /ai/testgen_review and POST /ai/testgen_repair endpoints."""

from __future__ import annotations

import ast
import asyncio
import json
import logging
from pathlib import Path
from typing import TYPE_CHECKING, Annotated, Any

import jinja2
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._markdown import (
    extract_code_block_with_context,
)
from codeflash_api.llm._models import ANTHROPIC_CLAUDE_HAIKU_4_5
from codeflash_api.optimize._context import validate_trace_id
from codeflash_api.testgen.schemas import (
    CoverageDetails,
    FunctionVerdict,
    TestGenErrorResponse,
    TestgenReviewRequest,
    TestgenReviewResponse,
    TestRepairRequest,
    TestRepairResponse,
    TestReview,
    TestSourceWithFailures,
)

if TYPE_CHECKING:
    from codeflash_api.llm._client import LLMClient

log = logging.getLogger(__name__)

router = APIRouter()

_PROMPTS_DIR = Path(__file__).parent / "prompts"
_JINJA_ENV = jinja2.Environment(  # noqa: S701
    loader=jinja2.FileSystemLoader(str(_PROMPTS_DIR)),
    keep_trailing_newline=True,
)

REVIEW_SYSTEM_TPL = _JINJA_ENV.get_template("review_system.md.j2")
REVIEW_USER_TPL = _JINJA_ENV.get_template("review_user.md.j2")
REPAIR_SYSTEM_TPL = _JINJA_ENV.get_template("repair_system.md.j2")
REPAIR_USER_TPL = _JINJA_ENV.get_template("repair_user.md.j2")

REVIEW_MODEL = ANTHROPIC_CLAUDE_HAIKU_4_5
REPAIR_MODEL = ANTHROPIC_CLAUDE_HAIKU_4_5

_MAX_REPAIR_RETRIES = 2


def _build_coverage_context(
    details: CoverageDetails | None,
) -> str:
    """
    Format coverage details for prompt injection.
    """
    if details is None:
        return ""

    parts: list[str] = []
    mf = details.main_function
    parts.append(f"Function `{mf.name}`: {mf.coverage:.0f}% coverage")
    if mf.unexecuted_lines:
        parts.append(f"  Unexecuted lines: {mf.unexecuted_lines}")
    if mf.unexecuted_branches:
        parts.append(f"  Unexecuted branches: {mf.unexecuted_branches}")

    if details.dependent_function:
        df = details.dependent_function
        parts.append(
            f"Dependent function `{df.name}`: {df.coverage:.0f}% coverage"
        )
        if df.unexecuted_lines:
            parts.append(f"  Unexecuted lines: {df.unexecuted_lines}")

    return "\n".join(parts)


def _parse_review_verdicts(
    content: str,
) -> list[FunctionVerdict]:
    """
    Parse function verdicts from LLM JSON response.
    """
    result = extract_code_block_with_context(content, language="json")
    json_str = result[1] if result else content

    try:
        data = json.loads(json_str)
    except (json.JSONDecodeError, ValueError):
        log.warning("Failed to parse review JSON: %s", content[:200])
        return []

    raw_functions: list[dict[str, Any]]
    if isinstance(data, dict):
        raw_functions = data.get("functions", [])
    elif isinstance(data, list):
        raw_functions = data
    else:
        return []

    verdicts: list[FunctionVerdict] = []
    for fn in raw_functions:
        if not isinstance(fn, dict):
            continue
        verdict = fn.get("verdict", "")
        if verdict == "repair":
            verdicts.append(
                FunctionVerdict(
                    function_name=fn.get("function_name", ""),
                    verdict="repair",
                    reason=fn.get("reason", ""),
                )
            )
    return verdicts


async def _review_single_test(
    llm_client: LLMClient,
    test: TestSourceWithFailures,
    function_name: str,
    function_source_code: str,
    trace_id: str,
    coverage_context: str,
    coverage_summary: str,
) -> TestReview:
    """
    Review a single test source and return verdicts.
    """
    failed_verdicts: list[FunctionVerdict] = [
        FunctionVerdict(
            function_name=fn,
            verdict="repair",
            reason=test.failure_messages.get(fn, "Failed behavioral test"),
        )
        for fn in test.failed_test_functions
    ]

    failed_note = ""
    if test.failed_test_functions:
        failed_names = ", ".join(
            f"`{fn}`" for fn in test.failed_test_functions
        )
        failed_note = (
            f"The following functions failed behavioral"
            f" tests and are already marked for repair:"
            f" {failed_names}. Do NOT re-review them."
        )

    system_prompt = REVIEW_SYSTEM_TPL.render()
    user_prompt = REVIEW_USER_TPL.render(
        function_name=function_name,
        function_source_code=function_source_code,
        test_source=test.test_source,
        coverage_context=coverage_context,
        coverage_summary=coverage_summary,
        failed_note=failed_note,
    )

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

    try:
        output = await llm_client.call(
            llm=REVIEW_MODEL,
            messages=messages,
            call_type="testgen_review",
            trace_id=trace_id,
        )
        ai_verdicts = _parse_review_verdicts(output.content)
    except Exception:
        log.exception(
            "Review LLM call failed: trace_id=%s test_index=%d",
            trace_id,
            test.test_index,
        )
        ai_verdicts = []

    already_failed = {v.function_name for v in failed_verdicts}
    combined = failed_verdicts + [
        v for v in ai_verdicts if v.function_name not in already_failed
    ]

    return TestReview(
        test_index=test.test_index,
        functions=combined,
    )


@router.post(
    "/ai/testgen_review",
    response_model=TestgenReviewResponse,
    responses={
        400: {"model": TestGenErrorResponse},
        422: {"model": TestGenErrorResponse},
    },
)
async def testgen_review(
    request: Request,
    data: TestgenReviewRequest,
    user: Annotated[AuthenticatedUser, Depends(require_auth)],
    _rate: Annotated[None, Depends(check_rate_limit)],
    _usage: Annotated[None, Depends(track_usage)],
) -> TestgenReviewResponse:
    """
    Review generated tests for quality issues.
    """
    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":
        return TestgenReviewResponse(reviews=[])

    llm_client: LLMClient = request.app.state.llm_client
    coverage_context = _build_coverage_context(
        data.coverage_details,
    )

    coros = [
        _review_single_test(
            llm_client,
            test,
            data.function_name,
            data.function_source_code,
            data.trace_id,
            coverage_context,
            data.coverage_summary,
        )
        for test in data.tests
    ]

    reviews = await asyncio.gather(*coros)

    return TestgenReviewResponse(reviews=list(reviews))


# ── Repair ──────────────────────────────────────────────────────────


def _extract_python_code(content: str) -> str | None:
    """
    Extract a Python code block from LLM output.
    """
    result = extract_code_block_with_context(content, language="python")
    if result is None:
        return None
    return result[1]


def _get_syntax_error(code: str) -> str:
    """
    Return the syntax error message, or empty string if valid.
    """
    try:
        ast.parse(code)
    except SyntaxError as exc:
        return f"Line {exc.lineno}: {exc.msg}"
    return ""


@router.post(
    "/ai/testgen_repair",
    response_model=TestRepairResponse,
    responses={
        400: {"model": TestGenErrorResponse},
        422: {"model": TestGenErrorResponse},
        500: {"model": TestGenErrorResponse},
    },
)
async def testgen_repair(
    request: Request,
    data: TestRepairRequest,
    user: Annotated[AuthenticatedUser, Depends(require_auth)],
    _rate: Annotated[None, Depends(check_rate_limit)],
    _usage: Annotated[None, Depends(track_usage)],
) -> TestRepairResponse:
    """
    Repair flagged test functions.
    """
    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.",
        )

    llm_client: LLMClient = request.app.state.llm_client
    coverage_context = _build_coverage_context(
        data.coverage_details,
    )

    system_prompt = REPAIR_SYSTEM_TPL.render()
    user_prompt = REPAIR_USER_TPL.render(
        function_name=data.function_name,
        function_source_code=data.function_source_code,
        test_source=data.test_source,
        functions_to_repair=data.functions_to_repair,
        previous_repair_errors=data.previous_repair_errors,
        coverage_context=coverage_context,
    )

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

    repaired_code: str | None = None

    for attempt in range(_MAX_REPAIR_RETRIES):
        try:
            output = await llm_client.call(
                llm=REPAIR_MODEL,
                messages=messages,
                call_type=(
                    "testgen_repair"
                    if attempt == 0
                    else "testgen_repair_retry"
                ),
                trace_id=data.trace_id,
            )
        except Exception:
            log.exception(
                "Repair LLM call failed: trace_id=%s attempt=%d",
                data.trace_id,
                attempt + 1,
            )
            raise HTTPException(
                status_code=422,
                detail="Could not generate repair. Please try again.",
            ) from None

        repaired_code = _extract_python_code(output.content)
        if repaired_code is None:
            log.warning(
                "No code block in repair response: trace_id=%s attempt=%d",
                data.trace_id,
                attempt + 1,
            )
            continue

        syntax_err = _get_syntax_error(repaired_code)
        if not syntax_err:
            break

        log.warning(
            "Syntax error in repair: trace_id=%s attempt=%d error=%s",
            data.trace_id,
            attempt + 1,
            syntax_err,
        )
        messages.append({"role": "assistant", "content": output.content})
        messages.append(
            {
                "role": "user",
                "content": (
                    f"Syntax error:\n```\n{syntax_err}\n```\n"
                    "Please return the complete corrected file"
                    " in a single Python code block."
                ),
            }
        )
        repaired_code = None

    if repaired_code is None:
        raise HTTPException(
            status_code=422,
            detail=(
                "Could not generate a valid repair"
                " after retries. Please try again."
            ),
        )

    return TestRepairResponse(
        generated_tests=repaired_code,
    )
