api_handler.py

"""
api_handler.py — AI Provider Integration + Smart Error Classification

Error types and handling strategy:
  RateLimitError    → exponential backoff, wait up to 60s, retry 6 times
  TokenLimitError   → caller should split chunk smaller (not retryable here)
  QuotaError        → not retryable, raise immediately with clear message
  NetworkError      → retry 4 times with backoff
  AuthError         → not retryable, raise immediately

Public API
----------
fetch_models(provider, api_key)                         -> list[str]
validate_api_key(provider, api_key)                     -> (bool, str)
call_ai(prompt, *, provider, model, api_key,
        system_prompt, timeout)                         -> str
classify_error(exception)                               -> ErrorKind
"""

import logging
import time
import re
import requests

logger = logging.getLogger(__name__)

# ─────────────────────────────────────────────────────────────────────────────
# Error classification
# ─────────────────────────────────────────────────────────────────────────────

class RateLimitError(Exception):
    """429 — slow down, retry after backoff."""
    def __init__(self, msg, retry_after: int = 0):
        super().__init__(msg)
        self.retry_after = retry_after   # seconds hint from header


class TokenLimitError(Exception):
    """Output/context token limit exceeded — caller must split chunk."""


class QuotaError(Exception):
    """Daily/monthly quota exhausted — cannot retry today."""


class AuthError(Exception):
    """Invalid API key — not retryable."""


class NetworkError(Exception):
    """Transient network failure — retry."""


def classify_error(exc: Exception):
    """
    Inspect an exception and return the most specific error subclass.
    Works for both Gemini (google-generativeai) and OpenRouter (openai SDK).
    """
    msg = str(exc).lower()
    code = getattr(exc, "status_code", None) or getattr(exc, "code", None)
    # Try to extract HTTP status from the message string if not on object
    if code is None:
        m = re.search(r'\b(4\d\d|5\d\d)\b', str(exc))
        if m:
            code = int(m.group(1))

    # Rate limit
    if code == 429 or "rate" in msg or "quota" in msg and "per" in msg:
        # Try to parse retry-after from message
        ra = 0
        m = re.search(r'retry.{0,10}after.{0,10}(\d+)', msg)
        if m:
            ra = int(m.group(1))
        return RateLimitError(str(exc), retry_after=ra)

    # Quota exhausted (daily limit)
    if code == 429 and ("daily" in msg or "exhausted" in msg or "exceeded" in msg):
        return QuotaError(str(exc))
    if "resource_exhausted" in msg or "quota_exceeded" in msg:
        return QuotaError(str(exc))

    # Token / context limit
    if any(k in msg for k in [
        "token", "context", "too long", "maximum context",
        "max_tokens", "finish_reason: length", "content_filter"
    ]):
        return TokenLimitError(str(exc))

    # Auth
    if code in (401, 403) or any(k in msg for k in ["api key", "unauthorized", "forbidden", "invalid key"]):
        return AuthError(str(exc))

    # Network
    if isinstance(exc, (requests.exceptions.ConnectionError,
                         requests.exceptions.Timeout,
                         ConnectionError, TimeoutError)):
        return NetworkError(str(exc))
    if code and code >= 500:
        return NetworkError(str(exc))

    return exc   # unknown — return original


# ─────────────────────────────────────────────────────────────────────────────
# Constants
# ─────────────────────────────────────────────────────────────────────────────

OPENROUTER_BASE   = "https://openrouter.ai/api/v1"
OPENROUTER_MODELS = f"{OPENROUTER_BASE}/models"
GEMINI_MODELS_URL = "https://generativelanguage.googleapis.com/v1beta/models"
DEFAULT_TIMEOUT   = 60

GEMINI_FALLBACK = [
    "gemini-2.0-flash", "gemini-2.0-flash-lite",
    "gemini-1.5-pro", "gemini-1.5-flash", "gemini-1.5-flash-8b",
]
OPENROUTER_FALLBACK = [
    "google/gemini-2.0-flash-exp:free",
    "meta-llama/llama-3.3-70b-instruct:free",
    "deepseek/deepseek-chat-v3-0324:free",
    "mistralai/mistral-7b-instruct:free",
    "qwen/qwen-2.5-72b-instruct:free",
]


# ─────────────────────────────────────────────────────────────────────────────
# Model listing
# ─────────────────────────────────────────────────────────────────────────────

