from __future__ import annotations

import ast
import re
import uuid
from pathlib import Path
from typing import TYPE_CHECKING

import libcst as cst
import sentry_sdk
from aiservice.analytics.posthog import ph
from aiservice.common_utils import parse_python_version
from aiservice.env_specific import (
    create_openai_client,
    debug_log_sensitive_data,
    debug_log_sensitive_data_from_callable,
)
from aiservice.models.aimodels import OPTIMIZE_MODEL
from authapp.auth import AuthBearer
from log_features.log_features import log_features
from ninja import NinjaAPI, Schema
from openai import OpenAIError
from openai.types.chat import ChatCompletionSystemMessageParam, ChatCompletionUserMessageParam
from pydantic import ValidationError
from testgen.instrumentation.edit_generated_test import parse_module_to_cst

from optimizer.models import CodeAndExplanation, CodeExplanationAndID
from optimizer.postprocess import optimizations_postprocessing_pipeline

if TYPE_CHECKING:
    from aiservice.models.aimodels import LLM
    from openai.types.chat import (
        ChatCompletionAssistantMessageParam,
        ChatCompletionFunctionMessageParam,
        ChatCompletionToolMessageParam,
    )

optimize_api = NinjaAPI(auth=AuthBearer(), urls_namespace="optimize")

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


async def optimize_python_code(
    user_id: str, source_code: str, n: int = 1, optimize_model: LLM = OPTIMIZE_MODEL
) -> list[CodeAndExplanation]:
    """Optimize the given python code for performance using OpenAI's GPT-4 model.

    Parameters
    ----------
    - source_code (str): The python code to optimize.
    - n (int): Number of optimization variants to generate. Default is 1.

    Returns: - List[Tuple[Union[str, None], Union[str, None]]]: A list of tuples where the first element is the
    optimized code and the second is the explanation.

    """
    print("/optimize: Optimizing python code.")
    debug_log_sensitive_data(f"Optimizing python code for user {user_id}:\n{source_code}")
    # TODO: Experiment with iterative approaches to optimization. Take the learnings from the testing phase into the
    #  next optimization iteration
    # TODO: Experiment with iterative chain-of-thought generation. ask what is the
    #  function doing and then ask it to describe how to speed it up and then generate optimization
    system_message = ChatCompletionSystemMessageParam(role="system", content=SYSTEM_PROMPT)
    user_message = ChatCompletionUserMessageParam(role="user", content=USER_PROMPT.format(source_code=source_code))
    messages: list[
        ChatCompletionSystemMessageParam
        | ChatCompletionUserMessageParam
        | ChatCompletionAssistantMessageParam
        | ChatCompletionToolMessageParam
        | ChatCompletionFunctionMessageParam
    ] = [system_message, user_message]
    async with create_openai_client() as openai_client:
        # TODO: Verify if the context window length is within the model capability
        try:
            output = await openai_client.with_options(max_retries=3).chat.completions.create(
                model=optimize_model.name, messages=messages, n=n
            )
        except OpenAIError as e:
            print("OpenAI Code Generation error ...")
            print(e)
            debug_log_sensitive_data(f"Failed to generate code for source:\n{source_code}")
            return []

    debug_log_sensitive_data(f"OpenAIClient optimization response:\n{output.model_dump_json(indent=2)}")

    if output.usage is not None:
        ph(
            user_id,
            "aiservice-optimize-openai-usage",
            properties={"model": optimize_model.name, "n": n, "usage": output.usage.json()},
        )
    results = [content for op in output.choices if (content := op.message.content)]
    optimized_code_and_explanations: list[CodeAndExplanation] = []
    for result in results:
        match = re.match(r"(.*)```python(?:\n|\\n)(.*?)```(.*)", result, re.DOTALL | re.MULTILINE)
        if match:
            code = match.group(2)
            explanation = match.group(1) + match.group(3)
            try:
                cst_module = parse_module_to_cst(code)
            except cst.ParserSyntaxError as e:
                # log exception with sentry
                sentry_sdk.capture_exception(e)
                debug_log_sensitive_data(f"ParserSyntaxError for source:\n{code}")
                debug_log_sensitive_data(f"Traceback: {e}")
                continue
            try:
                optimized_code_and_explanations.append(CodeAndExplanation(cst_module, explanation))
            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{code}")
                debug_log_sensitive_data(f"Traceback: {exc}")

    return optimized_code_and_explanations


