"""LLM integration with Featherless API.

Implements model fallback chain with 4 models + 3 full passes.
Logs every attempt to stderr for operational visibility.

Custom exceptions: LLMError (all models failed), _ModelFailure (single attempt).
"""

import os
import sys
import time
import requests
from datetime import datetime
from typing import Optional, Tuple, Dict, Any, List


DEFAULT_MODEL = "deepseek-ai/DeepSeek-V3.1"
DEFAULT_API_BASE = "https://api.featherless.ai/v1"
DEFAULT_TEMPERATURE = 0.7
DEFAULT_MAX_TOKENS = 600

TIMEOUT = 30
MAX_CHAIN_RETRIES = 3
CHAIN_RETRY_WAIT = 30
MAX_RETRIES_PER_MODEL = 1
RATE_LIMIT_BACKOFF = 5
GATEWAY_RETRY_BACKOFF = 3

FALLBACK_MODELS = [
    "deepseek-ai/DeepSeek-V3.1",
    "Qwen/Qwen3.6-27B",
    "google/gemma-4-31B-it",
    "mistralai/Mistral-7B-Instruct-v0.2",
]


class LLMError(Exception):
    def __init__(self, error_type, detail, status_code=None, model=None, attempts=None):
        self.error_type = error_type
        self.detail = detail
        self.status_code = status_code
        self.model = model
        self.attempts = attempts or []
        super().__init__(f"LLMError({error_type}): {detail}")


class _ModelFailure(Exception):
    def __init__(self, error_type, detail, status_code=None):
        self.error_type = error_type
        self.detail = detail
        self.status_code = status_code


def _now() -> str:
    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")


def _log(msg):
    print(f"[LLM {_now()}] {msg}", file=sys.stderr)


def _get_api_config() -> Dict[str, str]:
    return {
        "api_key": os.environ.get("FEATHERLESS_API_KEY", ""),
        "api_base": os.environ.get("FEATHERLESS_API_BASE", DEFAULT_API_BASE),
    }


def _try_model(
    model: str,
    prompt: str,
    temperature: float,
    max_tokens: int,
    api_key: str,
    api_base: str,
    timeout: int = TIMEOUT,
    max_retries: int = MAX_RETRIES_PER_MODEL,
) -> Tuple[str, Dict[str, Any]]:
    """Try one model with retry logic. Raises _ModelFailure on failure.

    Error handling order:
    1. HTTP timeout → _ModelFailure("timeout")
    2. Connection error → _ModelFailure("connection_error")
    3. 429 → retry once after 5s, then _ModelFailure("rate_limit")
    4. 502/503/504 → check for "at capacity":
       - "at capacity": _ModelFailure("capacity") immediately (no retry)
       - Otherwise: retry once after 3s, then _ModelFailure("gateway")
    5. Other non-200 → _ModelFailure("http_error")
    6. JSON parse failure → _ModelFailure("parse_error")
    7. Empty response → _ModelFailure("empty_response")
    """
    url = f"{api_base.rstrip('/')}/chat/completions"

    attempt = 0
    while attempt <= max_retries:
        try:
            start_time = time.time()

            resp = requests.post(
                url,
                headers={
                    "Authorization": f"Bearer {api_key}",
                    "Content-Type": "application/json",
                },
                json={
                    "model": model,
                    "messages": [{"role": "user", "content": prompt}],
                    "temperature": temperature,
                    "max_tokens": max_tokens,
                },
                timeout=timeout,
            )

            status = resp.status_code
            body_text = resp.text

            # 429 — rate limited
            if status == 429:
                if attempt < max_retries:
                    _log(f"{model} rate limited — retrying in {RATE_LIMIT_BACKOFF}s")
                    time.sleep(RATE_LIMIT_BACKOFF)
                    attempt += 1
                    continue
                raise _ModelFailure("rate_limit", f"HTTP 429 after {max_retries+1} attempts")

            # 502/503/504 — gateway errors
            if status in (502, 503, 504):
                body_lower = body_text.lower()
                if "at capacity" in body_lower or "capacity" in body_lower:
                    _log(f"{model} at capacity — skipping to next model")
                    raise _ModelFailure("capacity", "Model at capacity")

                if attempt < max_retries:
                    _log(f"{model} HTTP {status} — retrying in {GATEWAY_RETRY_BACKOFF}s")
                    time.sleep(GATEWAY_RETRY_BACKOFF)
                    attempt += 1
                    continue
                raise _ModelFailure("gateway", f"HTTP {status} after {max_retries+1} attempts")

            # Other HTTP errors
            if status != 200:
                raise _ModelFailure("http_error", f"HTTP {status}")

            # Parse JSON
            result = resp.json()
            latency_ms = int((time.time() - start_time) * 1000)

            # Extract content with content/reasoning fallback
            try:
                message = result["choices"][0]["message"]
            except (KeyError, IndexError):
                raise _ModelFailure("empty_response", "No choices in response")

            content = (message.get("content") or "").strip()
            reasoning = (message.get("reasoning") or "").strip()

            if not content and reasoning:
                content = reasoning
                _log(f"{model} used reasoning field fallback (GLM-5.1)")

            if not content:
                raise _ModelFailure("empty_response", "Both content and reasoning empty")

            metadata = {
                "model_used": model,
                "finish_reason": result["choices"][0].get("finish_reason", "unknown"),
                "latency_ms": latency_ms,
                "tokens_used": result.get("usage", {}).get("total_tokens", 0),
            }

            _log(f"{model} succeeded in {latency_ms}ms")
            return content, metadata

        except requests.Timeout:
            _log(f"{model} timed out ({timeout}s)")
            raise _ModelFailure("timeout", f"Request timed out after {timeout}s")

        except requests.ConnectionError as e:
            _log(f"{model} connection error: {e}")
            raise _ModelFailure("connection_error", str(e))

        except _ModelFailure:
            raise

        except Exception as e:
            if attempt < max_retries:
                _log(f"{model} unexpected error: {e} — retrying")
                time.sleep(1)
                attempt += 1
                continue
            _log(f"{model} unexpected error after {max_retries+1} attempts: {e}")
            raise _ModelFailure("unexpected", str(e))

    raise _ModelFailure("exhausted", "All retries exhausted")