def fetch_gemini_models(api_key: str) -> list:
    if not api_key or not api_key.strip():
        return GEMINI_FALLBACK.copy()
    try:
        resp = requests.get(
            GEMINI_MODELS_URL,
            params={"key": api_key.strip(), "pageSize": 100},
            timeout=15,
        )
        resp.raise_for_status()
        models = []
        for m in resp.json().get("models", []):
            name = m.get("name", "").replace("models/", "")
            if "generateContent" in m.get("supportedGenerationMethods", []):
                models.append(name)
        models.sort(key=lambda x: (0 if x.startswith("gemini-2") else
                                   1 if x.startswith("gemini-1.5") else 2, x))
        return models or GEMINI_FALLBACK.copy()
    except Exception as exc:
        logger.warning("Gemini model fetch: %s", exc)
        return GEMINI_FALLBACK.copy()


def fetch_openrouter_models(api_key: str) -> list:
    if not api_key or not api_key.strip():
        return OPENROUTER_FALLBACK.copy()
    try:
        resp = requests.get(
            OPENROUTER_MODELS,
            headers={
                "Authorization": f"Bearer {api_key.strip()}",
                "HTTP-Referer": "https://huggingface.co/spaces",
                "X-Title": "SubSync Myanmar Translator",
            },
            timeout=15,
        )
        resp.raise_for_status()
        raw = resp.json().get("data", [])
        raw.sort(key=lambda m: (-m.get("context_length", 0), m.get("id", "")))
        ids = [m["id"] for m in raw if "id" in m]
        return ids or OPENROUTER_FALLBACK.copy()
    except Exception as exc:
        logger.warning("OpenRouter model fetch: %s", exc)
        return OPENROUTER_FALLBACK.copy()


def fetch_models(provider: str, api_key: str) -> list:
    if provider == "gemini":
        return fetch_gemini_models(api_key)
    if provider == "openrouter":
        return fetch_openrouter_models(api_key)
    raise ValueError(f"Unknown provider: {provider!r}")


def validate_api_key(provider: str, api_key: str) -> tuple:
    if not api_key or not api_key.strip():
        return False, "API key လိုအပ်သည်"
    try:
        if provider == "gemini":
            resp = requests.get(GEMINI_MODELS_URL,
                                params={"key": api_key.strip(), "pageSize": 5},
                                timeout=15)
            if resp.status_code == 400: return False, "API key မမှန်ပါ (400)"
            if resp.status_code == 403: return False, "API key ခွင့်မပြု (403)"
            resp.raise_for_status()
            return True, f"✓ Valid — {len(resp.json().get('models',[]))} models"
        if provider == "openrouter":
            resp = requests.get(OPENROUTER_MODELS,
                                headers={"Authorization": f"Bearer {api_key.strip()}",
                                         "HTTP-Referer": "https://huggingface.co/spaces"},
                                timeout=15)
            if resp.status_code == 401: return False, "API key မမှန်ပါ (401)"
            resp.raise_for_status()
            return True, f"✓ Valid — {len(resp.json().get('data',[]))} models"
        return False, f"Unknown provider: {provider}"
    except requests.exceptions.ConnectionError:
        return False, "Network error — internet စစ်ဆေးပါ"
    except requests.exceptions.Timeout:
        return False, "Request timeout"
    except Exception as exc:
        return False, f"Error: {exc}"


# ─────────────────────────────────────────────────────────────────────────────
# Inference — with smart retry per error type
# ─────────────────────────────────────────────────────────────────────────────