class OptimizeSchema(Schema):
    source_code: str
    trace_id: str
    python_version: str
    experiment_metadata: dict[str, str] | None = None
    codeflash_version: str | None = None


class OptimizeResponseItemSchema(Schema):
    source_code: str
    explanation: str
    optimization_id: str


class OptimizeResponseSchema(Schema):
    optimizations: list[OptimizeResponseItemSchema]


class OptimizeErrorResponseSchema(Schema):
    error: str


@optimize_api.post(
    "/", response={200: OptimizeResponseSchema, 400: OptimizeErrorResponseSchema, 500: OptimizeErrorResponseSchema}
)
async def optimize(request, data: OptimizeSchema) -> tuple[int, OptimizeResponseSchema | OptimizeErrorResponseSchema]:
    ph(request.user, "aiservice-optimize-called")
    try:
        python_version: tuple[int, int, int] = parse_python_version(data.python_version)
    except:
        return 400, OptimizeErrorResponseSchema(
            error="Invalid Python version, it should look like 3.x.x. We only support Python 3.9 and above."
        )
    try:
        parsed = ast.parse(data.source_code, feature_version=python_version[:2])
        compile(data.source_code, "data.source_code", "exec")
        if not parsed.body:
            raise SyntaxError
    except SyntaxError:
        return 400, OptimizeErrorResponseSchema(
            error="Invalid source code. It is not valid Python code. Please check syntax of your code."
        )
    try:
        uuid_obj = uuid.UUID(data.trace_id, version=4)
        if not str(uuid_obj) == data.trace_id:
            raise ValueError
    except ValueError:
        return 400, OptimizeErrorResponseSchema(error="Invalid trace ID. Please provide a valid UUIDv4.")
    optimized_code_and_explanations = await optimize_python_code(request.user, data.source_code, n=10)
    if len(optimized_code_and_explanations) == 0:
        ph(request.user, "aiservice-optimize-no-optimizations-found")
        debug_log_sensitive_data(f"No optimizations found for source:\n{data.source_code}")
        return 500, OptimizeErrorResponseSchema(error="Error generating optimizations. Internal server error.")
    ph(
        request.user,
        "aiservice-optimize-optimizations-found",
        properties={"num_optimizations": len(optimized_code_and_explanations)},
    )
    traced_optimizations = [
        CodeExplanationAndID(cst_module=ce.cst_module, explanation=ce.explanation, id=str(uuid.uuid4()))
        for ce in optimized_code_and_explanations
    ]
    processed_optimizations: list[CodeExplanationAndID] = optimizations_postprocessing_pipeline(
        data.source_code, traced_optimizations
    )

    if request.tier is None:
        await log_features(
            trace_id=data.trace_id,
            user_id=request.user,
            original_code=data.source_code,
            optimizations_raw={cei.id: cei.cst_module.code for cei in traced_optimizations},
            optimizations_post={cei.id: cei.cst_module.code for cei in processed_optimizations},
            explanations_raw={cei.id: cei.explanation for cei in traced_optimizations},
            explanations_post={cei.id: cei.explanation for cei in processed_optimizations},
            experiment_metadata=data.experiment_metadata if data.experiment_metadata else None,
        )
    response = OptimizeResponseSchema(
        optimizations=[
            OptimizeResponseItemSchema(
                source_code=ce.cst_module.code, explanation=ce.explanation, optimization_id=ce.id
            )
            for ce in processed_optimizations
        ]
    )

    def log_response():
        debug_log_sensitive_data(f"Response:\n{response.json()}")
        for opt in response.optimizations:
            debug_log_sensitive_data(f"Optimized source:\n{opt.source_code}")
            debug_log_sensitive_data(f"Optimization explanation:\n{opt.explanation}")

    debug_log_sensitive_data_from_callable(log_response)
    ph(request.user, "aiservice-optimize-successful")
    return 200, response
