"""
ai_engine.py — BIRAS Unified AI Gateway
========================================
Single entry point for ALL LLM calls in the system.

Usage:
    from ai_engine import call_ai

    result = call_ai(
        prompt="Explain this trend...",
        model='openai',           # openai | gemini | deepseek | zai | groq | openrouter | auto
        max_tokens=300,
        fallback=True,            # try next provider if this one fails
    )

Design:
  - Provider-agnostic interface
  - Automatic fallback chain: openai → gemini → deepseek
  - In-process LRU cache (keyed on prompt hash + model)
  - Rate-limit safe (exponential back-off, max 2 retries)
  - Returns str | None (never raises)
"""

from __future__ import annotations

import os
import time
import hashlib
import logging
from functools import lru_cache
from typing import Any, Optional

logger = logging.getLogger('biras.ai')

# ── Provider registry ──────────────────────────────────────────────────────────
_PROVIDERS = {
    'openai': {
        'env_key':   'OPENAI_API_KEY',
        'base_url':  None,
        'model':     'gpt-4o-mini',
    },
    'deepseek': {
        'env_key':   'DEEPSEEK_API_KEY',
        'base_url':  'https://api.deepseek.com',
        'model':     'deepseek-chat',
    },
    'zai': {
        'env_key':   'ZAI_API_KEY',
        'base_url':  lambda: os.environ.get('ZAI_BASE_URL', 'https://api.zai.ai'),
        'model':     'glm-4-plus',
    },
    'groq': {
        'env_key':   'GROQ_API_KEY',
        'base_url':  'https://api.groq.com/openai/v1',
        'model':     'llama-3.3-70b-versatile',
    },
    'openrouter': {
        'env_key':   'OPENROUTER_API_KEY',
        'base_url':  'https://openrouter.ai/api/v1',
        'model':     'openrouter/free',
    },
    'github': {
        'env_key':   'GITHUB_API_KEY',
        'base_url':  'https://models.inference.ai.azure.com',
        'model':     'gpt-4o',
    },
    'huggingface': {
        'env_key':   'HUGGINGFACE_API_KEY',
        'base_url':  'https://router.huggingface.co/v1',
        'model':     'Qwen/Qwen2.5-72B-Instruct',
    },
    'sambanova': {
        'env_key':   'SAMBANOVA_API_KEY',
        'base_url':  'https://api.sambanova.ai/v1',
        'model':     'Meta-Llama-3.3-70B-Instruct',
    },
    'cerebras': {
        'env_key':   'CEREBRAS_API_KEY',
        'base_url':  'https://api.cerebras.ai/v1',
        'model':     'llama3.3-70b',
    },
    'together': {
        'env_key':   'TOGETHER_API_KEY',
        'base_url':  'https://api.together.xyz/v1',
        'model':     'meta-llama/Llama-3.2-11B-Vision-Instruct-Turbo',
    },
    'mistral': {
        'env_key':   'MISTRAL_API_KEY',
        'base_url':  'https://api.mistral.ai/v1',
        'model':     'mistral-small-latest',
    },
}

_FALLBACK_CHAIN = ['cerebras', 'groq', 'github', 'together', 'mistral', 'gemini', 'openai', 'sambanova', 'deepseek', 'openrouter', 'huggingface', 'zai']

# Simple in-memory cache: {hash → response_text}
_AI_CACHE: dict[str, str] = {}

# OpenAI client cache: avoid re-initializing on every call
_CLIENT_CACHE: dict[str, Any] = {}


def _cache_key(prompt: str, model: str) -> str:
    return hashlib.md5(f"{model}::{prompt}".encode(), usedforsecurity=False).hexdigest()


def _get_openai_client(api_key: str, base_url: Optional[str]) -> Any:
    """Return a cached OpenAI client. Creates a new one only when key/url changes."""
    from openai import OpenAI
    cache_key = f"{api_key[:12]}::{base_url or 'default'}"
    if cache_key not in _CLIENT_CACHE:
        kwargs: dict = {'api_key': api_key}
        if base_url:
            kwargs['base_url'] = base_url
        _CLIENT_CACHE[cache_key] = OpenAI(**kwargs)
    return _CLIENT_CACHE[cache_key]


