# Derived from https://github.com/openai/openai-cookbook/blob/main/examples/Unit_test_writing_using_a_multi-step_prompt.ipynb
from __future__ import annotations

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

import libcst as cst
import sentry_sdk
from ninja.errors import HttpError

from aiservice.analytics.posthog import ph
from aiservice.common.markdown_utils import extract_code_block, split_markdown_code
from aiservice.common_utils import safe_isort
from aiservice.env_specific import debug_log_sensitive_data, debug_log_sensitive_data_from_callable
from aiservice.llm import llm_client
from aiservice.llm_models import EXECUTE_MODEL, HAIKU_MODEL, OPENAI_MODEL
from aiservice.models.functions_to_optimize import FunctionToOptimize
from core.languages.python.cst_utils import any_ellipsis_in_cst, ellipsis_in_cst_not_types, parse_module_to_cst
from core.languages.python.optimizer.context_utils.context_helpers import is_multi_context
from core.languages.python.testgen.demo_hacks import try_demo_testgen
from core.languages.python.testgen.models import CostTracker, LLMOutputParseError
from core.languages.python.testgen.postprocessing.code_validator import (
    CodeValidationError,
    has_test_functions,
    validate_testgen_code,
)
from core.languages.python.testgen.postprocessing.postprocess_pipeline import postprocessing_testgen_pipeline
from core.languages.python.testgen.preprocessing.preprocess_pipeline import generate_notes
from core.languages.python.testgen.validate import instrument_tests, validate_request_data
from core.log_features.log_event import update_optimization_cost
from core.log_features.log_features import safe_log_features
from core.shared.jinja_utils import create_prompt_env
from core.shared.testgen_models import (
    TestGenDebugInfo,
    TestGenerationFailedError,
    TestGenErrorResponseSchema,
    TestGenResponseSchema,
)

if TYPE_CHECKING:
    from openai.types.chat import ChatCompletionMessageParam

    from aiservice.llm_models import LLM
    from authapp.auth import AuthenticatedRequest
    from core.shared.testgen_models import TestGenSchema


def select_model_for_test(test_index: int) -> tuple[LLM, str]:
    if test_index % 2 == 0:
        return OPENAI_MODEL, "OpenAI"
    return HAIKU_MODEL, "Anthropic"


_current_dir = Path(__file__).parent
_prompts_dir = _current_dir / "prompts"

_jinja_env = create_prompt_env(_prompts_dir)


def build_prompt(
    *,
    qualified_name: str,
    source_code: str,
    notes: str,
    function_name: str,
    unit_test_package: str,
    is_async: bool,
    is_numerical_code: bool | None = None,
    model_type: str = "openai",
) -> tuple[list[ChatCompletionMessageParam], str, str]:
    system_template = "generate_async_system.md.j2" if is_async else "generate_system.md.j2"

    system_prompt = _jinja_env.get_template(system_template).render(
        function_name=qualified_name, model_type=model_type, is_numerical_code=is_numerical_code
    )
    user_prompt = _jinja_env.get_template("compose_user.md.j2").render(
        notes=notes,
        is_numerical_code=is_numerical_code,
        is_async=is_async,
        unit_test_package=unit_test_package,
        function_name=function_name,
        function_code=source_code,
        package_comment="",
    )

    posthog_event_suffix = "async-" if is_async else ""
    error_context = "async " if is_async else ""
    execute_messages: list[ChatCompletionMessageParam] = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt},
    ]

    return execute_messages, posthog_event_suffix, error_context


def did_generate_ellipsis(source_code: str, generated_code: str) -> bool:
    generated_module = cst.parse_module(generated_code)
    if not ellipsis_in_cst_not_types(generated_module):
        return False
    if is_multi_context(source_code):
        return any(not any_ellipsis_in_cst(cst.parse_module(sc)) for sc in split_markdown_code(source_code).values())
    return not any_ellipsis_in_cst(cst.parse_module(source_code))


