"""ObsiGate AI — Multi-provider AI service for editor enhancement.

Supports: DeepSeek, OpenRouter, Google Gemini, Ollama, NVIDIA, QwenCloud, Xiaomi, Mistral.
Configured via environment variables.
"""

import json
import logging
import os
from pathlib import Path
from typing import Literal

import httpx

logger = logging.getLogger("obsigate.ai")

ProviderName = Literal["deepseek", "openrouter", "gemini", "ollama", "nvidia", "qwencloud", "xiaomi", "mistral"]

# Provider configurations — keys loaded from file or .env
AI_KEYS_FILE = Path("data/api_keys.json")
APP_CONFIG_FILE = Path(__file__).resolve().parent.parent / "data" / "config.json"

def _read_ai_keys() -> dict:
    if not AI_KEYS_FILE.exists():
        return {}
    try:
        return json.loads(AI_KEYS_FILE.read_text(encoding="utf-8"))
    except Exception:
        return {}


def _read_app_config() -> dict:
    """Read the persisted application config (``data/config.json``)."""
    if not APP_CONFIG_FILE.exists():
        return {}
    try:
        return json.loads(APP_CONFIG_FILE.read_text(encoding="utf-8"))
    except Exception:
        return {}

def get_ai_key(env_name: str) -> str:
    """Get AI key: stored file first, then .env fallback."""
    keys = _read_ai_keys()
    if keys.get(env_name):
        return keys[env_name]
    return os.getenv(env_name, "").strip()

def _load_provider_keys():
    """Load AI keys from stored file, falling back to .env."""
    providers = {
        "deepseek": {
            "api_key": get_ai_key("DEEPSEEK_API_KEY"),
            "base_url": "https://api.deepseek.com/v1",
            "model": os.getenv("DEEPSEEK_MODEL", "deepseek-chat"),
            "auth_header": "Bearer {api_key}",
        },
        "openrouter": {
            "api_key": get_ai_key("OPENROUTER_API_KEY"),
            "base_url": "https://openrouter.ai/api/v1",
            "model": os.getenv("OPENROUTER_MODEL", "openai/gpt-4o-mini"),
            "auth_header": "Bearer {api_key}",
        },
        "gemini": {
            "api_key": get_ai_key("GEMINI_API_KEY"),
            "base_url": "https://generativelanguage.googleapis.com/v1beta",
            "model": os.getenv("GEMINI_MODEL", "gemini-2.0-flash"),
            "auth_header": None,
        },
        "ollama": {
            "api_key": "ollama",
            "base_url": os.getenv("OLLAMA_BASE_URL", "http://ollama:11434/v1"),
            "model": os.getenv("OLLAMA_MODEL", "qwen2.5-coder:1.5b"),
            "auth_header": "Bearer {api_key}",
        },
        "nvidia": {
            "api_key": get_ai_key("NVIDIA_API_KEY"),
            "base_url": "https://integrate.api.nvidia.com/v1",
            "model": os.getenv("NVIDIA_MODEL", "meta/llama-3.1-405b-instruct"),
            "auth_header": "Bearer {api_key}",
        },
        "qwencloud": {
            "api_key": get_ai_key("QWENCLOUD_API_KEY"),
            "base_url": "https://dashscope.aliyuncs.com/compatible-mode/v1",
            "model": os.getenv("QWENCLOUD_MODEL", "qwen-max"),
            "auth_header": "Bearer {api_key}",
        },
        "xiaomi": {
            "api_key": get_ai_key("XIAOMI_API_KEY"),
            "base_url": os.getenv("XIAOMI_BASE_URL", "https://api.xiaomimimo.com/v1"),
            "model": os.getenv("XIAOMI_MODEL", "mimo-v2.5-pro"),
            # Xiaomi MiMo uses a dedicated `api-key` header (NOT Authorization: Bearer).
            # The `_call_deepseek_openrouter` helper substitutes {api_key} verbatim,
            # so we just emit the raw key value here.
            "auth_header": "{api_key}",
            "auth_header_name": "api-key",
        },
        "mistral": {
            "api_key": get_ai_key("MISTRAL_API_KEY"),
            "base_url": "https://api.mistral.ai/v1",
            "model": os.getenv("MISTRAL_MODEL", "mistral-large-latest"),
            "auth_header": "Bearer {api_key}",
        },
    }
    # Apply persisted per-provider model overrides (data/config.json).
    overrides = _read_app_config().get("ai_default_models") or {}
    if isinstance(overrides, dict):
        for name, model in overrides.items():
            if name in providers and isinstance(model, str) and model:
                providers[name]["model"] = model
    return providers


