"""Context overflow detection for LLM responses."""

from __future__ import annotations

import re
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from pi_agents.ai.types import AssistantMessage

OVERFLOW_PATTERNS: list[re.Pattern[str]] = [
    re.compile(r"prompt is too long", re.IGNORECASE),
    re.compile(r"input is too long for requested model", re.IGNORECASE),
    re.compile(r"exceeds the context window", re.IGNORECASE),
    re.compile(r"input token count.*exceeds the maximum", re.IGNORECASE),
    re.compile(r"maximum prompt length is \d+", re.IGNORECASE),
    re.compile(r"reduce the length of the messages", re.IGNORECASE),
    re.compile(r"maximum context length is \d+ tokens", re.IGNORECASE),
    re.compile(r"exceeds the limit of \d+", re.IGNORECASE),
    re.compile(r"exceeds the available context size", re.IGNORECASE),
    re.compile(r"greater than the context length", re.IGNORECASE),
    re.compile(r"context window exceeds limit", re.IGNORECASE),
    re.compile(r"exceeded model token limit", re.IGNORECASE),
    re.compile(r"too large for model with \d+ maximum context length", re.IGNORECASE),
    re.compile(r"model_context_window_exceeded", re.IGNORECASE),
    re.compile(r"context[_ ]length[_ ]exceeded", re.IGNORECASE),
    re.compile(r"too many tokens", re.IGNORECASE),
    re.compile(r"token limit exceeded", re.IGNORECASE),
]

# Cerebras-specific pattern for 400/413 status codes
_CEREBRAS_PATTERN = re.compile(r"^4(00|13)\s*(status code)?\s*\(no body\)", re.IGNORECASE)


def is_context_overflow(message: AssistantMessage, context_window: int | None = None) -> bool:
    """Check if an assistant message indicates a context overflow error.

    Handles two cases:
    1. Error-based overflow: Most providers return stopReason "error" with a specific error message.
    2. Silent overflow: Some providers accept overflow requests and return successfully.
       For these, we check if usage.input exceeds the context window.
    """
    # Case 1: Check error message patterns
    if message.stop_reason == "error" and message.error_message:
        if any(p.search(message.error_message) for p in OVERFLOW_PATTERNS):
            return True
        if _CEREBRAS_PATTERN.search(message.error_message):
            return True

    # Case 2: Silent overflow (z.ai style)
    if context_window and message.stop_reason == "stop":
        input_tokens = message.usage.input + message.usage.cache_read
        if input_tokens > context_window:
            return True

    return False


def get_overflow_patterns() -> list[re.Pattern[str]]:
    """Get the overflow patterns for testing purposes."""
    return list(OVERFLOW_PATTERNS)
