from __future__ import annotations

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

import libcst as cst
import sentry_sdk
from ninja.errors import HttpError
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 parse_python_version, validate_trace_id
from aiservice.env_specific import debug_log_sensitive_data, debug_log_sensitive_data_from_callable
from aiservice.llm import llm_client
from aiservice.llm_models import LLM, OPTIMIZE_MODEL
from authapp.user import get_user_by_id
from core.languages.python.optimizer.context_utils.optimizer_context import BaseOptimizerContext
from core.languages.python.optimizer.diff_patches_utils.diff import DiffMethod
from core.log_features.log_event import get_or_create_optimization_event
from core.log_features.log_features import safe_log_features
from core.shared.optimizer_config import MAX_OPTIMIZER_CALLS, get_model_distribution
from core.shared.optimizer_models import OptimizedCandidateSource
from core.shared.optimizer_schemas import (
    OptimizeErrorResponseSchema,
    OptimizeResponseItemSchema,
    OptimizeResponseSchema,
)

if TYPE_CHECKING:
    from openai.types.chat import ChatCompletionMessageParam

    from authapp.auth import AuthenticatedRequest
    from core.shared.optimizer_models import OptimizeSchema

# 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_SYSTEM_PROMPT = (current_dir / "async_system_prompt.md").read_text()
ASYNC_USER_PROMPT = (current_dir / "async_user_prompt.md").read_text()
JIT_INSTRUCTIONS = (current_dir / "jit_instructions.md").read_text()