def _get_models() -> List[str]:
    """Return the model chain from config or default."""
    # Check env var first
    env_models = os.environ.get("FEATHERLESS_MODELS_FALLBACK", "")
    if env_models:
        return [m.strip() for m in env_models.split(",") if m.strip()]

    # Check if there's a primary model set — insert it at position 0
    primary = os.environ.get("FEATHERLESS_MODEL", "")
    if primary and primary not in FALLBACK_MODELS:
        return [primary] + FALLBACK_MODELS

    return list(FALLBACK_MODELS)


def call_llm(
    prompt: str,
    model: Optional[str] = None,
    temperature: float = DEFAULT_TEMPERATURE,
    max_tokens: int = DEFAULT_MAX_TOKENS,
) -> Tuple[str, Dict[str, Any]]:
    """Call Featherless with model fallback chain.

    Tries up to 4 models across 3 passes (12 total attempts).
    Each pass: try all models. Between passes: wait 30s for capacity.

    Args:
        prompt: The prompt to send
        model: Override model (bypasses chain if set)
        temperature: LLM temperature
        max_tokens: Max output tokens

    Returns:
        Tuple of (content_string, metadata_dict)

    Raises:
        LLMError: All models failed after all passes
        ValueError: No API key configured
    """
    config = _get_api_config()
    api_key = config["api_key"]
    api_base = config["api_base"]

    if not api_key:
        raise ValueError("FEATHERLESS_API_KEY environment variable not set")

    # If caller overrides model, just use that (no chain)
    if model is not None:
        content, meta = _try_model(model, prompt, temperature, max_tokens, api_key, api_base)
        return content, meta

    models = _get_models()
    all_attempts = []

    for chain_pass in range(MAX_CHAIN_RETRIES):
        pass_errors = []

        for idx, mdl in enumerate(models):
            position = "primary" if idx == 0 else f"fallback_{idx}"

            try:
                content, meta = _try_model(mdl, prompt, temperature, max_tokens, api_key, api_base)
                meta["model_position"] = position
                meta["chain_pass"] = chain_pass + 1
                return content, meta

            except _ModelFailure as e:
                pass_errors.append({
                    "model": mdl,
                    "position": position,
                    "error_type": e.error_type,
                    "detail": e.detail,
                })
                _log(f"{position} ({mdl}) failed: {e.error_type}. Moving to next model...")

        all_attempts.extend(pass_errors)

        if chain_pass < MAX_CHAIN_RETRIES - 1:
            _log(f"Pass {chain_pass+1} complete — all {len(models)} models failed. Waiting {CHAIN_RETRY_WAIT}s...")
            time.sleep(CHAIN_RETRY_WAIT)

    raise LLMError(
        error_type="all_models_failed",
        detail=f"All {len(models)} models failed after {MAX_CHAIN_RETRIES} passes",
        attempts=all_attempts,
        model=models[0] if models else None,
    )


def health_check() -> Dict[str, Any]:
    """Test each model in the chain with a minimal prompt (max_tokens=5).

    Returns:
        Dict with 'models' (list of status dicts) and 'any_available' (bool).
    """
    config = _get_api_config()
    api_key = config["api_key"]
    api_base = config["api_base"]

    if not api_key:
        return {"models": [], "any_available": False, "error": "No API key configured"}

    models = _get_models()
    results = []

    for idx, mdl in enumerate(models):
        position = "primary" if idx == 0 else f"fallback_{idx}"
        try:
            _, meta = _try_model(mdl, "Hello", temperature=0.0, max_tokens=5,
                                api_key=api_key, api_base=api_base, timeout=15, max_retries=0)
            results.append({
                "model": mdl,
                "position": position,
                "status": "ok",
                "latency_ms": meta.get("latency_ms", 0),
            })
        except _ModelFailure as e:
            results.append({
                "model": mdl,
                "position": position,
                "status": e.error_type,
                "detail": e.detail,
            })

    available = sum(1 for r in results if r["status"] == "ok")
    return {
        "models": results,
        "any_available": available > 0,
        "available_count": available,
        "total_count": len(models),
    }