from __future__ import annotations

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

from ninja import NinjaAPI
from openai.types.chat import ChatCompletionSystemMessageParam, ChatCompletionUserMessageParam

from aiservice.analytics.posthog import ph
from aiservice.background import fire_and_forget
from aiservice.common.markdown_utils import split_markdown_code
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 OPTIMIZE_MODEL
from aiservice.validators.javascript_validator import validate_javascript_syntax, validate_typescript_syntax
from core.languages.java.optimizer_lp import optimize_java_code_line_profiler
from core.languages.js_ts.context_helpers import is_multi_context_js, is_multi_context_ts
from core.languages.js_ts.optimizer_lp import optimize_javascript_code_line_profiler
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 update_optimization_cost
from core.log_features.log_features import safe_log_features
from core.shared.optimizer_config import MAX_OPTIMIZER_LP_CALLS, get_model_distribution
from core.shared.optimizer_models import OptimizedCandidateSource, OptimizeSchemaLP
from core.shared.optimizer_schemas import (
    OptimizeErrorResponseSchema,
    OptimizeResponseItemSchema,
    OptimizeResponseSchema,
)

if TYPE_CHECKING:
    from openai.types.chat import ChatCompletionMessageParam

    from aiservice.llm_models import LLM
    from authapp.auth import AuthenticatedRequest


optimize_line_profiler_api = NinjaAPI(urls_namespace="optimize-line-profiler")


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


