# ============================================================
#  llm.py — Connexion Ollama + Groq
#  rolax1212 v2.0.2 by Roland
# ============================================================

import json
import os
import requests
from pathlib import Path
from .ui import console, ok, info, warn, err

OLLAMA_URL  = "http://localhost:11434/api/chat"
OLLAMA_TAGS = "http://localhost:11434/api/tags"
GROQ_URL    = "https://api.groq.com/openai/v1/chat/completions"

GROQ_MODELES = [
    {"id": "llama-3.3-70b-versatile",    "nom": "Llama 3.3 70B",    "tag": "★ puissant  (recommandé)"},
    {"id": "llama-3.1-8b-instant",       "nom": "Llama 3.1 8B",     "tag": "⚡ rapide   (très rapide)"},
    {"id": "mistral-saba-24b",           "nom": "Mistral Saba 24B", "tag": "★ puissant"},
    {"id": "gemma2-9b-it",               "nom": "Gemma2 9B",        "tag": "⚡ rapide"},
]

MODELES_PUISSANTS = ["gpt-oss", "llama3", "mistral", "mixtral", "70b", "120b"]
MODELES_LEGERS    = ["1.5b", "0.5b", "1b", "3b"]

# ── Chargement clé Groq ──────────────────────────────────────
def _charger_cle_groq():
    # 1. Variable d'environnement
    cle = os.environ.get("GROQ_API_KEY", "")
    if cle:
        return cle
    # 2. Fichier .env dans le dossier courant
    env = Path(".env")
    if env.exists():
        for ligne in env.read_text(encoding="utf-8").splitlines():
            if ligne.startswith("GROQ_API_KEY="):
                return ligne.split("=", 1)[1].strip().strip('"').strip("'")
    # 3. Fichier .env dans le home
    env_home = Path.home() / ".rolax1212.env"
    if env_home.exists():
        for ligne in env_home.read_text(encoding="utf-8").splitlines():
            if ligne.startswith("GROQ_API_KEY="):
                return ligne.split("=", 1)[1].strip().strip('"').strip("'")
    return ""


def _sauvegarder_cle_groq(cle):
    env_home = Path.home() / ".rolax1212.env"
    env_home.write_text(f"GROQ_API_KEY={cle}\n", encoding="utf-8")
    ok(f"Clé sauvegardée dans {env_home}")


# ── Ollama ───────────────────────────────────────────────────
def ollama_actif():
    try:
        requests.get(OLLAMA_TAGS, timeout=3)
        return True
    except Exception:
        return False


def lister_modeles_ollama():
    try:
        r = requests.get(OLLAMA_TAGS, timeout=5)
        return [m["name"] for m in r.json().get("models", [])]
    except Exception:
        return []


# ── Sélection provider ───────────────────────────────────────
def choisir_modele():
    console.print("\n[bold cyan]Provider :[/bold cyan]")
    console.print("  [bold]1.[/bold] 🌐 [cyan]Groq[/cyan]        [dim](cloud gratuit — ultra rapide)[/dim]")
    console.print("  [bold]2.[/bold] 🖥  [yellow]Ollama[/yellow]      [dim](local — nécessite bonne RAM)[/dim]")
    console.print()

    choix = console.input("[cyan]Choisis (1/2) [1]: [/cyan]").strip() or "1"

    if choix == "2":
        return _choisir_ollama()
    else:
        return _choisir_groq()


def _choisir_groq():
    cle = _charger_cle_groq()

    if not cle:
        warn("Clé Groq non trouvée.")
        console.print("[dim]Obtiens-la sur : https://console.groq.com/keys[/dim]")
        cle = console.input("[cyan]Colle ta clé Groq (gsk_...) : [/cyan]").strip()
        if not cle:
            err("Clé invalide.")
            return _choisir_ollama()
        _sauvegarder_cle_groq(cle)

    ok("Groq connecté")
    console.print("\n[bold cyan]Modèles Groq :[/bold cyan]")
    for i, m in enumerate(GROQ_MODELES, 1):
        console.print(f"  [bold]{i}.[/bold] {m['nom']}  [dim]{m['tag']}[/dim]")

    console.print()
    choix = console.input("[cyan]Choisis un modèle [1]: [/cyan]").strip() or "1"
    try:
        idx = int(choix) - 1
        modele = GROQ_MODELES[idx if 0 <= idx < len(GROQ_MODELES) else 0]
    except ValueError:
        modele = GROQ_MODELES[0]

    ok(f"Modèle : [bold]{modele['nom']}[/bold]")
    return f"groq:{modele['id']}:{cle}"


