"""Okna kontekstu modeli i planowanie budżetu tokenów (LOG-30/31).

Po co to istnieje: horoskop MA powstać niezależnie od objętości promptu. Żeby to
zagwarantować, trzeba wiedzieć dwie rzeczy o każdym modelu — ile zmieści na
wejściu (okno kontekstu) i ile maksymalnie wypisze na wyjściu. Bez tego łatwo
wysłać prompt, który wypełnia całe okno i **nie zostawia miejsca na odpowiedź** —
model kończy wtedy na `max_tokens` z pustą albo uciętą treścią.

Zasada naczelna: **zawsze rezerwuj miejsce na odpowiedź.** Budżet promptu liczy
się jako `okno_kontekstu − zarezerwowane_wyjście − margines`, nigdy odwrotnie.

Wartości są zaszyte jako rozsądne domyślne i nadpisywalne środowiskiem
(`<DOSTAWCA>_CONTEXT_WINDOW`, `<DOSTAWCA>_MAX_OUTPUT`) — modele wychodzą szybciej,
niż aktualizuje się ten plik.
"""
from __future__ import annotations

import os

# model -> (okno kontekstu, maksymalne wyjście) w tokenach
_MODEL_LIMITS: dict[str, tuple[int, int]] = {
    # Anthropic
    "claude-opus-4-8": (1_000_000, 128_000),
    "claude-opus-4-7": (1_000_000, 128_000),
    "claude-opus-4-6": (1_000_000, 128_000),
    "claude-sonnet-5": (1_000_000, 128_000),
    "claude-sonnet-4-6": (1_000_000, 128_000),
    "claude-fable-5": (1_000_000, 128_000),
    "claude-haiku-4-5": (200_000, 64_000),
    # OpenAI
    "gpt-4o": (128_000, 16_384),
    "gpt-4o-mini": (128_000, 16_384),
    "gpt-4.1": (1_000_000, 32_768),
    "gpt-4.1-mini": (1_000_000, 32_768),
    # lokalne (Ollama/vLLM) — zwykle małe okno, dlatego ostrożna domyślna
    "llama3.1": (8_192, 4_096),
    "llama3.2": (8_192, 4_096),
    "qwen2.5": (32_768, 8_192),
    "mistral": (32_768, 8_192),
}

# gdy modelu nie ma w tabeli — zachowawczo, żeby nie obiecywać nieistniejącego okna
_FALLBACK: dict[str, tuple[int, int]] = {
    "anthropic": (200_000, 32_000),
    "openai": (128_000, 16_384),
    "local": (8_192, 4_096),
}

# ile tokenów zostawiamy jako bufor na narzut protokołu i niedokładność liczenia
SAFETY_MARGIN = 2_000
# poniżej tylu tokenów wyjścia nie ma sensu wołać modelu — nie zmieści horoskopu
MIN_OUTPUT = 1_500
# powyżej tylu tokenów promptu ostrzegamy użytkownika (nadal pozwalając wysłać)
WARN_PROMPT_TOKENS = 90_000


def _env_int(provider: str, suffix: str) -> int | None:
    raw = os.getenv(f"{provider.upper()}_{suffix}")
    if not raw:
        return None
    try:
        value = int(raw)
    except ValueError:
        return None
    return value if value > 0 else None


def limits_for(provider: str, model: str) -> tuple[int, int]:
    """(okno kontekstu, maksymalne wyjście) dla modelu — z nadpisaniem z ENV.

    Dopasowanie po prefiksie, bo nazwy modeli lokalnych niosą tag (`llama3.1:8b`).
    """
    env_ctx = _env_int(provider, "CONTEXT_WINDOW")
    env_out = _env_int(provider, "MAX_OUTPUT")

    key = (model or "").strip().lower()
    known: tuple[int, int] | None = _MODEL_LIMITS.get(key)
    if known is None:
        for name, pair in _MODEL_LIMITS.items():
            if key.startswith(name):
                known = pair
                break
    if known is None:
        known = _FALLBACK.get(provider, _FALLBACK["local"])

    return (env_ctx or known[0], env_out or known[1])


def plan(provider: str, model: str, prompt_tokens: int,
         want_output: int | None = None) -> dict:
    """Ile tokenów wyjścia zamówić dla promptu tej wielkości.

    Zwraca plan z jawną diagnostyką — UI ma z czego zbudować ostrzeżenie, a błąd
    ma czym wytłumaczyć, dlaczego się nie udało.
    """
    context_window, model_max_output = limits_for(provider, model)
    room = context_window - prompt_tokens - SAFETY_MARGIN
    target = want_output or model_max_output
    max_output = max(0, min(model_max_output, target, room))

    warnings: list[str] = []
    if prompt_tokens > WARN_PROMPT_TOKENS:
        warnings.append(
            f"Prompt ma ~{prompt_tokens} tokenów — to dużo. Zapytanie zostanie wysłane, "
            f"ale potrwa dłużej i będzie odpowiednio kosztowne."
        )
    if max_output < MIN_OUTPUT:
        warnings.append(
            f"Po zmieszczeniu promptu zostaje tylko {max_output} tokenów na odpowiedź "
            f"(minimum {MIN_OUTPUT}). Zmniejsz budżet promptu albo wybierz model "
            f"z większym oknem kontekstu."
        )

    return {
        "provider": provider,
        "model": model,
        "context_window": context_window,
        "model_max_output": model_max_output,
        "prompt_tokens": prompt_tokens,
        "max_output": max_output,
        "fits": max_output >= MIN_OUTPUT,
        "warnings": warnings,
    }


def prompt_token_budget(provider: str, model: str, reserve_output: int | None = None) -> int:
    """Ile tokenów promptu wolno wysłać, ZAWSZE zostawiając miejsce na odpowiedź.

    To jest podstawa opcji „maksymalny kontekst modelu" w UI.
    """
    context_window, model_max_output = limits_for(provider, model)
    reserve = reserve_output or model_max_output
    return max(0, context_window - reserve - SAFETY_MARGIN)