from __future__ import annotations

import subprocess
from typing import List, Optional

import numpy as np
import requests


DEFAULT_OPENAI_BASE = "https://api.openai.com"


def _norm_provider(provider: str | None) -> str:
    p = (provider or "ollama").strip().lower()
    return p or "ollama"


def _norm_base(base: str | None, default: str) -> str:
    b = (base or "").strip()
    return (b or default).rstrip("/")


def ollama_pull_model(model: str) -> None:
    """Ensure the given model exists locally by running `ollama pull <model>`."""
    try:
        print(f"Ollama: prøver at hente modellen '{model}' (ollama pull)...")
        subprocess.run(["ollama", "pull", model], check=True)
    except FileNotFoundError:
        raise RuntimeError(
            "Kunne ikke finde 'ollama' kommandoen. Installer Ollama og sørg for at 'ollama' ligger i PATH."
        )
    except subprocess.CalledProcessError as e:
        raise RuntimeError(f"Kunne ikke pull'e modellen '{model}'. Ollama exit code: {e.returncode}")


def embed_batch(
    texts: List[str],
    *,
    provider: str,
    model: str,
    host: str = "http://localhost:11434",
    api_key: Optional[str] = None,
    api_base: Optional[str] = None,
    timeout: int = 120,
) -> np.ndarray:
    p = _norm_provider(provider)
    if p == "ollama":
        return _ollama_embed_batch(texts, model=model, host=host, timeout=timeout)
    if p == "openai":
        return _openai_embed_batch(texts, model=model, api_key=api_key, api_base=api_base, timeout=timeout)
    raise RuntimeError(f"Ukendt provider: {provider}")


def embed_one(
    text: str,
    *,
    provider: str,
    model: str,
    host: str = "http://localhost:11434",
    api_key: Optional[str] = None,
    api_base: Optional[str] = None,
    timeout: int = 120,
) -> np.ndarray:
    vecs = embed_batch(
        [text],
        provider=provider,
        model=model,
        host=host,
        api_key=api_key,
        api_base=api_base,
        timeout=timeout,
    )
    return np.array(vecs[0], dtype=np.float32)


def chat_completion(
    prompt: str,
    *,
    provider: str,
    model: str,
    host: str = "http://localhost:11434",
    api_key: Optional[str] = None,
    api_base: Optional[str] = None,
    timeout: int = 180,
) -> str:
    p = _norm_provider(provider)
    if p == "ollama":
        return _ollama_chat(prompt, model=model, host=host, timeout=timeout)
    if p == "openai":
        return _openai_chat(prompt, model=model, api_key=api_key, api_base=api_base, timeout=timeout)
    raise RuntimeError(f"Ukendt provider: {provider}")


def _openai_headers(api_key: Optional[str]) -> dict:
    key = (api_key or "").strip()
    if not key:
        raise RuntimeError("OpenAI API key mangler. Indtast din egen nøgle i appen.")
    return {
        "Authorization": f"Bearer {key}",
        "Content-Type": "application/json",
    }


def _openai_embed_batch(
    texts: List[str],
    *,
    model: str,
    api_key: Optional[str],
    api_base: Optional[str],
    timeout: int,
) -> np.ndarray:
    if not texts:
        return np.zeros((0, 0), dtype=np.float32)
    base = _norm_base(api_base, DEFAULT_OPENAI_BASE)
    url = f"{base}/v1/embeddings"
    payload = {"model": model, "input": texts}
    r = requests.post(url, headers=_openai_headers(api_key), json=payload, timeout=timeout)
    if not r.ok:
        raise RuntimeError(f"OpenAI embeddings fejlede: {r.status_code} {r.text[:300]}")
    data = r.json()
    items = data.get("data")
    if not isinstance(items, list):
        raise RuntimeError(f"Uventet OpenAI embeddings-svar: {data}")
    items_sorted = sorted(items, key=lambda x: int(x.get("index", 0)) if isinstance(x, dict) else 0)
    vecs: List[List[float]] = []
    for it in items_sorted:
        if not isinstance(it, dict):
            continue
        emb = it.get("embedding")
        if not isinstance(emb, list):
            continue
        vecs.append(emb)
    if not vecs:
        raise RuntimeError(f"Ingen embeddings i OpenAI-svar: {data}")
    return np.array(vecs, dtype=np.float32)


def _openai_chat(
    prompt: str,
    *,
    model: str,
    api_key: Optional[str],
    api_base: Optional[str],
    timeout: int,
) -> str:
    base = _norm_base(api_base, DEFAULT_OPENAI_BASE)
    url = f"{base}/v1/chat/completions"
    payload = {
        "model": model,
        "messages": [
            {
                "role": "system",
                "content": (
                    "Du er en hjælpsom dansk studiemakker, men du MÅ KUN bruge den givne KONTEKST. "
                    "Hvis svaret ikke kan udledes af KONTEKSTEN, skal du skrive: "
                    "'Jeg kan ikke finde det i dine noter (i de hentede kilder).' "
                    "Du må ikke bruge baggrundsviden. Når du påstår noget faktuelt, skal du altid henvise til "
                    "mindst én KILDE som [KILDE 1], [KILDE 2] osv. Hvis ingen KILDE passer, sig at det ikke står i noterne."
                ),
            },
            {"role": "user", "content": prompt},
        ],
        "temperature": 0.1,
    }
    r = requests.post(url, headers=_openai_headers(api_key), json=payload, timeout=timeout)
    if not r.ok:
        raise RuntimeError(f"OpenAI chat fejlede: {r.status_code} {r.text[:300]}")
    data = r.json()
    choices = data.get("choices")
    if not isinstance(choices, list) or not choices:
        raise RuntimeError(f"Uventet OpenAI chat-svar: {data}")
    msg = (choices[0].get("message") or {}) if isinstance(choices[0], dict) else {}
    content = msg.get("content")
    if isinstance(content, str):
        return content
    raise RuntimeError(f"Mangler content i OpenAI chat-svar: {data}")


