from __future__ import annotations

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

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

from aiservice.analytics.posthog import ph
from aiservice.common_utils import validate_trace_id
from aiservice.env_specific import debug_log_sensitive_data
from aiservice.llm import ADAPTIVE_OPTIMIZE_MODEL, calculate_llm_cost, call_llm
from authapp.auth import AuthenticatedRequest
from log_features.log_event import update_optimization_cost
from log_features.log_features import log_features
from core.languages.python.optimizer.context_utils.optimizer_context import OptimizeResponseItemSchema
from core.languages.python.optimizer.models import OptimizedCandidateSource

from .adaptive_optimizer_context import AdaptiveOptContext, AdaptiveOptContextData, AdaptiveOptRequestSchema

if TYPE_CHECKING:
    from openai.types.chat import ChatCompletionMessageParam

    from aiservice.llm 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()


async def perform_adaptive_optimize(
    user_id: str, ctx: AdaptiveOptContext, optimize_model: LLM = ADAPTIVE_OPTIMIZE_MODEL
) -> tuple[OptimizeResponseItemSchema | None, float | None, AdaptiveOptErrorResponseSchema | None]:
    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 call_llm(llm=optimize_model, messages=messages)
        llm_cost = calculate_llm_cost(output.raw_response, optimize_model)
    except Exception as e:
        logging.exception("Claude Code Generation error in adaptive_optimize")
        sentry_sdk.capture_exception(e)
        debug_log_sensitive_data(f"Failed to generate code for source:\n{ctx.data.original_source_code}")
        return None, None, AdaptiveOptErrorResponseSchema(error=str(e))
    debug_log_sensitive_data(f"ClaudeClient optimization response:\n{output.content}")
    if output.usage is not None:
        await asyncio.to_thread(
            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 = 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
            return (None, None, AdaptiveOptErrorResponseSchema(error="Invalid code generated " + str(extracted_code)))

        # 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, None  # noqa: TRY300
    except (ValueError, ValidationError, cst.ParserSyntaxError) as exc:
        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}")
        return None, None, AdaptiveOptErrorResponseSchema(error=str(exc))


@adaptive_optimize_api.post(
    "/",
    response={
        200: OptimizeResponseItemSchema,
        400: AdaptiveOptErrorResponseSchema,
        500: AdaptiveOptErrorResponseSchema,
    },
)
async def adaptive_optimize(
    request: AuthenticatedRequest, data: AdaptiveOptRequestSchema
) -> tuple[int, OptimizeResponseItemSchema | AdaptiveOptErrorResponseSchema]:
    await asyncio.to_thread(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.")

    adaptive_optimization_candidate, llm_cost, error = await perform_adaptive_optimize(user_id=request.user, ctx=ctx)
    total_llm_cost = 0.0
    if error:
        return 500, error
    if adaptive_optimization_candidate is None:
        return 500, AdaptiveOptErrorResponseSchema(error="Failed to generate optimization candidate")
    if llm_cost is not None:
        total_llm_cost += llm_cost

    if hasattr(request, "should_log_features") and request.should_log_features:
        await 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,
                }
            },
        )
    await update_optimization_cost(trace_id=trace_id, cost=total_llm_cost, user_id=request.user)
    return 200, adaptive_optimization_candidate
