from __future__ import annotations

import asyncio
import uuid
from collections.abc import Coroutine
from pathlib import Path
from typing import TYPE_CHECKING, Any

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.xml_utils import extract_xml_tag
from aiservice.common_utils import validate_trace_id
from aiservice.env_specific import debug_log_sensitive_data
from aiservice.llm import llm_client
from aiservice.llm_models import REFINEMENT_MODEL
from authapp.auth import AuthenticatedRequest
from core.languages.python.optimizer.context_utils.refiner_context import BaseRefinerContext, RefinementContextData
from core.log_features.log_event import update_optimization_cost
from core.log_features.log_features import log_features
from core.shared.optimizer_models import OptimizedCandidateSource
from core.shared.optimizer_schemas import OptimizeResponseItemSchema

if TYPE_CHECKING:
    from openai.types.chat import ChatCompletionMessageParam

    from aiservice.llm_models import LLM


refinement_api = NinjaAPI(urls_namespace="refinement")

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


async def refinement(  # noqa: D417
    user_id: str,
    optimization_id: str,
    ctx: BaseRefinerContext,
    trace_id: str = "",
    optimize_model: LLM = REFINEMENT_MODEL,
    call_sequence: int | None = None,
) -> RefinementIntermediateResponseItemschema | OptimizeErrorResponseSchema:
    """Optimize the given python code for performance using Anthropic's Claude 4 model.

    Parameters
    ----------
    :param user_id:
    :param optimization_id: the optimization id of the original candidate
    :param optimize_model: LLM for getting the refinements
    :param ctx: the refiner context (single or multi), has the data property which includes
        - speedup: original speedup for candidate
        - optimized_code_runtime: runtime in ns for candidate
        - optimization_id: id of optimization candidate
        - optimized_line_profiler_results: line profiler results for optimized candidate
        - optimized_source_code: source code of optimization candidate
        - read_only_dependency_code: unmodifiable code for better context
        - original_code_runtime: runtime in ns for baseline code
        - original_line_profiler_results: line profiler results for baseline code
        - original_source_code: baseline source code
        - optimized_explanation: Original explanation of the candidate

    Returns
    -------
    RefinementIntermediateResponseItemschema
    optimization_id: str = ""
    source_code
    explanation

    """
    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")

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

    try:
        output = await llm_client.call(
            llm=optimize_model,
            messages=messages,
            call_type="refinement",
            trace_id=trace_id or "",
            user_id=user_id,
            context=obs_context,
        )
        llm_cost = output.cost
    except Exception as e:
        debug_log_sensitive_data(f"Failed to generate code for source:\n{ctx.data.original_source_code}")
        return OptimizeErrorResponseSchema(error=str(e))
    debug_log_sensitive_data(f"ClaudeClient optimization response:\n{output.content}")
    if output.usage is not None:
        ph(
            user_id,
            "refinement-usage",
            properties={
                "model": optimize_model.name,
                "usage": {"input_tokens": output.usage.input_tokens, "output_tokens": output.usage.output_tokens},
            },
        )

    explanation = extract_xml_tag(output.content, "explanation")
    refined_explanation = explanation.strip() if explanation else output.content

    refined_optimization = ""
    try:
        refined_optimization = ctx.apply_patches_to_optimized_code(output.content)
    except (ValueError, ValidationError) as exc:
        sentry_sdk.capture_exception(exc)
        debug_log_sensitive_data(f"{type(exc).__name__} for source:\n{ctx.data.optimized_source_code}")
        debug_log_sensitive_data(f"Traceback: {exc}")
        refined_optimization = ""

    if not ctx.is_valid_refinement(refined_optimization):
        refined_optimization = ""

    return RefinementIntermediateResponseItemschema(
        parent_id=optimization_id,
        optimization_id=str(uuid.uuid4()),
        source_code=refined_optimization,
        explanation=refined_explanation,
        original_explanation=ctx.data.optimized_explanation,
        llm_cost=llm_cost,
    )


class RefinementRequestSchema(Schema):
    trace_id: str
    optimization_id: str
    original_source_code: str
    original_line_profiler_results: str = ""
    read_only_dependency_code: str
    optimized_source_code: str
    optimized_line_profiler_results: str = ""
    optimized_explanation: str
    original_code_runtime: str = ""
    optimized_code_runtime: str = ""
    speedup: str = ""
    python_version: str | None = None
    function_references: str | None = None
    call_sequence: int | None = None
    language: str = "python"  # python, javascript, typescript
    language_version: str | None = None
    rerun_trace_id: str | None = None


class OptimizeErrorResponseSchema(Schema):
    error: str


class RefinementIntermediateResponseItemschema(Schema):
    explanation: str
    parent_id: str
    optimization_id: str
    source_code: str
    original_explanation: str
    llm_cost: float


class Refinementschema(Schema):
    refinements: list[OptimizeResponseItemSchema]