async def generate_optimization_candidate(
    user_id: str,
    ctx: BaseOptimizerContext,
    trace_id: str,
    dependency_code: str | None = None,
    optimize_model: LLM = OPTIMIZE_MODEL,
    python_version: tuple[int, int, int] = (3, 12, 9),
    call_sequence: int | None = None,
) -> tuple[OptimizeResponseItemSchema | None, float | None, str]:
    """Optimize the given python code for performance using LLMs."""
    logging.info("/optimize: Optimizing python code.")
    debug_log_sensitive_data(f"Optimizing python code for user {user_id}:\n{ctx.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
    python_version_str = ".".join(str(x) for x in python_version)

    system_prompt = ctx.get_system_prompt(python_version_str)
    user_prompt = ctx.get_user_prompt(dependency_code or "", None)

    obs_context: dict[str, Any] | None = {"call_sequence": call_sequence} if call_sequence is not None else None

    messages: list[ChatCompletionMessageParam] = [
        ChatCompletionSystemMessageParam(role="system", content=system_prompt),
        ChatCompletionUserMessageParam(role="user", content=user_prompt),
    ]
    try:
        output = await llm_client.call(
            llm=optimize_model,
            messages=messages,
            call_type="optimization",
            trace_id=trace_id,
            user_id=user_id,
            python_version=python_version_str,
            context=obs_context,
        )
    except Exception:
        debug_log_sensitive_data(f"Failed to generate code for source:\n{ctx.source_code}")
        return None, None, optimize_model.name

    llm_cost = output.cost

    debug_log_sensitive_data_from_callable(
        lambda: f"LLM optimization response:\n{output.raw_response.model_dump_json(indent=2)}"
    )

    if output.raw_response.usage is not None:
        ph(
            user_id,
            "aiservice-optimize-openai-usage",
            properties={"model": optimize_model.name, "usage": output.raw_response.usage.json()},
        )

    ctx.extract_code_and_explanation_from_llm_res(output.content)
    try:
        res = await asyncio.to_thread(ctx.parse_and_generate_candidate_schema)
        if res is not None and ctx.is_valid_code():
            return res, llm_cost, optimize_model.name
    except (ValueError, ValidationError, cst.ParserSyntaxError) as e:
        sentry_sdk.capture_message(f"Error parsing optimization result: {e}")
        debug_log_sensitive_data(f"error for source:\n{ctx.source_code}")
        debug_log_sensitive_data(f"Traceback: {e}")

    return None, llm_cost, optimize_model.name


async def optimize_python_code(
    user_id: str,
    ctx: BaseOptimizerContext,
    trace_id: str,
    original_source_code: str,
    dependency_code: str | None = None,
    python_version: tuple[int, int, int] = (3, 12, 9),
    n_candidates: int = 0,
) -> tuple[list[OptimizeResponseItemSchema], float, dict[str, dict[str, str]], dict[str, str]]:
    """Run parallel optimizations with multiple models based on the distribution config.

    Returns:
        tuple containing:
        - list of optimization results
        - total LLM cost
        - dict of raw code/explanations keyed by optimization_id
        - dict mapping optimization_id to model name

    """
    # Create tasks for each model call
    tasks: list[
        tuple[asyncio.Task[tuple[OptimizeResponseItemSchema | None, float | None, str]], BaseOptimizerContext]
    ] = []
    call_sequence = 1

    if n_candidates == 0:
        return [], 0.0, {}, {}

    async with asyncio.TaskGroup() as tg:
        for model, num_calls in get_model_distribution(n_candidates, MAX_OPTIMIZER_CALLS):
            for _ in range(num_calls):
                # Each call needs its own context instance to avoid shared state issues
                # Use original_source_code (markdown format) to preserve file path info
                task_ctx = BaseOptimizerContext.get_dynamic_context(
                    ctx.base_system_prompt, ctx.base_user_prompt, original_source_code
                )
                task = tg.create_task(
                    generate_optimization_candidate(
                        user_id=user_id,
                        ctx=task_ctx,
                        trace_id=trace_id,
                        dependency_code=dependency_code,
                        optimize_model=model,
                        python_version=python_version,
                        call_sequence=call_sequence,
                    )
                )
                tasks.append((task, task_ctx))
                call_sequence += 1

    # Collect results
    optimization_results: list[OptimizeResponseItemSchema] = []
    total_cost = 0.0
    code_and_explanations: dict[str, dict[str, str]] = {}
    optimization_models: dict[str, str] = {}

    for task, task_ctx in tasks:
        result, cost, model_name = task.result()
        if cost:
            total_cost += cost
        if result is not None:
            optimization_results.append(result)
            optimization_models[result.optimization_id] = model_name
        # Collect raw code/explanations for logging
        for op_id, cei in task_ctx.code_and_explanation_before_post_processing.items():
            code_and_explanations[op_id] = {"code": cei.code, "explanation": cei.explanation}

    return optimization_results, total_cost, code_and_explanations, optimization_models


def validate_request_data(data: OptimizeSchema, ctx: BaseOptimizerContext) -> tuple[int, int, int]:
    if not data.source_code:
        raise HttpError(400, "Source code cannot be empty.")
    if not validate_trace_id(data.trace_id):
        raise HttpError(400, "Invalid trace ID. Please provide a valid UUIDv4.")

    if not data.python_version:
        raise HttpError(400, "Python version is required.")
    try:
        python_version = parse_python_version(data.python_version)
    except ValueError as e:
        raise HttpError(
            400, "Invalid Python version, it should look like 3.x.x. We only support Python 3.9 and above."
        ) from e

    try:
        ctx.validate_and_parse_source_code(data.source_code, feature_version=python_version[:2])
    except SyntaxError as e:
        raise HttpError(
            400, "Invalid source code. It is not valid Python code. Please check syntax of your code."
        ) from e

    return python_version


async def optimize_python(
    request: AuthenticatedRequest, data: OptimizeSchema
) -> tuple[int, OptimizeResponseSchema | OptimizeErrorResponseSchema]:
    """Optimize Python code for performance using LLMs."""
    system_prompt = ASYNC_SYSTEM_PROMPT if data.is_async else SYSTEM_PROMPT
    user_prompt = ASYNC_USER_PROMPT if data.is_async else USER_PROMPT
    if data.is_numerical_code:
        system_prompt += f"\n{JIT_INSTRUCTIONS}\n"
    ctx: BaseOptimizerContext = BaseOptimizerContext.get_dynamic_context(
        system_prompt, user_prompt, data.source_code, DiffMethod.NO_DIFF
    )
    ph(request.user, "aiservice-optimize-called")

    try:
        python_version = validate_request_data(data, ctx)
    except HttpError as e:
        e.add_note(f"Optimizer request validation error: {e.status_code} {e.message}")
        logging.exception("Optimizer request validation error: %s. trace_id=%s", e.message, data.trace_id)
        sentry_sdk.capture_exception(e)
        return e.status_code, OptimizeErrorResponseSchema(error=e.message)

    from core.languages.python.optimizer.demo_hacks import try_demo_optimize  # noqa: PLC0415

    demo_response = await try_demo_optimize(ctx)
    if demo_response is not None:
        async with asyncio.TaskGroup() as tg:
            event_task = tg.create_task(
                get_or_create_optimization_event(
                    event_type="no-pr",
                    user_id=request.user,
                    current_username=data.current_username,
                    repo_owner=data.repo_owner,
                    repo_name=data.repo_name,
                    trace_id=data.trace_id,
                    api_key_id=request.api_key_id,
                    metadata={
                        "codeflash_version": data.codeflash_version,
                        "num_optimizations": 4,
                        "experiment_metadata": data.experiment_metadata,
                    },
                    llm_cost=0,
                )
            )
        event, _created = event_task.result()
        for item in demo_response.optimizations:
            item.optimization_event_id = str(event.id) if event else None
        return 200, demo_response

    try:
        async with asyncio.TaskGroup() as tg:
            optimize_task = tg.create_task(
                optimize_python_code(
                    user_id=request.user,
                    ctx=ctx,
                    trace_id=data.trace_id,
                    original_source_code=data.source_code,
                    dependency_code=data.dependency_code,
                    python_version=python_version,
                    n_candidates=data.n_candidates,
                )
            )
            user_task = None
            if data.current_username is None:
                user_task = tg.create_task(get_user_by_id(request.user))
    except Exception as e:
        logging.exception("Error during optimization task or user retrieval. trace_id=%s", data.trace_id)
        sentry_sdk.capture_exception(e)
        return 500, OptimizeErrorResponseSchema(error="Error generating optimizations. Internal server error.")

    optimization_response_items, llm_cost, code_and_explanations, optimization_models = optimize_task.result()
    if user_task:
        user = await user_task
        if user and user.github_username:
            data.current_username = str(user.github_username)

    if len(optimization_response_items) == 0:
        ph(request.user, "aiservice-optimize-no-optimizations-found")
        debug_log_sensitive_data(f"No optimizations found for source:\n{data.source_code}")
        logging.error(
            "Could not generate any optimizations. trace_id=%s, repo=%s/%s, n_candidates=%d, source_len=%d",
            data.trace_id,
            data.repo_owner,
            data.repo_name,
            data.n_candidates,
            len(data.source_code) if data.source_code else 0,
        )
        return 422, OptimizeErrorResponseSchema(error="Could not generate any optimizations. Please try again.")
    ph(
        request.user,
        "aiservice-optimize-optimizations-found",
        properties={"num_optimizations": len(optimization_response_items)},
    )

    fire_and_forget(
        safe_log_features(
            trace_id=data.trace_id,
            user_id=request.user,
            original_code=data.source_code,
            dependency_code=data.dependency_code,
            optimizations_raw={op_id: cei["code"] for op_id, cei in code_and_explanations.items()},
            optimizations_post={cei.optimization_id: cei.source_code for cei in optimization_response_items},
            explanations_raw={op_id: cei["explanation"] for op_id, cei in code_and_explanations.items()},
            explanations_post={cei.optimization_id: cei.explanation for cei in optimization_response_items},
            experiment_metadata=data.experiment_metadata or None,
            optimizations_origin={
                cei.optimization_id: {
                    "source": OptimizedCandidateSource.OPTIMIZE,
                    "parent": None,
                    "model": optimization_models.get(cei.optimization_id, "unknown"),
                }
                for cei in optimization_response_items
            },
        )
    )

    event, _created = await get_or_create_optimization_event(
        event_type="no-pr",
        user_id=request.user,
        current_username=data.current_username,
        repo_owner=data.repo_owner,
        repo_name=data.repo_name,
        trace_id=data.trace_id,
        api_key_id=request.api_key_id,
        metadata={
            "codeflash_version": data.codeflash_version,
            "num_optimizations": len(optimization_response_items),
            "experiment_metadata": data.experiment_metadata,
        },
        llm_cost=llm_cost,
    )

    for item in optimization_response_items:
        item.optimization_event_id = str(event.id) if event else None

    response = OptimizeResponseSchema(optimizations=optimization_response_items)

    def log_response() -> str:
        parts = [f"Response:\n{response.model_dump_json()}"]
        for opt in response.optimizations:
            parts.append(f"Optimized source:\n{opt.source_code}")
            parts.append(f"Optimization explanation:\n{opt.explanation}")
        return "\n".join(parts)

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