def _choisir_ollama():
    info("Connexion à Ollama...")
    if not ollama_actif():
        err("Ollama ne répond pas. Lance : ollama serve")
        raise SystemExit(1)
    ok("Ollama actif")
    modeles = lister_modeles_ollama()
    if not modeles:
        err("Aucun modèle Ollama installé.")
        raise SystemExit(1)

    console.print("\n[bold cyan]Modèles Ollama :[/bold cyan]")
    for i, m in enumerate(modeles, 1):
        if any(p in m for p in MODELES_PUISSANTS):
            tag = "[yellow]★ puissant[/yellow]"
        elif any(p in m for p in MODELES_LEGERS):
            tag = "[green]⚡ rapide[/green]"
        else:
            tag = "[dim]standard[/dim]"
        console.print(f"  [bold]{i}.[/bold] {m}  {tag}")

    console.print()
    choix = console.input("[cyan]Choisis un modèle [1]: [/cyan]").strip() or "1"
    if choix.isdigit():
        idx = int(choix) - 1
        modele = modeles[idx] if 0 <= idx < len(modeles) else modeles[0]
    else:
        modele = next((m for m in modeles if choix.lower() in m.lower()), modeles[0])

    ok(f"Modèle : [bold]{modele}[/bold]")
    return modele


# ── Envoi message ────────────────────────────────────────────
def demander(messages, modele, stream=True, timeout=300):
    if modele.startswith("groq:"):
        _, model_id, cle = modele.split(":", 2)
        return _groq(messages, model_id, cle, stream, timeout)
    else:
        return _ollama(messages, modele, stream, timeout)


def demander_conversation(messages, modele):
    return demander(messages, modele, stream=True, timeout=300)


# ── Groq ─────────────────────────────────────────────────────
def _groq(messages, model_id, cle, stream=True, timeout=60):
    headers = {
        "Authorization": f"Bearer {cle}",
        "Content-Type":  "application/json"
    }
    payload = {
        "model":       model_id,
        "messages":    messages,
        "stream":      stream,
        "temperature": 0.2,
        "max_tokens":  4096,
    }
    try:
        if stream:
            return _groq_stream(headers, payload, timeout)
        else:
            r = requests.post(GROQ_URL, headers=headers, json=payload, timeout=timeout)
            r.raise_for_status()
            return r.json()["choices"][0]["message"]["content"]
    except requests.exceptions.HTTPError as e:
        if e.response.status_code == 401:
            err("Clé Groq invalide — va sur https://console.groq.com/keys")
        else:
            err(f"Groq erreur {e.response.status_code}")
        return None
    except Exception as e:
        err(f"Groq erreur : {e}")
        return None


def _groq_stream(headers, payload, timeout=60):
    reponse        = ""
    affiche_entete = False
    est_dans_code  = False

    try:
        with requests.post(GROQ_URL, headers=headers, json=payload, stream=True, timeout=timeout) as r:
            r.raise_for_status()
            for ligne in r.iter_lines():
                if not ligne:
                    continue
                ligne_str = ligne.decode("utf-8") if isinstance(ligne, bytes) else ligne
                if ligne_str.startswith("data: "):
                    data_str = ligne_str[6:]
                    if data_str == "[DONE]":
                        break
                    try:
                        data  = json.loads(data_str)
                        token = data["choices"][0].get("delta", {}).get("content", "")
                        if not token:
                            continue
                        reponse += token

                        if "```" in token:
                            est_dans_code = not est_dans_code

                        if not est_dans_code and "```" not in token:
                            if not affiche_entete:
                                console.print("\n[bold cyan]Jarvis >[/bold cyan] ", end="")
                                affiche_entete = True
                            print(token, end="", flush=True)
                    except json.JSONDecodeError:
                        continue
        if affiche_entete:
            print()
    except Exception as e:
        print()
        err(f"Stream Groq interrompu : {e}")

    return reponse or None


# ── Ollama ───────────────────────────────────────────────────
def _ollama(messages, modele, stream=True, timeout=300):
    payload = {
        "model":    modele,
        "messages": messages,
        "stream":   stream,
        "options":  {"temperature": 0.2, "num_predict": 4096, "num_ctx": 8192}
    }
    try:
        if stream:
            return _ollama_stream(payload, timeout)
        else:
            r = requests.post(OLLAMA_URL, json=payload, timeout=timeout)
            return r.json().get("message", {}).get("content", "")
    except requests.exceptions.ConnectionError:
        err("Ollama ne répond pas.")
        return None
    except requests.exceptions.Timeout:
        err(f"Timeout {timeout}s.")
        return None
    except Exception as e:
        err(f"Erreur LLM : {e}")
        return None


def _ollama_stream(payload, timeout=300):
    reponse        = ""
    affiche_entete = False
    est_dans_code  = False

    try:
        with requests.post(OLLAMA_URL, json=payload, stream=True, timeout=timeout) as r:
            for ligne in r.iter_lines():
                if not ligne:
                    continue
                try:
                    data  = json.loads(ligne)
                    token = data.get("message", {}).get("content", "")
                    reponse += token

                    if "```" in token:
                        est_dans_code = not est_dans_code

                    if not est_dans_code and "```" not in token:
                        if not affiche_entete:
                            console.print("\n[bold cyan]Jarvis >[/bold cyan] ", end="")
                            affiche_entete = True
                        print(token, end="", flush=True)

                    if data.get("done"):
                        break
                except json.JSONDecodeError:
                    continue
        if affiche_entete:
            print()
    except Exception as e:
        print()
        err(f"Stream Ollama interrompu : {e}")

    return reponse or None
