"""Auto-configure context window based on known model limits."""
from __future__ import annotations

from dataclasses import replace

from logos_adk.config import ContextWindowConfig

_MODEL_LIMITS: list[tuple[str, int]] = [
    # Anthropic
    ("claude-opus", 200_000),
    ("claude-sonnet", 200_000),
    ("claude-haiku", 200_000),
    # OpenAI
    ("gpt-4o", 128_000),
    ("gpt-4-turbo", 128_000),
    ("gpt-4-1", 1_000_000),
    ("gpt-3.5-turbo", 16_385),
    ("o1", 200_000),
    ("o3", 200_000),
    ("o4-mini", 200_000),
    # Google
    ("gemini-2.0", 1_000_000),
    ("gemini-1.5-pro", 2_000_000),
    ("gemini-1.5-flash", 1_000_000),
    # DeepSeek
    ("deepseek-chat", 64_000),
    ("deepseek-reasoner", 64_000),
]


def get_model_context_limit(model: str) -> int:
    """Return known context limit for a model, or 0 if unknown."""
    if not model:
        return 0
    best_match = 0
    best_prefix_len = 0
    for prefix, limit in _MODEL_LIMITS:
        if model.startswith(prefix) and len(prefix) > best_prefix_len:
            best_match = limit
            best_prefix_len = len(prefix)
    return best_match


def auto_configure_context_window(
    config: ContextWindowConfig,
    *,
    model: str,
) -> ContextWindowConfig:
    """Auto-configure context window if not explicitly set.

    If max_context_tokens is 0 (disabled) and the model is recognized,
    returns a new ContextWindowConfig with the model's limit applied.
    Otherwise returns a copy of the original config unchanged.
    """
    if config.max_context_tokens > 0:
        return replace(config)

    limit = get_model_context_limit(model)
    if limit <= 0:
        return replace(config)

    return replace(config, max_context_tokens=limit)