def _smart_retry_call(fn, max_attempts: int = 6):
    """
    Call fn() with smart per-error-type retry strategy:
      RateLimitError  → wait retry_after (or exp backoff 5s→120s), retry up to 6x
      NetworkError    → exp backoff 2s→30s, retry up to 4x
      TokenLimitError → raise immediately (caller splits chunk)
      QuotaError      → raise immediately (user must wait/switch key)
      AuthError       → raise immediately (wrong key)
    """
    rate_delay = 5    # starting backoff for rate limits
    net_delay  = 2    # starting backoff for network errors
    rate_tries = 0
    net_tries  = 0

    for attempt in range(1, max_attempts + 1):
        try:
            return fn()

        except Exception as raw_exc:
            classified = classify_error(raw_exc)

            if isinstance(classified, (TokenLimitError, QuotaError, AuthError)):
                raise classified from raw_exc

            if isinstance(classified, RateLimitError):
                rate_tries += 1
                wait = classified.retry_after if classified.retry_after > 0 else rate_delay
                rate_delay = min(rate_delay * 2, 120)
                if rate_tries >= max_attempts:
                    raise classified from raw_exc
                logger.warning("Rate limit — waiting %ds (attempt %d/%d)",
                               wait, attempt, max_attempts)
                time.sleep(wait)
                continue

            if isinstance(classified, NetworkError):
                net_tries += 1
                if net_tries >= 4:
                    raise classified from raw_exc
                logger.warning("Network error — waiting %ds (attempt %d)",
                               net_delay, attempt)
                time.sleep(net_delay)
                net_delay = min(net_delay * 2, 30)
                continue

            # Unknown error — limited retry
            if attempt >= 3:
                raise
            time.sleep(3)

    raise RuntimeError(f"All {max_attempts} attempts failed")


def _call_gemini(prompt: str, *, model: str, api_key: str,
                 system_prompt: str = "", timeout: int = DEFAULT_TIMEOUT) -> str:
    import google.generativeai as genai
    from google.generativeai.types import HarmCategory, HarmBlockThreshold

    genai.configure(api_key=api_key.strip())
    gen_cfg = genai.GenerationConfig(temperature=0.3, max_output_tokens=8192)
    safety  = {
        HarmCategory.HARM_CATEGORY_HARASSMENT:        HarmBlockThreshold.BLOCK_NONE,
        HarmCategory.HARM_CATEGORY_HATE_SPEECH:       HarmBlockThreshold.BLOCK_NONE,
        HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
        HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
    }
    gemini_m = genai.GenerativeModel(
        model_name=model, generation_config=gen_cfg,
        safety_settings=safety,
        system_instruction=system_prompt or None,
    )
    return _smart_retry_call(lambda: gemini_m.generate_content(prompt).text)


def _call_openrouter(prompt: str, *, model: str, api_key: str,
                     system_prompt: str = "", timeout: int = DEFAULT_TIMEOUT) -> str:
    from openai import OpenAI
    client = OpenAI(
        base_url=OPENROUTER_BASE, api_key=api_key.strip(),
        default_headers={
            "HTTP-Referer": "https://huggingface.co/spaces",
            "X-Title": "SubSync Myanmar Translator",
        },
        timeout=timeout,
    )
    messages = []
    if system_prompt:
        messages.append({"role": "system", "content": system_prompt})
    messages.append({"role": "user", "content": prompt})

    def _call():
        resp = client.chat.completions.create(
            model=model, messages=messages, temperature=0.3, max_tokens=8192
        )
        return resp.choices[0].message.content or ""

    return _smart_retry_call(_call)


def call_ai(prompt: str, *, provider: str, model: str, api_key: str,
            system_prompt: str = "", timeout: int = DEFAULT_TIMEOUT) -> str:
    """
    Unified AI call with smart error classification and retry.
    Raises:
      RateLimitError  — hit rate limit after all retries
      TokenLimitError — chunk too long for model (caller should split)
      QuotaError      — daily quota exhausted (cannot retry)
      AuthError       — invalid API key
      RuntimeError    — other failure
    """
    if not api_key or not api_key.strip():
        raise AuthError("API key မပါပါ — Settings တွင် ထည့်ပါ")
    if not model:
        raise ValueError("Model မရွေးရသေးပါ — Settings တွင် ရွေးပါ")
    try:
        if provider == "gemini":
            return _call_gemini(prompt, model=model, api_key=api_key,
                                system_prompt=system_prompt, timeout=timeout)
        if provider == "openrouter":
            return _call_openrouter(prompt, model=model, api_key=api_key,
                                    system_prompt=system_prompt, timeout=timeout)
        raise ValueError(f"Unknown provider: {provider!r}")
    except (RateLimitError, TokenLimitError, QuotaError, AuthError, ValueError):
        raise
    except Exception as exc:
        classified = classify_error(exc)
        if isinstance(classified, Exception) and classified is not exc:
            raise classified from exc
        raise RuntimeError(f"AI call failed ({provider}): {exc}") from exc