"""Context window resolution, estimation, and overflow helpers."""

from __future__ import annotations

from dataclasses import dataclass
import json
import math
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from memnixa.config import AppConfig, ProviderConfig
    from memnixa.providers.base import ProviderSpec


_SAFE_FALLBACK_CONTEXT_WINDOW_TOKENS = 16384
_SAFE_FALLBACK_MAX_OUTPUT_TOKENS = 4096
_OVERFLOW_KEYWORDS = (
    "maximum context length",
    "context length exceeded",
    "context_length_exceeded",
    "too many tokens",
    "prompt is too long",
    "prompt too long",
    "input is too long",
    "context window",
    "exceeds the context",
    "token limit exceeded",
    "上下文长度",
    "超出上下文",
    "超出最大长度",
)


@dataclass(frozen=True, slots=True)
class ContextWindowResolution:
    """Resolved limits for one provider/model profile."""

    context_window_tokens: int
    max_output_tokens: int
    compact_threshold_tokens: int
    warn_threshold_tokens: int
    input_budget_tokens: int
    source: str


def resolve_context_window(
    config: "AppConfig",
    profile: "ProviderConfig",
    spec: "ProviderSpec",
) -> ContextWindowResolution:
    """Resolve usable context limits for one provider profile."""
    context_window_tokens = _first_positive_int(
        profile.context_window_tokens,
        spec.default_context_window_tokens,
        config.context_window_tokens,
        fallback=_SAFE_FALLBACK_CONTEXT_WINDOW_TOKENS,
    )
    max_output_tokens = _first_positive_int(
        profile.max_output_tokens,
        spec.default_max_output_tokens,
        config.max_output_tokens,
        fallback=_SAFE_FALLBACK_MAX_OUTPUT_TOKENS,
    )
    warn_ratio = _normalize_ratio(config.context_warn_threshold_ratio, fallback=0.8)
    compact_ratio = _normalize_ratio(
        config.context_compact_threshold_ratio, fallback=0.88
    )
    if compact_ratio < warn_ratio:
        compact_ratio = warn_ratio
    safety_margin_tokens = max(0, int(config.context_safety_margin_tokens))
    input_budget_tokens = max(
        1,
        context_window_tokens - max_output_tokens - safety_margin_tokens,
    )
    warn_threshold_tokens = max(1, math.floor(input_budget_tokens * warn_ratio))
    compact_threshold_tokens = max(1, math.floor(input_budget_tokens * compact_ratio))
    source = "profile" if profile.context_window_tokens else (
        "preset" if spec.default_context_window_tokens else "config"
    )
    return ContextWindowResolution(
        context_window_tokens=context_window_tokens,
        max_output_tokens=max_output_tokens,
        compact_threshold_tokens=compact_threshold_tokens,
        warn_threshold_tokens=warn_threshold_tokens,
        input_budget_tokens=input_budget_tokens,
        source=source,
    )


def estimate_prompt_tokens(
    messages: list[dict[str, Any]],
    tools: list[dict[str, Any]] | None = None,
) -> int:
    """Estimate full prompt token usage conservatively."""
    total = 0
    for message in messages:
        total += estimate_message_tokens(message)
    if tools:
        total += estimate_json_tokens(tools)
    return max(1, total)


def estimate_message_tokens(message: dict[str, Any]) -> int:
    """Estimate one chat message's prompt contribution."""
    payload_parts: list[str] = [str(message.get("role") or "message")]
    content = message.get("content")
    if isinstance(content, str):
        payload_parts.append(content)
    elif isinstance(content, list):
        for item in content:
            if isinstance(item, dict) and item.get("type") == "text":
                payload_parts.append(str(item.get("text") or ""))
            else:
                payload_parts.append(json.dumps(item, ensure_ascii=False))
    elif content is not None:
        payload_parts.append(json.dumps(content, ensure_ascii=False))

    for key in ("name", "tool_call_id"):
        value = message.get(key)
        if value:
            payload_parts.append(str(value))
    tool_calls = message.get("tool_calls")
    if tool_calls:
        payload_parts.append(json.dumps(tool_calls, ensure_ascii=False))

    return estimate_text_tokens("\n".join(part for part in payload_parts if part))


def estimate_json_tokens(value: Any) -> int:
    """Estimate tokens for a JSON payload."""
    return estimate_text_tokens(json.dumps(value, ensure_ascii=False, sort_keys=True))


def estimate_text_tokens(text: str) -> int:
    """Estimate text tokens conservatively across ASCII and CJK-heavy payloads."""
    if not text:
        return 1
    ascii_count = sum(1 for char in text if char.isascii())
    non_ascii_count = max(0, len(text) - ascii_count)
    whitespace_count = sum(1 for char in text if char.isspace())
    visible_ascii = max(0, ascii_count - whitespace_count)
    estimated = (
        visible_ascii / 3.2
        + whitespace_count / 12
        + non_ascii_count * 1.15
    )
    return max(1, math.ceil(estimated))


def is_context_overflow_error_text(text: str | None) -> bool:
    """Return True when provider error text looks like a context overflow."""
    if not text:
        return False
    lowered = text.strip().lower()
    return any(keyword in lowered for keyword in _OVERFLOW_KEYWORDS)


def _first_positive_int(*values: int | None, fallback: int) -> int:
    for value in values:
        if isinstance(value, int) and value > 0:
            return value
    return fallback


def _normalize_ratio(value: float | int, *, fallback: float) -> float:
    if not isinstance(value, (float, int)):
        return fallback
    normalized = float(value)
    if normalized <= 0:
        return fallback
    return min(0.99, normalized)
