from __future__ import annotations

import time
from typing import Literal

from ..config import mcp_config

ThrottleCategory = Literal["upload", "status", "download_url", "marketplace_list"]

_CATEGORY_LIMITS: dict[ThrottleCategory, tuple[int, int]] = {
    "upload": (mcp_config.throttle.upload_window_ms, mcp_config.throttle.upload_max_requests),
    "status": (mcp_config.throttle.status_window_ms, mcp_config.throttle.status_max_requests),
    "download_url": (mcp_config.throttle.download_url_window_ms, mcp_config.throttle.download_url_max_requests),
    "marketplace_list": (60 * 60 * 1000, 5),  # 5 per hour
}

_TOOL_CATEGORIES: dict[str, ThrottleCategory] = {
    "upload_file": "upload",
    "upload_from_url": "upload",
    "upload_and_wait": "upload",
    "get_upload_status": "status",
    "get_download_url": "download_url",
    "list_on_marketplace": "marketplace_list",
}

_windows: dict[str, list[float]] = {}


def check_throttle(tool_name: str, token_prefix: str) -> tuple[bool, str]:
    """Returns (ok, error_message). error_message is empty on success."""
    category = _TOOL_CATEGORIES.get(tool_name)
    if category is None:
        return True, ""

    window_ms, max_requests = _CATEGORY_LIMITS[category]
    key = f"{token_prefix}:{category}"
    now = time.monotonic() * 1000
    cutoff = now - window_ms

    timestamps = _windows.get(key, [])
    timestamps = [t for t in timestamps if t > cutoff]

    if len(timestamps) >= max_requests:
        oldest = timestamps[0]
        retry_after_ms = oldest + window_ms - now
        retry_after_sec = max(1, int(retry_after_ms / 1000) + 1)
        _windows[key] = timestamps
        return (
            False,
            f"Rate limit: max {max_requests} {tool_name} calls per {window_ms // 60_000} minutes. "
            f"Retry after {retry_after_sec}s.",
        )

    timestamps.append(now)
    _windows[key] = timestamps
    return True, ""
