"""Test repair endpoint.

Takes a generated test file with specific flagged functions and returns a
repaired version with those functions rewritten.

Currently enabled for Python only.
"""

from __future__ import annotations

import ast
import logging
from pathlib import Path
from typing import Any

import sentry_sdk
from jinja2 import Environment, FileSystemLoader, StrictUndefined
from ninja import NinjaAPI
from ninja.errors import HttpError
from openai.types.chat import (
    ChatCompletionAssistantMessageParam,
    ChatCompletionMessageParam,
    ChatCompletionSystemMessageParam,
    ChatCompletionUserMessageParam,
)

from aiservice.analytics.posthog import ph
from aiservice.common.markdown_utils import extract_code_block
from aiservice.common_utils import validate_trace_id
from aiservice.llm import llm_client
from aiservice.llm_models import HAIKU_MODEL
from authapp.auth import AuthenticatedRequest
from core.shared.testgen_review.models import TestRepairErrorSchema, TestRepairResponseSchema, TestRepairSchema

testgen_repair_api = NinjaAPI(urls_namespace="testgen_repair")

_prompts_dir = Path(__file__).parent / "prompts"
_jinja_env = Environment(loader=FileSystemLoader(_prompts_dir), keep_trailing_newline=True, undefined=StrictUndefined)  # noqa: S701


def _extract_and_validate(llm_output: str) -> tuple[str | None, Any]:
    """Extract a Python code block from LLM output and validate its syntax.

    Returns (code_string, parsed_cst) on success, or (code_string, None) if
    extraction succeeded but syntax is invalid, or (None, None) if no code
    block was found.
    """
    from core.languages.python.cst_utils import parse_module_to_cst  # noqa: PLC0415

    code = extract_code_block(llm_output)
    if code is None:
        return None, None
    try:
        ast.parse(code)
    except SyntaxError:
        return code, None
    try:
        return code, parse_module_to_cst(code)
    except Exception:  # noqa: BLE001
        return code, None


def _get_syntax_error(code: str) -> str:
    try:
        ast.parse(code)
    except SyntaxError as e:
        return f"{e.msg} (line {e.lineno})"
    return "Unknown syntax error"


