"""Python test generation pipeline — prompt, LLM, validate, postprocess."""

from __future__ import annotations

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

import jinja2

from codeflash_api.languages.python._cst_utils import (
    any_ellipsis_in_cst,
    ellipsis_in_cst_not_types,
    parse_module_to_cst,
)
from codeflash_api.languages.python._markdown import (
    extract_code_block,
    is_multi_context,
    split_markdown_code,
)
from codeflash_api.llm._models import (
    ANTHROPIC_CLAUDE_HAIKU_4_5,
    EXECUTE_MODEL,
    LLM,
)
from codeflash_api.testgen._validate import (
    CodeValidationError,
    has_test_functions,
    validate_testgen_code,
)

if TYPE_CHECKING:
    from codeflash_api.llm._client import LLMClient

log = logging.getLogger(__name__)

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


def select_model_for_test(
    test_index: int,
) -> tuple[LLM, str]:
    """
    Select model based on test_index: even → OpenAI, odd → Haiku.
    """
    if test_index % 2 == 0:
        return EXECUTE_MODEL, "openai"
    return ANTHROPIC_CLAUDE_HAIKU_4_5, "anthropic"


def _generate_notes(source_code: str) -> str:
    """
    Generate preprocessing notes for the LLM prompt.
    """
    notes: list[str] = []
    if re.search(r"\btorch\b", source_code):
        notes.append(
            "Do not generate torch tensors with size greater than 100MB."
        )
    if not notes:
        return ""
    return "Notes:\n" + "\n".join(notes)


def _did_generate_ellipsis(
    source_code: str,
    generated_code: str,
) -> bool:
    """
    Check if the LLM produced ellipsis that was not in the source.
    """
    generated_module = parse_module_to_cst(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(parse_module_to_cst(sc))
            for sc in split_markdown_code(source_code).values()
        )
    return not any_ellipsis_in_cst(parse_module_to_cst(source_code))


def build_prompt(
    *,
    function_name: str,
    source_code: str,
    unit_test_package: str,
    is_async: bool,
    model_type: str = "openai",
) -> list[dict[str, Any]]:
    """
    Build the messages list for the testgen LLM call.
    """
    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=function_name,
        model_type=model_type,
        is_numerical_code=False,
    )
    user_prompt = _JINJA_ENV.get_template("compose_user.md.j2").render(
        notes=_generate_notes(source_code),
        is_numerical_code=False,
        is_async=is_async,
        unit_test_package=unit_test_package,
        function_name=function_name,
        function_code=source_code,
        package_comment="",
    )

    return [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt},
    ]


def _parse_python_version(
    version_str: str,
) -> tuple[int, int, int]:
    """
    Parse a version string like '3.12.1' into a tuple.
    """
    parts = version_str.split(".")
    if len(parts) < 2:
        msg = f"Invalid Python version: {version_str}"
        raise ValueError(msg)
    major = int(parts[0])
    minor = int(parts[1])
    patch = int(parts[2]) if len(parts) > 2 else 0
    if major != 3 or minor < 9:
        msg = f"Unsupported Python version: {version_str}"
        raise ValueError(msg)
    return (major, minor, patch)


def _parse_and_validate_llm_output(
    response_content: str,
    source_code: str,
    python_version: tuple[int, int, int],
) -> str:
    """
    Extract code from LLM response and validate it.
    """
    if "```python" not in response_content:
        msg = "LLM response did not contain a Python code block."
        raise CodeValidationError(
            msg,
            initial_code=response_content,
            validation_error=msg,
        )

    code = extract_code_block(response_content)
    if code is None:
        msg = "No Python code block found in LLM response."
        raise CodeValidationError(
            msg,
            initial_code=response_content,
            validation_error=msg,
        )

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

    if _did_generate_ellipsis(source_code, cleaned_code):
        msg = "Ellipsis in generated test code, regenerating..."
        raise SyntaxError(msg)

    return cleaned_code


async def generate_tests(
    llm_client: LLMClient,
    *,
    source_code: str,
    function_name: str,
    test_framework: str,
    python_version_str: str,
    is_async: bool,
    trace_id: str,
    user_id: str,
    test_index: int,
) -> str:
    """
    Generate regression tests for a Python function.
    """
    python_version = _parse_python_version(python_version_str)
    model, model_type = select_model_for_test(test_index)

    messages = build_prompt(
        function_name=function_name,
        source_code=source_code,
        unit_test_package=test_framework,
        is_async=is_async,
        model_type=model_type,
    )

    response = await llm_client.call(
        llm=model,
        messages=messages,
        call_type="test_generation",
        trace_id=trace_id,
        user_id=user_id,
    )

    validated_code = _parse_and_validate_llm_output(
        response_content=response.content,
        source_code=source_code,
        python_version=python_version,
    )

    tree = ast.parse(
        validated_code,
        feature_version=python_version[:2],
    )
    if not has_test_functions(tree):
        msg = (
            "No test functions in the generated test code."
            f" trace_id={trace_id}"
        )
        raise CodeValidationError(
            msg,
            final_code=validated_code,
            validation_error=("No test functions found after postprocessing"),
        )

    return validated_code