# ─────────────────────────────────────────────────────────────────────────────
def _call_openai_compatible(
    prompt: str,
    api_key: str,
    base_url: Optional[str],
    model: str,
    max_tokens: int,
) -> Optional[str]:
    """Call any OpenAI-compatible REST API (OpenAI, DeepSeek, ZAI)."""
    try:
        # Re-use cached client — avoids TCP session overhead on every call
        client = _get_openai_client(api_key, base_url)
        resp = client.chat.completions.create(
            model=model,
            messages=[
                {
                    'role': 'system',
                    'content': (
                        'You are an expert academic researcher and scientific writer. '
                        'Respond in formal English. Be concise and precise.'
                    ),
                },
                {'role': 'user', 'content': prompt},
            ],
            max_tokens=max_tokens,
            temperature=0.4,
        )
        text = resp.choices[0].message.content or ''
        return text.strip() or None
    except Exception as exc:
        logger.warning("OpenAI-compatible call failed (model=%s): %s", model, exc)
        return None


def _call_gemini(prompt: str, max_tokens: int) -> Optional[str]:
    """Call Google Gemini API."""
    api_key = os.environ.get('GEMINI_API_KEY', '')
    if not api_key:
        return None
    try:
        import google.generativeai as genai
        genai.configure(api_key=api_key)
        model = genai.GenerativeModel('gemini-2.0-flash')
        resp = model.generate_content(
            prompt,
            generation_config=genai.types.GenerationConfig(
                max_output_tokens=max_tokens,
                temperature=0.4,
            ),
        )
        text = resp.text or ''
        return text.strip() or None
    except Exception as exc:
        logger.warning("Gemini call failed: %s", exc)
        return None


def _try_provider(provider: str, prompt: str, max_tokens: int) -> Optional[str]:
    """Dispatch to a specific provider. Returns None on failure."""
    if provider == 'gemini':
        return _call_gemini(prompt, max_tokens)

    cfg = _PROVIDERS.get(provider)
    if not cfg:
        logger.warning("Unknown AI provider: %s", provider)
        return None

    api_key = os.environ.get(cfg['env_key'], '')
    if not api_key:
        logger.debug("Provider '%s' skipped (no API key)", provider)
        return None

    base_url = cfg['base_url']
    if callable(base_url):
        base_url = base_url()

    return _call_openai_compatible(
        prompt=prompt,
        api_key=api_key,
        base_url=base_url,
        model=cfg['model'],
        max_tokens=max_tokens,
    )


# ─────────────────────────────────────────────────────────────────────────────
def call_ai(
    prompt: str,
    model: str = 'auto',
    max_tokens: int = 400,
    fallback: bool = True,
    use_cache: bool = True,
    retries: int = 1,
) -> Optional[str]:
    """
    Unified AI gateway for BIRAS.

    Args:
        prompt:     The text prompt to send.
        model:      'openai' | 'gemini' | 'deepseek' | 'zai' | 'groq' | 'openrouter' | 'github' | 'huggingface' | 'sambanova' | 'auto'
                    'auto' → tries DEFAULT_AI_MODEL, then fallback chain.
        max_tokens: Maximum response tokens.
        fallback:   If True and primary model fails, try fallback chain.
        use_cache:  If True, cache responses in memory (keyed on prompt hash).
        retries:    Number of retry attempts with exponential back-off.

    Returns:
        Response string, or None if all providers fail.
    """
    if not prompt or len(prompt.strip()) < 5:
        return None

    # ── Cache lookup ────────────────────────────────────────────────────
    cache_enabled = use_cache and os.getenv('ENABLE_AI_CACHE', 'true').lower() == 'true'
    key = _cache_key(prompt, model)
    if cache_enabled and key in _AI_CACHE:
        logger.debug("AI cache hit (model=%s, chars=%d)", model, len(prompt))
        return _AI_CACHE[key]

    # ── Build provider order ────────────────────────────────────────────
    default_model = os.getenv('DEFAULT_AI_MODEL', 'openai')
    if model == 'auto':
        order = [default_model] + [p for p in _FALLBACK_CHAIN if p != default_model]
    else:
        order = [model] + (
            [p for p in _FALLBACK_CHAIN if p != model] if fallback else []
        )

    # ── Try each provider ───────────────────────────────────────────────
    for provider in order:
        for attempt in range(retries + 1):
            result = _try_provider(provider, prompt, max_tokens)
            if result and len(result) > 10:
                logger.info("AI response via %s (%d chars)", provider, len(result))
                if cache_enabled:
                    _AI_CACHE[key] = result
                return result
            if attempt < retries:
                wait = 2 ** attempt
                logger.debug("Retrying provider %s in %ds...", provider, wait)
                time.sleep(wait)

    logger.warning("All AI providers failed for prompt: %s...", prompt[:60])
    return None


def clear_ai_cache() -> int:
    """Clear the in-memory AI cache. Returns number of entries cleared."""
    n = len(_AI_CACHE)
    _AI_CACHE.clear()
    logger.info("AI cache cleared (%d entries)", n)
    return n