PROVIDERS = _load_provider_keys()

DEFAULT_PROVIDER: ProviderName = os.getenv("AI_DEFAULT_PROVIDER", "deepseek")  # type: ignore


def get_default_provider() -> str:
    """Resolve the default provider: ``data/config.json`` > env > ``deepseek``."""
    provider: str = str(_read_app_config().get("ai_default_provider") or os.getenv("AI_DEFAULT_PROVIDER", "deepseek"))
    return provider if provider in PROVIDERS else "deepseek"


def reload_ai_config() -> str:
    """Reload provider keys and model overrides from disk into ``PROVIDERS`` in place.

    Mutating ``PROVIDERS`` (rather than rebinding it) keeps references held by
    other modules valid. Returns the resolved default provider.
    """
    global DEFAULT_PROVIDER
    PROVIDERS.clear()
    PROVIDERS.update(_load_provider_keys())
    DEFAULT_PROVIDER = get_default_provider()  # type: ignore[assignment]
    logger.info(f"AI config reloaded (default provider: {DEFAULT_PROVIDER})")
    return DEFAULT_PROVIDER


def _get_provider_config(provider: ProviderName | None = None) -> dict:
    """Get provider config, falling back to default if requested provider unavailable."""
    default = get_default_provider()
    p = provider or default
    if p not in PROVIDERS:
        p = default
    cfg = PROVIDERS[p]
    if not cfg["api_key"]:
        # Try next available provider
        for alt in PROVIDERS:
            if PROVIDERS[alt]["api_key"]:
                p = alt
                cfg = PROVIDERS[alt]
                break
    return {"name": p, **cfg}


def _build_headers(cfg: dict) -> dict:
    """Build HTTP headers for an OpenAI-compatible provider config.

    Most providers use ``Authorization: Bearer KEY``. Some (Xiaomi MiMo) use a
    dedicated header like ``api-key: KEY`` — supported via the
    ``auth_header_name`` key in PROVIDERS (defaults to ``Authorization``).
    """
    header_name = cfg.get("auth_header_name") or "Authorization"
    header_value = cfg["auth_header"].format(api_key=cfg["api_key"])
    if header_name == "Authorization" and not header_value.lower().startswith("bearer "):
        header_value = "Bearer " + header_value
    return {
        header_name: header_value,
        "Content-Type": "application/json",
    }


async def _call_deepseek_openrouter(prompt: str, system: str, provider: ProviderName | None = None,
                                     temperature: float = 0.7, max_tokens: int = 2048) -> str:
    """Call OpenAI-compatible API (DeepSeek, OpenRouter, Xiaomi MiMo, etc.)."""
    cfg = _get_provider_config(provider)
    # Debug: log masked key to diagnose 401
    key_preview = cfg["api_key"][:8] + "..." + cfg["api_key"][-4:] if len(cfg["api_key"]) > 12 else "***"
    logger.info(f"AI call: provider={cfg['name']} model={cfg['model']} key={key_preview}")
    headers = _build_headers(cfg)
    payload = {
        "model": cfg["model"],
        "messages": [
            {"role": "system", "content": system},
            {"role": "user", "content": prompt},
        ],
        "temperature": temperature,
        "max_tokens": max_tokens,
    }
    async with httpx.AsyncClient(timeout=60.0) as client:
        resp = await client.post(
            f"{cfg['base_url']}/chat/completions",
            headers=headers,
            json=payload,
        )
        resp.raise_for_status()
        data = resp.json()
        return data["choices"][0]["message"]["content"].strip()


