from __future__ import annotations

from dataclasses import dataclass

import openai
from openai import OpenAI
from tenacity import (
    retry,
    retry_if_exception_type,
    stop_after_attempt,
    wait_exponential,
)

from ..genome.hypothesis import HypothesisAgentGenome, ModelTier

# Modelli configurati per Phase 1 — tutti via OpenRouter
MODEL_TIER_S = "anthropic/claude-opus-4-7"
MODEL_TIER_A = "anthropic/claude-sonnet-4-6"
MODEL_TIER_B = "anthropic/claude-sonnet-4-6"
MODEL_TIER_C = "qwen/qwen-2.5-72b-instruct"
MODEL_TIER_D = "meta-llama/llama-3.3-70b-instruct"
OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1"

# Errori transient: retry. RateLimit/Auth/InvalidRequest: NO retry.
_RETRYABLE_EXCEPTIONS: tuple[type[BaseException], ...] = (
    openai.APIConnectionError,
    openai.APITimeoutError,
    openai.InternalServerError,
)


@dataclass(frozen=True)
class CompletionResult:
    text: str
    input_tokens: int
    output_tokens: int
    tier: ModelTier
    model: str


class LLMClient:
    def __init__(
        self,
        openrouter_api_key: str,
        model_tier_s: str = MODEL_TIER_S,
        model_tier_a: str = MODEL_TIER_A,
        model_tier_b: str = MODEL_TIER_B,
        model_tier_c: str = MODEL_TIER_C,
        model_tier_d: str = MODEL_TIER_D,
        openrouter_base_url: str = OPENROUTER_BASE_URL,
    ) -> None:
        self.model_tier_s = model_tier_s
        self.model_tier_a = model_tier_a
        self.model_tier_b = model_tier_b
        self.model_tier_c = model_tier_c
        self.model_tier_d = model_tier_d
        self.openrouter_base_url = openrouter_base_url
        self._tier_models: dict[ModelTier, str] = {
            ModelTier.S: model_tier_s,
            ModelTier.A: model_tier_a,
            ModelTier.B: model_tier_b,
            ModelTier.C: model_tier_c,
            ModelTier.D: model_tier_d,
        }
        self._client = OpenAI(api_key=openrouter_api_key, base_url=openrouter_base_url)

    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1.0, min=2.0, max=10.0),
        retry=retry_if_exception_type(_RETRYABLE_EXCEPTIONS),
        reraise=True,
    )
    def complete(
        self,
        genome: HypothesisAgentGenome,
        system: str,
        user: str,
        max_tokens: int = 2000,
    ) -> CompletionResult:
        model = self._tier_models[genome.model_tier]
        resp = self._client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": system},
                {"role": "user", "content": user},
            ],
            temperature=genome.temperature,
            top_p=genome.top_p,
            max_tokens=max_tokens,
        )
        usage = resp.usage
        assert usage is not None
        return CompletionResult(
            text=resp.choices[0].message.content or "",
            input_tokens=usage.prompt_tokens,
            output_tokens=usage.completion_tokens,
            tier=genome.model_tier,
            model=model,
        )