from __future__ import annotations

import asyncio
import logging
from typing import TYPE_CHECKING, Any

import libcst as cst

from codeflash_api.llm._client import LLMOutputUnparseableError
from codeflash_api.llm._models import (
    ANTHROPIC_CLAUDE_SONNET_4_5,
    LLM,
    OPENAI_GPT_5_MINI,
)
from codeflash_api.optimize._context import BaseOptimizerContext

if TYPE_CHECKING:
    from codeflash_api.llm._client import LLMClient
    from codeflash_api.optimize.schemas import OptimizeResponseItem

log = logging.getLogger(__name__)

OPENAI_MODEL: LLM = OPENAI_GPT_5_MINI
ANTHROPIC_MODEL: LLM = ANTHROPIC_CLAUDE_SONNET_4_5

MAX_OPTIMIZER_CALLS = 6


def get_model_distribution(
    total_calls: int, max_calls: int
) -> list[tuple[LLM, int]]:
    """
    Split calls between OpenAI and Anthropic models.
    """
    final_total = min(total_calls, max_calls)
    claude_calls = (final_total - 1) // 2
    gpt_calls = final_total - claude_calls
    return [
        (OPENAI_MODEL, gpt_calls),
        (ANTHROPIC_MODEL, claude_calls),
    ]


async def generate_optimization_candidate(
    llm_client: LLMClient,
    user_id: str,
    ctx: BaseOptimizerContext,
    trace_id: str,
    *,
    dependency_code: str | None = None,
    optimize_model: LLM = OPENAI_MODEL,
    python_version: tuple[int, int, int] = (3, 12, 9),
    call_sequence: int | None = None,
    baseline_runtime_ns: int | None = None,
    loop_count: int | None = None,
    line_profiler_results: str | None = None,
    test_input_examples: str | None = None,
) -> tuple[OptimizeResponseItem | None, float | None, str]:
    """
    Generate a single optimization candidate via LLM.
    """
    log.info("Generating optimization candidate")
    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 "",
        line_profiler_results,
        baseline_runtime_ns=baseline_runtime_ns,
        loop_count=loop_count,
        test_input_examples=test_input_examples,
    )

    messages: list[dict[str, Any]] = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt},
    ]
    try:
        output = await llm_client.call(
            llm=optimize_model,
            messages=messages,
            call_type="optimization",
            trace_id=trace_id,
        )
    except LLMOutputUnparseableError as e:
        return None, e.cost, optimize_model.name
    except Exception:
        log.exception("Failed to generate optimization")
        return None, None, optimize_model.name

    llm_cost = output.cost

    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, cst.ParserSyntaxError):
        log.warning("Error parsing optimization result", exc_info=True)

    return None, llm_cost, optimize_model.name


async def optimize_python_code(
    llm_client: LLMClient,
    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,
    baseline_runtime_ns: int | None = None,
    loop_count: int | None = None,
    line_profiler_results: str | None = None,
    test_input_examples: str | None = None,
) -> tuple[
    list[OptimizeResponseItem],
    float,
    dict[str, dict[str, str]],
    dict[str, str],
]:
    """
    Run parallel optimizations with multiple models.
    """
    tasks: list[
        tuple[
            asyncio.Task[
                tuple[
                    OptimizeResponseItem | 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):
                task_ctx = BaseOptimizerContext.get_dynamic_context(
                    ctx.base_system_prompt,
                    ctx.base_user_prompt,
                    original_source_code,
                )
                lp_for_this_call = (
                    line_profiler_results if call_sequence % 2 == 1 else None
                )
                task = tg.create_task(
                    generate_optimization_candidate(
                        llm_client,
                        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,
                        baseline_runtime_ns=baseline_runtime_ns,
                        loop_count=loop_count,
                        line_profiler_results=lp_for_this_call,
                        test_input_examples=test_input_examples,
                    )
                )
                tasks.append((task, task_ctx))
                call_sequence += 1

    optimization_results: list[OptimizeResponseItem] = []
    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
        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,
    )
