#!/usr/bin/env python3
"""
AI Router — service functions for chat completions.

Primary: Ollama Cloud (deepseek-v4-flash, free).
Fallback: OpenRouter (various models via OPENROUTER_API_KEY).

Both endpoints are OpenAI-compatible /v1/chat/completions.
Called by app/routers/chat.py and app/routers/wallet_clustering_router.py
as ``ai_router.chat_completion(...)`` and
``ai_router.stream_chat_completion(...)``.
"""

import base64
import json
import logging
import os
import time
from collections.abc import AsyncGenerator
from typing import Any

import httpx
from fastapi import APIRouter

logger = logging.getLogger("rmi.ai_router")

router = APIRouter(tags=["AI Router"])

# Decode base64 LLM key if present, otherwise use plain LLM_API_KEY
if os.getenv("LLM_API_KEY_B64"):
    os.environ["LLM_API_KEY"] = base64.b64decode(os.getenv("LLM_API_KEY_B64")).decode()

# ── Provider config ─────────────────────────────────────────────
OLLAMA_KEY = os.getenv("OLLAMA_API_KEY", "")
OLLAMA_URL = "https://ollama.com/v1/chat/completions"
OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "deepseek-v4-flash")

OPENROUTER_KEY = os.getenv("OPENROUTER_API_KEY", "")
OPENROUTER_URL = "https://openrouter.ai/api/v1/chat/completions"
OPENROUTER_MODEL = os.getenv("OPENROUTER_MODEL", "deepseek/deepseek-chat")

# Tier → model mapping (best effort; Ollama Cloud ignores non-model fields)
TIER_MODELS: dict[str, str] = {
    "T0": "deepseek-v4-pro",
    "T1": "deepseek-v4-pro",
    "T2": "deepseek-v4-flash",
    "T3": "deepseek-v4-flash",
    "T4": "deepseek-v4-flash",
}

MODEL_TIERS = {
    "T0": {"name": "Ultra", "models": ["deepseek-v4-pro"], "max_cost_per_1k": 0.0},
    "T1": {"name": "Premium", "models": ["deepseek-v4-pro"], "max_cost_per_1k": 0.0},
    "T2": {"name": "Standard", "models": ["deepseek-v4-flash"], "max_cost_per_1k": 0.0},
    "T3": {"name": "Fast", "models": ["deepseek-v4-flash"], "max_cost_per_1k": 0.0},
    "T4": {"name": "Free", "models": ["deepseek-v4-flash"], "max_cost_per_1k": 0.0},
}

PROVIDERS = {
    "ollama": {"url": OLLAMA_URL, "key_env": "OLLAMA_API_KEY", "model": OLLAMA_MODEL, "rpm": 100},
    "openrouter": {"url": OPENROUTER_URL, "key_env": "OPENROUTER_API_KEY", "model": OPENROUTER_MODEL, "rpm": 100},
}


def _resolve_model(model: str | None, tier: str | None) -> str:
    """Pick the model: explicit model > tier mapping > default."""
    if model:
        return model
    if tier and tier in TIER_MODELS:
        return TIER_MODELS[tier]
    return OLLAMA_MODEL


async def chat_completion(
    *,
    messages: list[dict[str, str]],
    model: str | None = None,
    tier: str | None = None,
    temperature: float = 0.3,
    max_tokens: int = 500,
    timeout: float = 30.0,
) -> dict[str, Any]:
    """Non-streaming chat completion.

    Returns dict with keys: content, model, _provider, _latency_ms, usage, error (on failure).
    Tries Ollama Cloud first, falls back to OpenRouter.
    """
    resolved = _resolve_model(model, tier)
    start = time.time()

    # Provider 1: Ollama Cloud
    if OLLAMA_KEY:
        try:
            async with httpx.AsyncClient(timeout=timeout) as client:
                resp = await client.post(
                    OLLAMA_URL,
                    headers={
                        "Authorization": f"Bearer {OLLAMA_KEY}",
                        "Content-Type": "application/json",
                    },
                    json={
                        "model": resolved,
                        "messages": messages,
                        "temperature": temperature,
                        "max_tokens": max_tokens,
                        "stream": False,
                    },
                )
                if resp.status_code == 200:
                    data = resp.json()
                    content = data.get("choices", [{}])[0].get("message", {}).get("content", "")
                    return {
                        "content": content,
                        "model": data.get("model", resolved),
                        "_provider": "ollama",
                        "_latency_ms": round((time.time() - start) * 1000, 1),
                        "usage": data.get("usage", {}),
                    }
                logger.warning("ai_router ollama status=%s body=%s", resp.status_code, resp.text[:200])
        except Exception as e:
            logger.warning("ai_router ollama error: %s", e)

    # Provider 2: OpenRouter fallback
    if OPENROUTER_KEY:
        try:
            async with httpx.AsyncClient(timeout=timeout) as client:
                resp = await client.post(
                    OPENROUTER_URL,
                    headers={
                        "Authorization": f"Bearer {OPENROUTER_KEY}",
                        "Content-Type": "application/json",
                    },
                    json={
                        "model": resolved,
                        "messages": messages,
                        "temperature": temperature,
                        "max_tokens": max_tokens,
                        "stream": False,
                    },
                )
                if resp.status_code == 200:
                    data = resp.json()
                    content = data.get("choices", [{}])[0].get("message", {}).get("content", "")
                    return {
                        "content": content,
                        "model": data.get("model", resolved),
                        "_provider": "openrouter",
                        "_latency_ms": round((time.time() - start) * 1000, 1),
                        "usage": data.get("usage", {}),
                    }
                logger.warning("ai_router openrouter status=%s body=%s", resp.status_code, resp.text[:200])
        except Exception as e:
            logger.warning("ai_router openrouter error: %s", e)

    return {
        "content": "",
        "model": resolved,
        "_provider": "none",
        "_latency_ms": round((time.time() - start) * 1000, 1),
        "usage": {},
        "error": "No AI provider available (set OLLAMA_API_KEY or OPENROUTER_API_KEY)",
    }