def parse_and_validate_llm_output(
    response_content: str,
    source_code: str,
    python_version: tuple[int, int, int],
    error_context: str,
    function_to_optimize: FunctionToOptimize | None = None,
    module_path: str | None = None,
) -> str:
    code: str | None = None
    try:
        if "```python" not in response_content:
            sentry_sdk.capture_message("LLM response did not contain a Python code block:\n" + response_content)
            raise LLMOutputParseError(
                "LLM response did not contain a Python code block.", raw_llm_output=response_content
            )

        code = extract_code_block(response_content)
        if code is None:
            raise LLMOutputParseError(
                "No Python code block found in the LLM response.", raw_llm_output=response_content
            )

        cleaned_code = validate_testgen_code(code, python_version[:2])

        if did_generate_ellipsis(source_code, cleaned_code):
            msg = f"Ellipsis in generated {error_context}test code, regenerating..."
            raise SyntaxError(msg)

        return cleaned_code  # noqa: TRY300
    except CodeValidationError as e:
        # Re-raise with raw LLM output added
        e.raw_llm_output = response_content  # type: ignore[attr-defined]
        sentry_sdk.capture_exception(e)
        raise
    except LLMOutputParseError:
        raise
    except Exception as e:
        sentry_sdk.capture_exception(e)
        raise


async def generate_and_validate_test_code(
    messages: list[ChatCompletionMessageParam],
    model: LLM,
    source_code: str,
    python_version: tuple[int, int, int],
    error_context: str,
    cost_tracker: CostTracker,
    user_id: str,
    posthog_event_suffix: str,
    trace_id: str = "",
    call_sequence: int | None = None,
    function_to_optimize: FunctionToOptimize | None = None,
    module_path: str | None = None,
    test_module_path: str | None = None,
    helper_function_names: list[str] | None = None,
    is_async: bool = False,
    test_index: int | None = None,
) -> tuple[str, str]:
    obs_context: dict[str, Any] | None = (
        {
            "call_sequence": call_sequence,
            "test_index": test_index,
            "module_path": module_path,
            "test_module_path": test_module_path,
            "helper_function_names": helper_function_names,
            "is_async": is_async,
            "function_to_optimize": {
                "function_name": function_to_optimize.function_name,
                "file_path": function_to_optimize.file_path,
                "qualified_name": function_to_optimize.qualified_name,
                "starting_line": function_to_optimize.starting_line,
                "ending_line": function_to_optimize.ending_line,
            }
            if function_to_optimize is not None
            else None,
        }
        if call_sequence is not None
        else None
    )
    response = await llm_client.call(
        llm=model,
        messages=messages,
        call_type="test_generation",
        trace_id=trace_id,
        user_id=user_id,
        python_version=".".join(str(v) for v in python_version),
        context=obs_context,
    )

    cost = response.cost
    cost_tracker.add(cost)

    debug_log_sensitive_data_from_callable(
        lambda: f"LLM {error_context}execute response:\n{response.raw_response.model_dump_json(indent=2)}"
    )

    if response.raw_response.usage:
        ph(
            user_id,
            f"aiservice-testgen-{posthog_event_suffix}execute-openai-usage",
            properties={"model": model.name, "usage": response.raw_response.usage.model_dump_json()},
        )

    raw_llm_content = response.content

    # Parse and validate
    validated_code = parse_and_validate_llm_output(
        response_content=raw_llm_content,
        source_code=source_code,
        python_version=python_version,
        error_context=error_context,
        function_to_optimize=function_to_optimize,
        module_path=module_path,
    )

    return validated_code, raw_llm_content