async def optimize_python_code_line_profiler_single(
    user_id: str,
    trace_id: str,
    line_profiler_results: str,
    ctx: BaseOptimizerContext,
    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 line profile.")
    debug_log_sensitive_data(f"Optimizing python code for user {user_id}:\n{ctx.source_code}")

    python_version_str = ".".join(str(x) for x in python_version)

    # 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_prompt = ctx.get_system_prompt(python_version_str=python_version_str)
    user_prompt = ctx.get_user_prompt(dependency_code or "", line_profiler_results)
    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")
    # TODO: Verify if the context window length is within the model capability

    obs_context: dict[str, Any] = {}
    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="line_profiler",
            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"OpenAIClient optimization response:\n{output.raw_response.model_dump_json(indent=2)}"
    )

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

    ctx.extract_code_and_explanation_from_llm_res(output.content)
    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

    return None, llm_cost, optimize_model.name


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

    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

    """
    if n_candidates == 0:
        return [], 0.0, {}, {}
    model_distribution = get_model_distribution(n_candidates, MAX_OPTIMIZER_LP_CALLS)

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

    async with asyncio.TaskGroup() as tg:
        for model, num_calls in model_distribution:
            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(
                    optimize_python_code_line_profiler_single(
                        user_id=user_id,
                        trace_id=trace_id,
                        line_profiler_results=line_profiler_results,
                        ctx=task_ctx,
                        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


@optimize_line_profiler_api.post(
    "/",
    response={
        200: OptimizeResponseSchema,
        400: OptimizeErrorResponseSchema,
        422: OptimizeErrorResponseSchema,
        500: OptimizeErrorResponseSchema,
    },
)
async def optimize(
    request: AuthenticatedRequest, data: OptimizeSchemaLP
) -> tuple[int, OptimizeResponseSchema | OptimizeErrorResponseSchema]:
    if data.rerun_trace_id:
        from core.shared.replay import get_rerun_record, rerun_optimize  # noqa: PLC0415

        record = await get_rerun_record(data.rerun_trace_id)
        if record is not None:
            result = rerun_optimize(record, "OPTIMIZE_LP")
            if result is not None:
                return 200, result

    ph(request.user, "aiservice-optimize-called")
    ctx: BaseOptimizerContext = BaseOptimizerContext.get_dynamic_context(
        SYSTEM_PROMPT, USER_PROMPT, data.source_code, DiffMethod.NO_DIFF
    )
    # Route based on language
    language = data.language.lower() if data.language else "python"
    is_javascript = language in ("javascript", "typescript")
    is_python = language == "python"

    if is_python:
        if not data.python_version:
            return 400, OptimizeErrorResponseSchema(error="Python version is required.")
        try:
            python_version: tuple[int, int, int] = parse_python_version(data.python_version)
        except:  # noqa: E722
            return 400, OptimizeErrorResponseSchema(
                error="Invalid Python version, it should look like 3.x.x. We only support Python 3.9 and above."
            )
        try:
            ctx.validate_and_parse_source_code(code=data.source_code, feature_version=python_version[:2])
        except SyntaxError:
            return 400, OptimizeErrorResponseSchema(
                error="Invalid source code. It is not valid Python code. Please check syntax of your code."
            )
    if not validate_trace_id(data.trace_id):
        return 400, OptimizeErrorResponseSchema(error="Invalid trace ID. Please provide a valid UUIDv4.")
    if data.line_profiler_results is None:
        return 400, OptimizeErrorResponseSchema(error="Line profiler results are required.")

    if is_javascript:
        # JavaScript/TypeScript path
        lang_name = "TypeScript" if language == "typescript" else "JavaScript"

        # Check if multi-file context
        is_multi_file = (
            is_multi_context_ts(data.source_code) if language == "typescript" else is_multi_context_js(data.source_code)
        )

        if is_multi_file:
            # Validate each file in the multi-file context
            file_to_code = split_markdown_code(data.source_code, language)
            if not file_to_code:
                return 400, OptimizeErrorResponseSchema(
                    error=f"Invalid source code format. Expected multi-file {lang_name} markdown format."
                )

            for file_path, code in file_to_code.items():
                if language == "typescript":
                    is_valid, error = validate_typescript_syntax(code)
                else:
                    is_valid, error = validate_javascript_syntax(code)

                if not is_valid:
                    return 400, OptimizeErrorResponseSchema(
                        error=f"Invalid source code in {file_path}. It is not valid {lang_name}: {error}"
                    )
        else:
            # Single file validation
            if language == "typescript":
                is_valid, error = validate_typescript_syntax(data.source_code)
            else:
                is_valid, error = validate_javascript_syntax(data.source_code)

            if not is_valid:
                return 400, OptimizeErrorResponseSchema(
                    error=f"Invalid source code. It is not valid {lang_name} code. Error: {error}"
                )

        language_version = data.language_version or "ES2022"

        (optimization_response_items, llm_cost, optimization_models) = await optimize_javascript_code_line_profiler(
            user_id=request.user,
            trace_id=data.trace_id,
            source_code=data.source_code,
            line_profiler_results=data.line_profiler_results,
            dependency_code=data.dependency_code,
            language_version=language_version,
            language=language,
            n_candidates=data.n_candidates,
        )
        # JavaScript path doesn't have code_and_explanations dict like Python
        code_and_explanations: dict[str, dict[str, str]] = {}

    elif language == "java":
        # Java path
        from aiservice.validators.java_validator import validate_java_syntax  # noqa: PLC0415
        from core.languages.java.demo_hacks import try_demo_optimize_java_lp  # noqa: PLC0415
        from core.languages.java.optimizer import is_multi_context_java  # noqa: PLC0415

        demo_response = await try_demo_optimize_java_lp(data.source_code)
        if demo_response is not None:
            return 200, demo_response

        is_multi_file = is_multi_context_java(data.source_code)

        if is_multi_file:
            file_to_code = split_markdown_code(data.source_code, "java")
            if not file_to_code:
                return 400, OptimizeErrorResponseSchema(
                    error="Invalid source code format. Expected multi-file Java markdown format."
                )
            for file_path, code in file_to_code.items():
                is_valid, error = validate_java_syntax(code)
                if not is_valid:
                    return 400, OptimizeErrorResponseSchema(
                        error=f"Invalid source code in {file_path}. It is not valid Java: {error}"
                    )
        else:
            is_valid, error = validate_java_syntax(data.source_code)
            if not is_valid:
                return 400, OptimizeErrorResponseSchema(
                    error=f"Invalid source code. It is not valid Java code. Error: {error}"
                )

        language_version = data.language_version or "17"

        (optimization_response_items, llm_cost, optimization_models) = await optimize_java_code_line_profiler(
            user_id=request.user,
            trace_id=data.trace_id,
            source_code=data.source_code,
            line_profiler_results=data.line_profiler_results,
            dependency_code=data.dependency_code,
            language_version=language_version,
            n_candidates=data.n_candidates,
        )
        code_and_explanations = {}

    else:
        # Python path (default)
        system_prompt = SYSTEM_PROMPT
        if data.is_numerical_code:
            system_prompt += f"\n{JIT_INSTRUCTIONS}\n"
        ctx = BaseOptimizerContext.get_dynamic_context(system_prompt, USER_PROMPT, data.source_code, DiffMethod.NO_DIFF)
        try:
            python_version = parse_python_version(data.python_version or "3.12.0")
        except:  # noqa: E722
            return 400, OptimizeErrorResponseSchema(
                error="Invalid Python version, it should look like 3.x.x. We only support Python 3.9 and above."
            )
        try:
            ctx.validate_and_parse_source_code(code=data.source_code, feature_version=python_version[:2])
        except SyntaxError:
            return 400, OptimizeErrorResponseSchema(
                error="Invalid source code. It is not valid Python code. Please check syntax of your code."
            )

        (
            optimization_response_items,
            llm_cost,
            code_and_explanations,
            optimization_models,
        ) = await optimize_python_code_line_profiler(
            user_id=request.user,
            trace_id=data.trace_id,
            line_profiler_results=data.line_profiler_results,
            ctx=ctx,
            original_source_code=data.source_code,
            dependency_code=data.dependency_code,
            n_candidates=data.n_candidates,
            python_version=python_version,
        )

    if len(optimization_response_items) == 0:
        await update_optimization_cost(trace_id=data.trace_id, cost=llm_cost, user_id=request.user)
        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 (line profiler). trace_id=%s, n_candidates=%d, source_len=%d, has_line_profiler=%s",
            data.trace_id,
            data.n_candidates,
            len(data.source_code) if data.source_code else 0,
            bool(data.line_profiler_results),
        )
        return 422, OptimizeErrorResponseSchema(error="Error generating optimizations. Internal server error.")
    ph(
        request.user,
        "aiservice-optimize-optimizations-found",
        properties={"num_optimizations": len(optimization_response_items), "language": language},
    )

    fire_and_forget(update_optimization_cost(trace_id=data.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,
                original_code=data.source_code,
                dependency_code=data.dependency_code,
                line_profiler_results=data.line_profiler_results,
                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_LP,
                        "parent": None,
                        "model": optimization_models.get(cei.optimization_id, "unknown"),
                    }
                    for cei in optimization_response_items
                },
            )
        )

    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