async def _call_gemini(prompt: str, system: str, temperature: float = 0.7, max_tokens: int = 2048) -> str:
    """Call Google Gemini API."""
    cfg = PROVIDERS["gemini"]
    url = f"{cfg['base_url']}/models/{cfg['model']}:generateContent?key={cfg['api_key']}"
    payload = {
        "system_instruction": {"parts": [{"text": system}]},
        "contents": [{"parts": [{"text": prompt}]}],
        "generationConfig": {
            "temperature": temperature,
            "maxOutputTokens": max_tokens,
        },
    }
    async with httpx.AsyncClient(timeout=60.0) as client:
        resp = await client.post(url, json=payload)
        resp.raise_for_status()
        data = resp.json()
        return data["candidates"][0]["content"]["parts"][0]["text"].strip()


async def ai_complete(prompt: str, provider: ProviderName | None = None) -> str:
    """Generic AI completion. Routes to appropriate provider."""
    cfg = _get_provider_config(provider)
    if cfg["name"] == "gemini":
        return await _call_gemini(prompt, "You are a helpful assistant.")
    # All other providers use OpenAI-compatible format
    return await _call_deepseek_openrouter(prompt, "You are a helpful assistant.", provider)


# ── Specialized AI actions ──

SYSTEM_PROMPT = """You are an AI assistant integrated into ObsiGate, a knowledge management tool. 
Your responses should be direct and concise. When editing text, return ONLY the modified text, 
no explanations or markdown fences."""


async def ai_improve_writing(text: str, provider: ProviderName | None = None) -> str:
    """Improve writing quality while preserving meaning."""
    return await _call_deepseek_openrouter(
        f"Improve the following text. Fix grammar, clarity, and flow. Preserve the original language and meaning.\n\n{text}",
        SYSTEM_PROMPT, provider, temperature=0.3,
    )


async def ai_fix_spelling(text: str, provider: ProviderName | None = None) -> str:
    """Fix spelling and grammar errors."""
    return await _call_deepseek_openrouter(
        f"Fix all spelling and grammar errors in this text. Return only the corrected text.\n\n{text}",
        SYSTEM_PROMPT, provider, temperature=0.1,
    )


async def ai_make_shorter(text: str, provider: ProviderName | None = None) -> str:
    """Make text more concise."""
    return await _call_deepseek_openrouter(
        f"Make this text shorter and more concise while preserving the key information.\n\n{text}",
        SYSTEM_PROMPT, provider, temperature=0.3,
    )


async def ai_make_longer(text: str, provider: ProviderName | None = None) -> str:
    """Expand text with more detail."""
    return await _call_deepseek_openrouter(
        f"Expand this text with more detail, examples, or explanation while keeping the same tone.\n\n{text}",
        SYSTEM_PROMPT, provider, temperature=0.7, max_tokens=4096,
    )


async def ai_simplify(text: str, provider: ProviderName | None = None) -> str:
    """Simplify language."""
    return await _call_deepseek_openrouter(
        f"Simplify this text. Use clearer, more straightforward language. Avoid jargon.\n\n{text}",
        SYSTEM_PROMPT, provider, temperature=0.3,
    )


async def ai_change_tone(text: str, tone: str, provider: ProviderName | None = None) -> str:
    """Change the tone of the text."""
    return await _call_deepseek_openrouter(
        f"Rewrite this text in a {tone} tone. Preserve the original meaning.\n\n{text}",
        SYSTEM_PROMPT, provider, temperature=0.5,
    )