async def generate_regression_tests_from_function(
    source_code: str,
    qualified_name: str,
    user_id: str,
    function_name: str,
    python_version: tuple[int, int, int],
    data: TestGenSchema,
    unit_test_package: str = "pytest",
    execute_model: LLM = EXECUTE_MODEL,
    is_async: bool = False,  # noqa: FBT001, FBT002
    trace_id: str = "",
    call_sequence: int | None = None,
    test_index: int | None = None,
    model_type: str = "openai",
) -> tuple[str, str, str, str]:
    execute_messages, posthog_event_suffix, error_context = build_prompt(
        qualified_name=qualified_name,
        source_code=source_code,
        notes=generate_notes(source_code),
        function_name=function_name,
        unit_test_package=unit_test_package,
        is_async=is_async,
        is_numerical_code=data.is_numerical_code,
        model_type=model_type,
    )

    cost_tracker = CostTracker()
    try:
        validated_code, _raw_llm_content = await generate_and_validate_test_code(
            messages=execute_messages,
            model=execute_model,
            source_code=source_code,
            python_version=python_version,
            error_context=error_context,
            cost_tracker=cost_tracker,
            user_id=user_id,
            posthog_event_suffix=posthog_event_suffix,
            trace_id=trace_id,
            call_sequence=call_sequence,
            function_to_optimize=data.function_to_optimize,
            module_path=data.module_path,
            test_module_path=data.test_module_path,
            helper_function_names=data.helper_function_names,
            is_async=data.function_to_optimize.is_async or data.is_async or False,
            test_index=test_index,
        )

        processed_cst, display_cst = postprocessing_testgen_pipeline(
            parse_module_to_cst(validated_code),
            data.helper_function_names or [],
            data.function_to_optimize,
            data.module_path,
            data.source_code_being_tested,
        )

        generated_test_source = processed_cst.code
        raw_display_source = display_cst.code

        instrumented_behavior_tests, instrumented_perf_tests = instrument_tests(
            generated_test_source, data, python_version
        )

        if instrumented_behavior_tests is None or instrumented_perf_tests is None:
            msg = (
                f"There was an error detected in the function to optimize, is it valid Python code? trace_id={trace_id}"
            )
            logging.error(msg)
            raise TestGenerationFailedError(
                msg,
                debug_info={
                    "stage": "instrumentation",
                    "final_code": generated_test_source,
                    "validation_error": "Instrumentation failed - function may not be valid Python",
                },
            )

        sorted_generated_tests = safe_isort(generated_test_source)

        try:
            parse_module_to_cst(sorted_generated_tests)
            generated_test_source = sorted_generated_tests
        except Exception:  # noqa: BLE001
            sentry_sdk.capture_message("isort caused a syntax error in testgen; returning un-sorted code.")

        sorted_display_source = safe_isort(raw_display_source)
        try:
            parse_module_to_cst(sorted_display_source)
            raw_display_source = sorted_display_source
        except Exception:  # noqa: BLE001
            pass
        tree = ast.parse(generated_test_source, feature_version=python_version[:2])
        if not has_test_functions(tree):  # sanity check, after all the processing somehow no test functions
            msg = f"No test functions were found in the generated test code. trace_id={trace_id}"
            logging.error(msg)
            raise TestGenerationFailedError(
                msg,
                debug_info={
                    "stage": "postprocessing",
                    "final_code": generated_test_source,
                    "validation_error": "No test functions found after postprocessing",
                },
            )
        return generated_test_source, instrumented_behavior_tests, instrumented_perf_tests, raw_display_source  # noqa: TRY300
    except CodeValidationError as e:
        msg = f"Failed to generate valid {error_context}test code after {cost_tracker.calls} tries. trace_id={trace_id}"
        logging.exception(msg)
        debug_info = e.to_debug_dict()
        debug_info["raw_llm_output"] = getattr(e, "raw_llm_output", None)
        raise TestGenerationFailedError(msg, debug_info=debug_info) from e
    except LLMOutputParseError as e:
        msg = f"Failed to parse LLM output for {error_context}test code after {cost_tracker.calls} tries. trace_id={trace_id}"
        logging.exception(msg)
        raise TestGenerationFailedError(
            msg,
            debug_info={
                "stage": "llm_generation",
                "raw_llm_output": e.raw_llm_output,
                "initial_code": e.code,
                "validation_error": str(e),
            },
        ) from e
    except (SyntaxError, ValueError, cst.ParserSyntaxError) as e:
        msg = f"Failed to generate valid {error_context}test code after {cost_tracker.calls} tries. trace_id={trace_id}"
        logging.exception(msg)
        raise TestGenerationFailedError(msg, debug_info={"stage": "unknown", "validation_error": str(e)}) from e
    finally:
        if cost_tracker.total > 0:
            await update_optimization_cost(trace_id=trace_id, cost=cost_tracker.total, user_id=user_id)