async def stream_chat_completion(
    *,
    messages: list[dict[str, str]],
    model: str | None = None,
    tier: str | None = None,
    temperature: float = 0.3,
    max_tokens: int = 500,
    timeout: float = 30.0,
) -> AsyncGenerator[str, None]:
    """Streaming chat completion via SSE.

    Yields content tokens as strings. On error, yields ``[ERROR: ...]``.
    Tries Ollama Cloud first, falls back to OpenRouter.
    """
    resolved = _resolve_model(model, tier)

    # Provider 1: Ollama Cloud (streaming)
    if OLLAMA_KEY:
        try:
            async with httpx.AsyncClient(timeout=timeout) as client, client.stream(
                "POST",
                OLLAMA_URL,
                headers={
                    "Authorization": f"Bearer {OLLAMA_KEY}",
                    "Content-Type": "application/json",
                },
                json={
                    "model": resolved,
                    "messages": messages,
                    "temperature": temperature,
                    "max_tokens": max_tokens,
                    "stream": True,
                },
            ) as resp:
                if resp.status_code == 200:
                    async for line in resp.aiter_lines():
                        if line.startswith("data: ") and line != "data: [DONE]":
                            try:
                                chunk = json.loads(line[6:])
                                delta = chunk.get("choices", [{}])[0].get("delta", {})
                                content = delta.get("content", "")
                                if content:
                                    yield content
                            except json.JSONDecodeError:
                                pass
                    return
                logger.warning("ai_router ollama stream status=%s", resp.status_code)
        except Exception as e:
            logger.warning("ai_router ollama stream error: %s", e)

    # Provider 2: OpenRouter streaming fallback
    if OPENROUTER_KEY:
        try:
            async with httpx.AsyncClient(timeout=timeout) as client, client.stream(
                "POST",
                OPENROUTER_URL,
                headers={
                    "Authorization": f"Bearer {OPENROUTER_KEY}",
                    "Content-Type": "application/json",
                },
                json={
                    "model": resolved,
                    "messages": messages,
                    "temperature": temperature,
                    "max_tokens": max_tokens,
                    "stream": True,
                },
            ) as resp:
                if resp.status_code == 200:
                    async for line in resp.aiter_lines():
                        if line.startswith("data: ") and line != "data: [DONE]":
                            try:
                                chunk = json.loads(line[6:])
                                delta = chunk.get("choices", [{}])[0].get("delta", {})
                                content = delta.get("content", "")
                                if content:
                                    yield content
                            except json.JSONDecodeError:
                                pass
                    return
                logger.warning("ai_router openrouter stream status=%s", resp.status_code)
        except Exception as e:
            logger.warning("ai_router openrouter stream error: %s", e)

    yield "[ERROR: No AI provider available (set OLLAMA_API_KEY or OPENROUTER_API_KEY)]"


# ── HTTP endpoints (kept for compat, now backed by real service) ──


@router.post("/ai/completions")
async def ai_completions(request: dict[str, Any]) -> dict[str, Any]:
    """AI completion via optimal provider routing."""
    return await chat_completion(
        messages=request.get("messages", []),
        model=request.get("model"),
        tier=request.get("tier"),
        temperature=request.get("temperature", 0.3),
        max_tokens=request.get("max_tokens", 500),
    )


@router.post("/ai/chat")
async def ai_chat(request: dict[str, Any]) -> dict[str, Any]:
    """AI chat endpoint with provider fallback."""
    return await chat_completion(
        messages=request.get("messages", []),
        model=request.get("model"),
        tier=request.get("tier"),
        temperature=request.get("temperature", 0.3),
        max_tokens=request.get("max_tokens", 500),
    )


@router.get("/ai/providers")
async def list_providers() -> dict[str, Any]:
    """List available AI providers."""
    return {"providers": list(PROVIDERS.keys())}


@router.get("/ai/models")
async def list_models() -> dict[str, Any]:
    """List available models by tier."""
    return {"tiers": MODEL_TIERS}