import os
import streamlit as st
from openai import OpenAI
from VerifierApp.localisation_module import t

_EINFRA_BASE_URL = 'https://llm.ai.e-infra.cz/v1/'
_GEMINI_BASE_URL = 'https://generativelanguage.googleapis.com/v1beta/openai/'
_API_TIMEOUT = 10

# ── Capability registry ───────────────────────────────────────────────────────
# The OpenAI-compatible APIs return NO capability metadata in /v1/models,
# so capabilities are determined by matching substrings of the model ID.
# Update these lists when new models are deployed.
#
#   "vision"  — model accepts image binary content in the prompt
#   "pdf"     — model accepts raw PDF binary (application/pdf) directly;
#               if False, PDFs are pre-extracted to plain text via pypdf
#
_CAPABILITY_PATTERNS: dict[str, list[str]] = {
    "vision": [
        "llama-4",
        "pixtral",
        "gemini",
        "gpt-4o",
        "gpt-4-turbo",
        "claude-3",
        "claude-sonnet",
        "claude-haiku",
        "claude-opus",
        "gemma4",
        "gemma-4",
        "glm-4v",
        "qwen-vl",
        "qwen2-vl",
        "qwen2.5-vl",
        "internvl",
        "phi-3-vision",
        "phi-4-vision",
    ],
    "pdf": [
        "gemini",
        "gpt-4o",
        "gpt-4-turbo",
        "claude-3",
        "claude-sonnet",
        "claude-haiku",
        "claude-opus",
    ],
}


def get_model_capabilities(model_id: str) -> dict[str, bool]:
    """Return capability flags for *model_id*.

    Checks the user's saved overrides first (``data/model_capabilities.json``),
    then falls back to pattern-matching against :data:`_CAPABILITY_PATTERNS`.

    Returns ``{"vision": bool, "pdf": bool}``.
    Safe to call from background threads — no Streamlit dependency.
    """
    from VerifierApp.config import load_model_capability_overrides
    overrides = load_model_capability_overrides()
    name = model_id.lower()

    result: dict[str, bool] = {
        cap: any(pattern in name for pattern in patterns)
        for cap, patterns in _CAPABILITY_PATTERNS.items()
    }

    model_override = overrides.get(model_id, {})
    pdf_override = model_override.get("pdf", "auto")
    if pdf_override == "native":
        result["pdf"] = True
    elif pdf_override == "extract":
        result["pdf"] = False

    return result


@st.cache_data(ttl=7200, show_spinner=False)
def fetch_available_models() -> list[str]:
    """Fetch model IDs from e-INFRA and (optionally) Google Gemini."""
    from dotenv import load_dotenv
    load_dotenv()

    model_ids = []

    try:
        client = OpenAI(
            base_url=_EINFRA_BASE_URL,
            api_key=os.getenv('E_INFRA_API_TOKEN'),
            timeout=_API_TIMEOUT,
        )
        model_ids += [m.id for m in client.models.list().data]
    except Exception as e:
        st.warning(t("models_fetch_error", error=e))

    gemini_token = os.getenv('GEMINI_API_TOKEN')
    if gemini_token:
        try:
            gemini_client = OpenAI(
                base_url=_GEMINI_BASE_URL,
                api_key=gemini_token,
                timeout=_API_TIMEOUT,
            )
            model_ids += [
                m.id.removeprefix("models/") for m in gemini_client.models.list().data
            ]
        except Exception as e:
            st.warning(t("models_gemini_fetch_error", error=e))

    if not model_ids:
        return ["qwen3.5-122b", "llama3.3:latest"]

    return sorted(set(model_ids))


@st.cache_data(ttl=7200, show_spinner=False)
def fetch_model_capabilities() -> dict[str, dict[str, bool]]:
    """Return ``{model_id: {"vision": bool, "pdf": bool}}`` for every known model.

    Includes user overrides.  Cached alongside the model list.
    """
    return {mid: get_model_capabilities(mid) for mid in fetch_available_models()}