async def testgen_python(
    request: AuthenticatedRequest, data: TestGenSchema
) -> tuple[int, TestGenResponseSchema | TestGenErrorResponseSchema]:
    """Generate Python tests using LLMs."""
    ph(request.user, "aiservice-testgen-called")

    try:
        python_version = validate_request_data(data)
    except HttpError as e:
        e.add_note(f"Testgen request validation error: {e.status_code} {e.message}")
        sentry_sdk.capture_exception(e)
        return e.status_code, TestGenErrorResponseSchema(error=e.message)

    if (demo_result := await try_demo_testgen(data, python_version)) is not None:
        return 200, demo_result

    logging.info("/testgen: Generating tests...")
    try:
        debug_log_sensitive_data(f"Generating tests for function {data.function_to_optimize.function_name}")

        # Using different LLMs for different test_index values to get more diverse tests
        test_index = data.test_index if data.test_index is not None else 0
        execute_model, model_source = select_model_for_test(test_index)
        logging.info(
            f"Using {model_source} model ({execute_model.name}) for test_index {test_index} to generate diverse tests"
        )

        (
            generated_test_source,
            instrumented_behavior_tests,
            instrumented_perf_tests,
            raw_display_source,
        ) = await generate_regression_tests_from_function(
            source_code=data.source_code_being_tested,
            qualified_name=data.function_to_optimize.qualified_name,
            user_id=request.user,
            function_name=data.function_to_optimize.qualified_name,
            python_version=python_version,
            data=data,
            unit_test_package=data.test_framework,
            is_async=data.function_to_optimize.is_async or data.is_async or False,
            trace_id=data.trace_id,
            call_sequence=data.call_sequence,
            execute_model=execute_model,
            test_index=test_index,
            model_type=execute_model.model_type,
        )

        ph(request.user, "aiservice-testgen-tests-generated")

        if hasattr(request, "should_log_features") and request.should_log_features:
            await safe_log_features(
                trace_id=data.trace_id,
                user_id=request.user,
                generated_tests=[raw_display_source],
                instrumented_generated_tests=[instrumented_behavior_tests],
                instrumented_perf_tests=[instrumented_perf_tests],
                test_index=test_index,
                test_framework=data.test_framework,
                metadata={
                    "test_timeout": data.test_timeout,
                    "function_to_optimize": data.function_to_optimize.qualified_name,
                    "file_path": data.function_to_optimize.file_path,
                },
            )

        return 200, TestGenResponseSchema(
            generated_tests=raw_display_source,
            instrumented_behavior_tests=instrumented_behavior_tests,
            instrumented_perf_tests=instrumented_perf_tests,
            raw_generated_tests=raw_display_source,
        )

    except TestGenerationFailedError as e:
        logging.exception("Test generation failed. trace_id=%s", data.trace_id)
        sentry_sdk.capture_exception(e)
        # Return detailed debug info for self-healing
        debug_info = None
        if e.debug_info:
            lines_removed = e.debug_info.get("lines_removed")
            debug_info = TestGenDebugInfo(
                stage=str(e.debug_info.get("stage", "unknown")),
                raw_llm_output=str(v) if (v := e.debug_info.get("raw_llm_output")) else None,
                initial_code=str(v) if (v := e.debug_info.get("initial_code")) else None,
                fixed_code=str(v) if (v := e.debug_info.get("fixed_code")) else None,
                final_code=str(v) if (v := e.debug_info.get("final_code")) else None,
                lines_removed=int(lines_removed) if lines_removed is not None else None,
                validation_error=str(v) if (v := e.debug_info.get("validation_error")) else None,
            )
        return 422, TestGenErrorResponseSchema(error=str(e), trace_id=data.trace_id, debug_info=debug_info)
    except Exception as e:
        logging.exception("Test generation failed. trace_id=%s", data.trace_id)
        sentry_sdk.capture_exception(e)
        return 500, TestGenErrorResponseSchema(
            error="Error generating tests. Internal server error.", trace_id=data.trace_id
        )