def _ollama_embed_batch(texts: List[str], *, model: str, host: str, timeout: int) -> np.ndarray:
    """
    Returns embeddings as float32 array shape (n, d).
    Supports both /api/embed and /api/embeddings styles.
    """
    url_embed = f"{host.rstrip('/')}/api/embed"
    url_embeddings = f"{host.rstrip('/')}/api/embeddings"

    r = requests.post(url_embed, json={"model": model, "input": texts}, timeout=timeout)
    if r.ok:
        data = r.json()
        vecs = data.get("embeddings") or data.get("embedding")
        if vecs is None:
            raise RuntimeError(f"Uventet svar fra {url_embed}: {data}")
        return np.array(vecs, dtype=np.float32)

    if r.status_code == 404 and "model" in r.text and "not found" in r.text:
        ollama_pull_model(model)
        r_retry = requests.post(url_embed, json={"model": model, "input": texts}, timeout=timeout)
        if r_retry.ok:
            data = r_retry.json()
            vecs = data.get("embeddings") or data.get("embedding")
            if vecs is None:
                raise RuntimeError(f"Uventet svar fra {url_embed}: {data}")
            return np.array(vecs, dtype=np.float32)

    vec_list: List[List[float]] = []
    for t in texts:
        r2 = requests.post(url_embeddings, json={"model": model, "prompt": t}, timeout=timeout)
        if not r2.ok:
            if r2.status_code == 404 and "model" in r2.text and "not found" in r2.text:
                ollama_pull_model(model)
                r2_retry = requests.post(url_embeddings, json={"model": model, "prompt": t}, timeout=timeout)
                if not r2_retry.ok:
                    raise RuntimeError(
                        f"Ollama embeddings fejlede efter pull: {r2_retry.status_code} {r2_retry.text[:200]}"
                    )
                data2 = r2_retry.json()
                v = data2.get("embedding")
                if v is None:
                    raise RuntimeError(f"Uventet svar fra {url_embeddings}: {data2}")
                vec_list.append(v)
                continue

            raise RuntimeError(f"Ollama embeddings fejlede: {r2.status_code} {r2.text[:200]}")
        data2 = r2.json()
        v = data2.get("embedding")
        if v is None:
            raise RuntimeError(f"Uventet svar fra {url_embeddings}: {data2}")
        vec_list.append(v)
    return np.array(vec_list, dtype=np.float32)


def _ollama_chat(prompt: str, *, model: str, host: str, timeout: int) -> str:
    url = f"{host.rstrip('/')}/api/chat"
    payload = {
        "model": model,
        "messages": [
            {
                "role": "system",
                "content": (
                    "Du er en hjælpsom dansk studiemakker, men du MÅ KUN bruge den givne KONTEKST. "
                    "Hvis svaret ikke kan udledes af KONTEKSTEN, skal du skrive: "
                    "'Jeg kan ikke finde det i dine noter (i de hentede kilder).' "
                    "Du må ikke bruge baggrundsviden. Når du påstår noget faktuelt, skal du altid henvise til mindst én "
                    "KILDE som [KILDE 1], [KILDE 2] osv. Hvis ingen KILDE passer, sig at det ikke står i noterne."
                ),
            },
            {"role": "user", "content": prompt},
        ],
        "stream": False,
        "options": {"temperature": 0.1},
    }
    r = requests.post(url, json=payload, timeout=timeout)
    if r.status_code == 404 and "model" in r.text and "not found" in r.text:
        fallbacks = ["mistral", "mistral-small3.1"]
        for fb in fallbacks:
            if fb == model:
                continue
            payload_fb = dict(payload)
            payload_fb["model"] = fb
            r_fb = requests.post(url, json=payload_fb, timeout=timeout)
            if r_fb.ok:
                data = r_fb.json()
                return (data.get("message") or {}).get("content") or ""
        ollama_pull_model(model)
        r_retry = requests.post(url, json=payload, timeout=timeout)
        if r_retry.ok:
            data = r_retry.json()
            return (data.get("message") or {}).get("content") or ""
        raise RuntimeError(f"Ollama chat fejlede efter pull: {r_retry.status_code} {r_retry.text[:200]}")
    if not r.ok:
        raise RuntimeError(f"Ollama chat fejlede: {r.status_code} {r.text[:200]}")
    data = r.json()
    return (data.get("message") or {}).get("content") or ""
