from __future__ import annotations

import asyncio
from collections.abc import Coroutine
from pathlib import Path
from typing import Any

import sentry_sdk
from ninja import NinjaAPI
from openai.types.chat import (
    ChatCompletionMessageParam,
    ChatCompletionSystemMessageParam,
    ChatCompletionUserMessageParam,
)
from packaging import version

from aiservice.analytics.posthog import ph
from authapp.auth import AuthenticatedRequest
from aiservice.common.markdown_utils import wrap_code_in_markdown
from aiservice.common.xml_utils import extract_xml_tag
from aiservice.common_utils import validate_trace_id
from aiservice.env_specific import IS_PRODUCTION, debug_log_sensitive_data
import stamina

from aiservice.llm import LLMOutputUnparseable, llm_client
from aiservice.llm_models import EXPLANATIONS_MODEL, LLM
from core.languages.python.explanations.models import (
    ExplanationsErrorResponseSchema,
    ExplanationsResponseSchema,
    ExplanationsSchema,
)
from core.log_features.log_event import update_optimization_cost
from core.log_features.log_features import log_features
from core.shared.jinja_utils import create_prompt_env

_PARSED_0182 = version.parse("0.18.2")

explanations_api = NinjaAPI(urls_namespace="explanations")

_PROMPT_DIR = Path(__file__).parent / "prompts"
_jinja_env = create_prompt_env(_PROMPT_DIR, trim_blocks=True, lstrip_blocks=True)

SYSTEM_PROMPT_TEMPLATE = _jinja_env.get_template("system_prompt.md.j2")
USER_PROMPT_TEMPLATE = _jinja_env.get_template("user_prompt.md.j2")


@stamina.retry(on=LLMOutputUnparseable, attempts=2)
async def explain_optimizations(
    user_id: str, data: ExplanationsSchema, explanations_model: LLM = EXPLANATIONS_MODEL
) -> tuple[ExplanationsResponseSchema, float]:
    # Avoid building potentially very large debug strings when logging is disabled.
    if not IS_PRODUCTION:
        debug_log_sensitive_data(f"Generating an explanation for {user_id}:\n{data.optimized_code}")

    # Parse the incoming version once per call (faster than parsing the constant each time).
    parsed_codeflash_version = version.parse(data.codeflash_version)
    if parsed_codeflash_version <= _PARSED_0182 and data.annotated_tests:
        data.annotated_tests = wrap_code_in_markdown(data.annotated_tests)

    include_throughput = data.original_throughput is not None and data.optimized_throughput is not None
    include_concurrency = data.original_concurrency_ratio is not None and data.optimized_concurrency_ratio is not None

    template_vars = {
        "model_type": explanations_model.model_type,
        "original_source_code": data.source_code,
        "original_line_profiler_results": data.original_line_profiler_results or "[No profiler results available]",
        "optimized_source_code": data.optimized_code,
        "optimized_line_profiler_results": data.optimized_line_profiler_results or "[No profiler results available]",
        "original_code_runtime": data.original_code_runtime,
        "optimized_code_runtime": data.optimized_code_runtime,
        "speedup": data.speedup,
        "annotated_tests": data.annotated_tests,
        "read_only_dependency_code": data.dependency_code or "[No read only code present]",
        "original_explanation": data.original_explanation,
        "python_version": data.python_version or "Not Available",
        "function_references": data.function_references or "Not Available",
        "include_throughput": include_throughput,
        "include_concurrency": include_concurrency,
        "acceptance_reason": data.acceptance_reason,
    }
    if include_throughput:
        template_vars["original_throughput"] = data.original_throughput
        template_vars["optimized_throughput"] = data.optimized_throughput
        template_vars["throughput_improvement"] = (
            data.throughput_improvement or "[Unable to calculate throughput improvement]"
        )
    if include_concurrency:
        template_vars["original_concurrency_ratio"] = data.original_concurrency_ratio
        template_vars["optimized_concurrency_ratio"] = data.optimized_concurrency_ratio
        template_vars["concurrency_improvement"] = (
            data.concurrency_improvement or "[Unable to calculate concurrency improvement]"
        )

    system_prompt = SYSTEM_PROMPT_TEMPLATE.render(template_vars)
    user_prompt = USER_PROMPT_TEMPLATE.render(template_vars)

    system_message = ChatCompletionSystemMessageParam(role="system", content=system_prompt)
    user_message = ChatCompletionUserMessageParam(role="user", content=user_prompt)
    if not IS_PRODUCTION:
        debug_log_sensitive_data(f"{system_prompt}{user_prompt}")
    messages: list[ChatCompletionMessageParam] = [system_message, user_message]

    obs_context: dict[str, str | float | int] = {"optimization_id": data.optimization_id, "speedup": data.speedup}
    if data.call_sequence is not None:
        obs_context["call_sequence"] = data.call_sequence

    try:
        output = await llm_client.call(
            llm=explanations_model,
            messages=messages,
            call_type="explanation",
            trace_id=data.trace_id,
            user_id=user_id,
            context=obs_context,
        )
    except Exception as e:
        raise LLMOutputUnparseable(str(e)) from e
    if not IS_PRODUCTION:
        debug_log_sensitive_data(f"AIClient optimization response:\n{output.content}")
    if output.usage is not None:
        ph(
            user_id,
            "aiservice-optimize-openai-usage",
            properties={
                "model": explanations_model.name,
                "n": 1,
                "usage": {"input_tokens": output.usage.input_tokens, "output_tokens": output.usage.output_tokens},
            },
        )
    return ExplanationsResponseSchema(explanation=output.content), output.cost


@explanations_api.post(
    "/",
    response={
        200: ExplanationsResponseSchema,
        400: ExplanationsErrorResponseSchema,
        422: ExplanationsErrorResponseSchema,
        500: ExplanationsErrorResponseSchema,
    },
)
async def explain(
    request: AuthenticatedRequest, data: ExplanationsSchema
) -> tuple[int, ExplanationsResponseSchema | ExplanationsErrorResponseSchema]:
    ph(request.user, "aiservice-explain-called")
    if not validate_trace_id(data.trace_id):
        return 400, ExplanationsErrorResponseSchema(error="Invalid trace ID. Please provide a valid UUIDv4.")
    try:
        explanation_response, llm_cost = await explain_optimizations(request.user, data)
    except LLMOutputUnparseable:
        ph(request.user, "Explanation not generated, revert to old explanation")
        debug_log_sensitive_data("No explanation was generated")
        return 422, ExplanationsErrorResponseSchema(error="Error generating optimizations. Internal server error.")
    ph(request.user, "explanation generated", properties={"explanation": explanation_response})
    # parse xml tag for explanation
    explanation = extract_xml_tag(explanation_response.explanation, "explain")
    if not explanation:
        return 422, ExplanationsErrorResponseSchema(error="Failed to parse explanation from LLM response.")
    coros: list[Coroutine[Any, Any, Any]] = [
        update_optimization_cost(trace_id=data.trace_id, cost=llm_cost, user_id=request.user)
    ]
    if hasattr(request, "should_log_features") and request.should_log_features:
        coros.append(log_features(trace_id=data.trace_id, user_id=request.user, final_explanation=explanation))
    results = await asyncio.gather(*coros, return_exceptions=True)
    for coro_result in results:
        if isinstance(coro_result, BaseException):
            sentry_sdk.capture_exception(coro_result)
    response = ExplanationsResponseSchema(explanation=explanation)
    return 200, response