@testgen_repair_api.post(
    "/", response={200: TestRepairResponseSchema, 400: TestRepairErrorSchema, 500: TestRepairErrorSchema}
)
async def testgen_repair(
    request: AuthenticatedRequest, data: TestRepairSchema
) -> tuple[int, TestRepairResponseSchema | TestRepairErrorSchema]:
    if data.language != "python":
        return 400, TestRepairErrorSchema(error="Test repair is only supported for Python")

    if not validate_trace_id(data.trace_id):
        return 400, TestRepairErrorSchema(error="Invalid trace ID. Please provide a valid UUIDv4.")

    ph(request.user, "aiservice-testgen-repair-called")

    try:
        from core.shared.testgen_review.review import _build_coverage_context  # noqa: PLC0415

        # module_source_code is the full file; function_source_code is just the function for the LLM prompt
        module_source = data.module_source_code or data.function_source_code

        coverage_context = (
            _build_coverage_context(data.coverage_details, data.function_source_code) if data.coverage_details else ""
        )

        _codeflash_env_vars = ("CODEFLASH_TRACER_DISABLE", "CODEFLASH_LOOP_INDEX", "CODEFLASH_TEST_ITERATION")
        has_codeflash_env_vars = any(var in module_source for var in _codeflash_env_vars)
        system_prompt = _jinja_env.get_template("testgen_repair_system_prompt.md.j2").render(
            has_codeflash_env_vars=has_codeflash_env_vars
        )
        user_prompt = _jinja_env.get_template("testgen_repair_user_prompt.md.j2").render(
            function_name=data.function_to_optimize.function_name,
            function_source_code=data.function_source_code,
            test_source=data.test_source,
            functions_to_repair=data.functions_to_repair,
            coverage_context=coverage_context,
            previous_repair_errors=data.previous_repair_errors,
        )

        messages: list[ChatCompletionMessageParam] = [
            ChatCompletionSystemMessageParam(role="system", content=system_prompt),
            ChatCompletionUserMessageParam(role="user", content=user_prompt),
        ]

        obs_context: dict[str, Any] = {}
        if data.call_sequence is not None:
            obs_context["call_sequence"] = data.call_sequence

        response = await llm_client.call(
            llm=HAIKU_MODEL,
            messages=messages,
            call_type="testgen_repair",
            trace_id=data.trace_id,
            user_id=request.user,
            context=obs_context,
        )

        cost = response.cost
        logging.debug("testgen_repair LLM cost: %s", cost)

        from core.languages.python.testgen.instrumentation.edit_generated_test import (  # noqa: PLC0415
            find_and_replace_function,
        )
        from core.languages.python.testgen.postprocessing.postprocess_pipeline import (  # noqa: PLC0415
            postprocessing_testgen_pipeline,
        )
        from core.languages.python.testgen.validate import instrument_tests, validate_request_data  # noqa: PLC0415
        from core.shared.testgen_models import TestGenSchema  # noqa: PLC0415

        repaired_code, repaired_cst = _extract_and_validate(response.content.strip())

        if repaired_cst is None:
            syntax_error = _get_syntax_error(repaired_code) if repaired_code else "No code block found"
            logging.info("Repair attempt produced invalid syntax, retrying: %s", syntax_error)
            messages.append(ChatCompletionAssistantMessageParam(role="assistant", content=response.content))
            messages.append(
                ChatCompletionUserMessageParam(
                    role="user",
                    content=f"The code you returned has a syntax error:\n```\n{syntax_error}\n```\n"
                    "Please return the complete corrected file in a single ```python code block.",
                )
            )
            retry_response = await llm_client.call(
                llm=HAIKU_MODEL,
                messages=messages,
                call_type="testgen_repair_retry",
                trace_id=data.trace_id,
                user_id=request.user,
                context=obs_context,
            )
            cost += retry_response.cost
            repaired_code, repaired_cst = _extract_and_validate(retry_response.content.strip())

        if repaired_cst is None:
            logging.warning("LLM returned syntactically invalid code after retry")
            return 500, TestRepairErrorSchema(error="LLM returned syntactically invalid code")
        from core.languages.python.cst_utils import parse_module_to_cst  # noqa: PLC0415

        original_cst = parse_module_to_cst(data.test_source)

        # Extract repaired function nodes by name
        repaired_functions: dict[str, Any] = {}
        for node in repaired_cst.body:
            if hasattr(node, "body") and hasattr(node.body, "__iter__"):
                for child in node.body:
                    if hasattr(child, "name") and hasattr(child.name, "value"):
                        repaired_functions[child.name.value] = child
            if hasattr(node, "name") and hasattr(node.name, "value"):
                repaired_functions[node.name.value] = node

        # Replace only flagged functions in the original
        merged_cst = original_cst
        for func_to_repair in data.functions_to_repair:
            if func_to_repair.function_name in repaired_functions:
                merged_cst = find_and_replace_function(
                    merged_cst, func_to_repair.function_name, repaired_functions[func_to_repair.function_name]
                )

        # Run postprocessing and instrumentation on the merged result
        testgen_data = TestGenSchema(
            source_code_being_tested=module_source,
            function_to_optimize=data.function_to_optimize,
            helper_function_names=data.helper_function_names or [],
            module_path=data.module_path,
            test_module_path=data.test_module_path,
            test_framework=data.test_framework,
            test_timeout=data.test_timeout,
            trace_id=data.trace_id,
            language_version=data.language_version,
            language=data.language,
            codeflash_version=data.codeflash_version,
        )

        python_version = validate_request_data(testgen_data)

        processed_cst, display_cst = postprocessing_testgen_pipeline(
            merged_cst, data.helper_function_names or [], data.function_to_optimize, data.module_path, module_source
        )
        processed_code = processed_cst.code
        display_code = display_cst.code

        instrumented_behavior, instrumented_perf = instrument_tests(processed_code, testgen_data, python_version)

        if instrumented_behavior is None or instrumented_perf is None:
            return 500, TestRepairErrorSchema(error="Failed to instrument repaired tests")

        ph(request.user, "aiservice-testgen-repair-completed")
        return 200, TestRepairResponseSchema(
            generated_tests=display_code,
            instrumented_behavior_tests=instrumented_behavior,
            instrumented_perf_tests=instrumented_perf,
        )

    except HttpError:
        raise
    except Exception as e:
        logging.exception("Error in testgen_repair")
        sentry_sdk.capture_exception(e)
        return 500, TestRepairErrorSchema(error="Internal server error")
