from __future__ import annotations

import asyncio
import logging
from pathlib import Path
from typing import TYPE_CHECKING

import libcst as cst
import sentry_sdk
import stamina
from ninja import NinjaAPI, Schema
from openai.types.chat import ChatCompletionSystemMessageParam, ChatCompletionUserMessageParam
from pydantic import ValidationError

from aiservice.analytics.posthog import ph
from aiservice.background import fire_and_forget
from aiservice.common_utils import validate_trace_id
from aiservice.env_specific import debug_log_sensitive_data
from aiservice.llm import LLMOutputUnparseable, llm_client
from aiservice.llm_models import ADAPTIVE_OPTIMIZE_MODEL
from authapp.auth import AuthenticatedRequest
from core.log_features.log_event import update_optimization_cost
from core.log_features.log_features import safe_log_features
from core.shared.optimizer_models import OptimizedCandidateSource
from core.shared.optimizer_schemas import OptimizeResponseItemSchema

from .adaptive_optimizer_context import AdaptiveOptContext, AdaptiveOptContextData, AdaptiveOptRequestSchema

logger = logging.getLogger(__name__)

if TYPE_CHECKING:
    from openai.types.chat import ChatCompletionMessageParam

    from aiservice.llm_models import LLM


class AdaptiveOptErrorResponseSchema(Schema):
    error: str


class AdaptiveOptResponseItemschema(Schema):
    optimization_id: str
    parent_id: str
    source_code: str
    explanation: str


adaptive_optimize_api = NinjaAPI(urls_namespace="adaptive_optimize")

# Get the directory of the current file
current_dir = Path(__file__).parent
SYSTEM_PROMPT = (current_dir / "ADAPTIVE_OPTIMIZER_SYSTEM_PROMPT.md").read_text()

USER_PROMPT = (current_dir / "ADAPTIVE_OPTIMIZER_USER_PROMPT.md").read_text()


@stamina.retry(on=LLMOutputUnparseable, attempts=2)
async def perform_adaptive_optimize(
    user_id: str, ctx: AdaptiveOptContext, trace_id: str = "", optimize_model: LLM = ADAPTIVE_OPTIMIZE_MODEL
) -> tuple[OptimizeResponseItemSchema, float]:
    system_prompt = ctx.get_system_prompt()
    user_prompt = ctx.get_user_prompt()

    system_message = ChatCompletionSystemMessageParam(role="system", content=system_prompt)
    user_message = ChatCompletionUserMessageParam(role="user", content=user_prompt)
    messages: list[ChatCompletionMessageParam] = [system_message, user_message]
    debug_log_sensitive_data(f"This was the user prompt\n {user_prompt}\n")
    try:
        output = await llm_client.call(
            llm=optimize_model, messages=messages, call_type="adaptive_optimize", trace_id=trace_id, user_id=user_id
        )
        llm_cost = output.cost
    except Exception as e:
        logger.exception("adaptive_optimize LLM call failed: trace_id=%s, user_id=%s", trace_id, user_id)
        debug_log_sensitive_data(f"Failed to generate code for source:\n{ctx.data.original_source_code}")
        raise LLMOutputUnparseable(str(e)) from e
    debug_log_sensitive_data(f"ClaudeClient optimization response:\n{output.content}")
    if output.usage is not None:
        ph(
            user_id,
            "adaptive_optimize-usage",
            properties={
                "model": optimize_model.name,
                "usage": {"input_tokens": output.usage.input_tokens, "output_tokens": output.usage.output_tokens},
            },
        )

    # Regex doesn't work yet in extracting everything else other than the search replace block
    llm_res = output.content

    try:
        ctx.extract_code_and_explanation_from_llm_res(llm_res)
        new_opt = await asyncio.to_thread(ctx.parse_and_generate_candidate_schema)
        if not new_opt or not ctx.is_valid_code():
            extracted_code = ctx.extracted_code_and_expl.code if ctx.extracted_code_and_expl else None
            logger.error("adaptive_optimize invalid code: trace_id=%s, user_id=%s", trace_id, user_id)
            raise LLMOutputUnparseable("Invalid code generated " + str(extracted_code), cost=llm_cost)

        # the parent is the last candidate in the previous optimizations
        last_optimization_id = ctx.data.attempts[-1].optimization_id
        new_opt.parent_id = last_optimization_id
        return new_opt, llm_cost  # noqa: TRY300
    except (ValueError, ValidationError, cst.ParserSyntaxError) as exc:
        logger.exception("adaptive_optimize parsing failed: trace_id=%s, user_id=%s", trace_id, user_id)
        sentry_sdk.capture_exception(exc)
        debug_log_sensitive_data(f"{type(exc).__name__} for source:\n{ctx.data.original_source_code}")
        debug_log_sensitive_data(f"Traceback: {exc}")
        raise LLMOutputUnparseable(str(exc), cost=llm_cost) from exc


@adaptive_optimize_api.post(
    "/",
    response={
        200: OptimizeResponseItemSchema,
        400: AdaptiveOptErrorResponseSchema,
        422: AdaptiveOptErrorResponseSchema,
        500: AdaptiveOptErrorResponseSchema,
    },
)
async def adaptive_optimize(
    request: AuthenticatedRequest, data: AdaptiveOptRequestSchema
) -> tuple[int, OptimizeResponseItemSchema | AdaptiveOptErrorResponseSchema]:
    ph(request.user, "aiservice-adaptive_optimize-called")
    ctx_data = AdaptiveOptContextData(
        original_source_code=data.original_source_code, attempts=data.candidates, python_version_str="3.12"
    )
    ctx = AdaptiveOptContext(ctx_data=ctx_data, base_system_prompt=SYSTEM_PROMPT, base_user_prompt=USER_PROMPT)
    trace_id = data.trace_id
    if not validate_trace_id(trace_id):
        return 400, AdaptiveOptErrorResponseSchema(error="Invalid trace ID. Please provide a valid UUIDv4.")
    try:
        adaptive_optimization_candidate, llm_cost = await perform_adaptive_optimize(
            user_id=request.user, ctx=ctx, trace_id=trace_id
        )
        if adaptive_optimization_candidate is None:
            logger.error("adaptive_optimize endpoint returning 500: trace_id=%s, no candidate generated", trace_id)
            return 500, AdaptiveOptErrorResponseSchema(error="Failed to generate optimization candidate")

    except LLMOutputUnparseable as e:
        return 422, AdaptiveOptErrorResponseSchema(error=str(e))

    fire_and_forget(update_optimization_cost(trace_id=trace_id, cost=llm_cost, user_id=request.user))
    if hasattr(request, "should_log_features") and request.should_log_features:
        fire_and_forget(
            safe_log_features(
                trace_id=data.trace_id,
                user_id=request.user,
                optimizations_raw={
                    adaptive_optimization_candidate.optimization_id: adaptive_optimization_candidate.source_code
                },
                explanations_raw={
                    adaptive_optimization_candidate.optimization_id: adaptive_optimization_candidate.explanation
                },
                optimizations_origin={
                    adaptive_optimization_candidate.optimization_id: {
                        "source": OptimizedCandidateSource.ADAPTIVE,
                        "parent": adaptive_optimization_candidate.parent_id,
                    }
                },
            )
        )
    return 200, adaptive_optimization_candidate