async def ai_translate(text: str, target_lang: str, provider: ProviderName | None = None) -> str:
    """Translate text to target language."""
    # Gemini is better at translation
    if provider is None and PROVIDERS["gemini"]["api_key"]:
        return await _call_gemini(
            f"Translate the following text to {target_lang}. Return only the translation.\n\n{text}",
            "You are a professional translator. Translate accurately and naturally.",
            temperature=0.1,
        )
    return await _call_deepseek_openrouter(
        f"Translate the following text to {target_lang}. Return only the translation.\n\n{text}",
        "You are a professional translator. Translate accurately and naturally.",
        provider, temperature=0.1,
    )


async def ai_explain(text: str, provider: ProviderName | None = None) -> str:
    """Explain the selected text."""
    return await _call_deepseek_openrouter(
        f"Explain the following text clearly and concisely:\n\n{text}",
        SYSTEM_PROMPT, provider, temperature=0.3,
    )


async def ai_summarize(text: str, provider: ProviderName | None = None) -> str:
    """Summarize the selected text."""
    return await _call_deepseek_openrouter(
        f"Summarize the following text concisely:\n\n{text}",
        SYSTEM_PROMPT, provider, temperature=0.3,
    )


async def ai_continue_writing(text: str, provider: ProviderName | None = None) -> str:
    """Continue writing from the selected text."""
    return await _call_deepseek_openrouter(
        f"Continue writing from where this text leaves off. Match the style and tone:\n\n{text}",
        SYSTEM_PROMPT, provider, temperature=0.7, max_tokens=4096,
    )


async def ai_custom_rewrite(text: str, instruction: str, provider: ProviderName | None = None) -> str:
    """Rewrite text based on a custom instruction."""
    return await _call_deepseek_openrouter(
        f"Rewrite the following text according to this instruction: {instruction}\n\n{text}",
        SYSTEM_PROMPT, provider, temperature=0.5,
    )


async def ai_convert_to_list(text: str, provider: ProviderName | None = None) -> str:
    """Convert paragraph text to a markdown list."""
    return await _call_deepseek_openrouter(
        f"Convert this text into a well-organized markdown bullet list. Extract key points.\n\n{text}",
        SYSTEM_PROMPT, provider, temperature=0.2,
    )


async def ai_convert_to_table(text: str, provider: ProviderName | None = None) -> str:
    """Convert text to a markdown table."""
    return await _call_deepseek_openrouter(
        f"Convert this information into a markdown table. Choose appropriate columns.\n\n{text}",
        SYSTEM_PROMPT, provider, temperature=0.2,
    )


async def ai_generate_frontmatter(text: str, provider: ProviderName | None = None) -> str:
    """Generate YAML frontmatter for a markdown document."""
    return await _call_deepseek_openrouter(
        f"Generate YAML frontmatter for this markdown document. Include: titre, tags (as list), catégorie, statut, date. Return ONLY the YAML between --- markers.\n\n{text}",
        SYSTEM_PROMPT, provider, temperature=0.3,
    )


async def ai_inline_complete(text: str, provider: ProviderName | None = None) -> str:
    """Inline completion — suggest a short continuation of the text before the cursor."""
    return await _call_deepseek_openrouter(
        "Continue the text below in the same language. Reply with ONLY the "
        "continuation: no repetition, no quotes, no explanation, at most one "
        "short sentence. If the text ends with a partial word, finish that word.\n\n"
        + text,
        SYSTEM_PROMPT, provider, temperature=0.2, max_tokens=128,
    )


async def ai_convert_to_canvas(text: str, provider: ProviderName | None = None) -> str:
    """Convert text to a Mermaid diagram or canvas representation."""
    return await _call_deepseek_openrouter(
        f"Convert this content into a Mermaid.js diagram if applicable, or a structured outline. Choose the best format.\n\n{text}",
        SYSTEM_PROMPT, provider, temperature=0.3, max_tokens=4096,
    )