@refinement_api.post("/", response={200: Refinementschema, 400: Refinementschema, 500: Refinementschema})
async def refine(
    request: AuthenticatedRequest, data: list[RefinementRequestSchema]
) -> tuple[int, Refinementschema | OptimizeErrorResponseSchema]:
    ph(request.user, "aiservice-refinement-called")
    trace_id = data[0].trace_id
    if not validate_trace_id(trace_id):
        return 400, OptimizeErrorResponseSchema(error="Invalid trace ID. Please provide a valid UUIDv4.")

    if data[0].rerun_trace_id:
        from core.shared.replay import get_rerun_record, rerun_refinement  # noqa: PLC0415

        record = await get_rerun_record(data[0].rerun_trace_id)
        if record is not None:
            parent_ids = [item.optimization_id for item in data]
            items = rerun_refinement(record, parent_ids)
            if items:
                return 200, Refinementschema(refinements=items)

    refinement_coros = []
    for i, item in enumerate(data):
        ctx_data = RefinementContextData(
            original_source_code=item.original_source_code,
            original_line_profiler_results=item.original_line_profiler_results,
            original_code_runtime=item.original_code_runtime,
            optimized_source_code=item.optimized_source_code,
            read_only_dependency_code=item.read_only_dependency_code,
            optimized_line_profiler_results=item.optimized_line_profiler_results,
            optimized_code_runtime=item.optimized_code_runtime,
            speedup=item.speedup,
            optimized_explanation=item.optimized_explanation,
            python_version=item.python_version,
            function_references=item.function_references,
            language=item.language,
            language_version=item.language_version,
        )
        ctx = BaseRefinerContext.get_dynamic_context(
            ctx_data=ctx_data, base_system_prompt=SYSTEM_PROMPT, base_user_prompt=USER_PROMPT
        )
        sequence = item.call_sequence if item.call_sequence is not None else i + 1
        refinement_coros.append(
            refinement(
                user_id=request.user,
                optimization_id=item.optimization_id,
                ctx=ctx,
                trace_id=trace_id,
                call_sequence=sequence,
            )
        )
    refinement_data = await asyncio.gather(*refinement_coros)
    # simple filtering mechanism, remove empty strings and remove duplicates after removing trailing and leading whitespaces, validate with libcst
    filtered_refined_optimizations = []
    source_code_set = set()
    total_llm_cost = 0.0
    for elem in refinement_data:
        if isinstance(elem, OptimizeErrorResponseSchema):
            continue
        total_llm_cost += elem.llm_cost
        try:
            # Use validate_code_syntax which handles both Python and JS/TS
            ctx.validate_code_syntax(elem.source_code)
        except cst.ParserSyntaxError as e:
            # log exception with sentry
            sentry_sdk.capture_exception(e)
            debug_log_sensitive_data(f"ParserSyntaxError for source:\n{elem.source_code}")
            debug_log_sensitive_data(f"Traceback: {e}")
            continue
        except (ValueError, ValidationError) as exc:
            # Another one bites the Pydantic validation dust
            sentry_sdk.capture_exception(exc)
            debug_log_sensitive_data(f"{type(exc).__name__} for source:\n{elem.source_code}")
            debug_log_sensitive_data(f"Traceback: {exc}")
            continue
        if (elem.source_code.strip() not in source_code_set) and elem.source_code != "":
            source_code_set.add(elem.source_code.strip())
            filtered_refined_optimizations.append(elem)
    # Parallel DB writes: log_features + update_optimization_cost run concurrently.
    # gather(return_exceptions=True) ensures both complete even if one fails.
    coros: list[Coroutine[Any, Any, Any]] = [
        update_optimization_cost(trace_id=trace_id, cost=total_llm_cost, user_id=request.user)
    ]
    if hasattr(request, "should_log_features") and request.should_log_features:
        coros.append(
            log_features(
                trace_id=trace_id,
                user_id=request.user,
                optimizations_raw={
                    cei.optimization_id: cei.source_code
                    for cei in refinement_data
                    if not isinstance(cei, OptimizeErrorResponseSchema)
                },
                optimizations_post={cei.optimization_id: cei.source_code for cei in filtered_refined_optimizations},
                explanations_raw={
                    cei.optimization_id: cei.explanation
                    for cei in refinement_data
                    if not isinstance(cei, OptimizeErrorResponseSchema)
                },
                explanations_post={cei.optimization_id: cei.explanation for cei in filtered_refined_optimizations},
                optimizations_origin={
                    cei.optimization_id: {"source": OptimizedCandidateSource.REFINE, "parent": cei.parent_id}
                    for cei in filtered_refined_optimizations
                },
            )
        )
    results = await asyncio.gather(*coros, return_exceptions=True)
    for result in results:
        if isinstance(result, BaseException):
            sentry_sdk.capture_exception(result)
    return 200, Refinementschema(
        refinements=[
            OptimizeResponseItemSchema(
                source_code=x.source_code,
                explanation=x.explanation,
                optimization_id=x.optimization_id,
                parent_id=x.parent_id,
            )
            for x in filtered_refined_optimizations
        ]
    )
