import os
import re
import gradio as gr
from huggingface_hub import InferenceClient
from difflib import SequenceMatcher

# =======================
#  MODELLI DISPONIBILI
# =======================

AVAILABLE_MODELS = {
    "Llama 3.2 (1B)": "fanherodev/Llama-3.2-1B-Instruct:featherless-ai",
    "Llama 3.1 (8B)": "meta-llama/Llama-3.1-8B-Instruct:ovhcloud",
    # "Llama 3.3 (70B)": "meta-llama/Llama-3.3-70B-Instruct:ovhcloud",
    "ministral (24B)": "huihui-ai/Mistral-Small-24B-Instruct-2501-abliterated:featherless-ai",
    # "Qwen 2 (2B)": "e-palmisano/Qwen2-1.5B-ITA-Instruct:featherless-ai",
    "Qwen 2 (5B)": "DeepMount00/Qwen2-1.5B-Ita:featherless-ai",
    # "Qwen 2 (72B)": "Qwen/Qwen2.5-VL-72B-Instruct:ovhcloud",
    "Qwen 3 (30B)": "Qwen/Qwen3-Coder-30B-A3B-Instruct:ovhcloud",
    "OpenAI (20B)": "openai/gpt-oss-20b:ovhcloud",
    "Kimi K2 (32B)": "moonshotai/Kimi-K2-Instruct-0905:groq",
    # "Swiss AI Apertus (70B)": "swiss-ai/Apertus-70B-Instruct-2509:publicai",
}

# =======================
#  TOKEN HF
# =======================

HF_TOKEN = os.getenv("HF_TOKEN")
if HF_TOKEN is None:
    raise ValueError(
        "Devi impostare la variabile d'ambiente HF_TOKEN con il tuo token Hugging Face (permessi READ)."
    )

# =======================
#  LETTURA GLOSSARIO TOON
# =======================

def parse_toon_glossario(path: str):
    """
    Legge un file TOON del tipo:
    [N]{term,definition}:
    TERM1,Definizione 1...
    TERM2,Definizione 2...
    e restituisce una lista di dict: {"term": ..., "definition": ...}
    """
    if not os.path.exists(path):
        print(f"⚠️ File glossario TOON non trovato: {path}")
        return []

    with open(path, "r", encoding="utf-8") as f:
        lines = [l.strip() for l in f.readlines() if l.strip()]

    if not lines:
        return []

    # prima riga è l'header tipo [N]{term,definition}:
    header = lines[0]
    rows = lines[1:]

    entries = []
    for row in rows:
        # ignora eventuali righe di commento
        if row.startswith("#"):
            continue
        # spezza solo sulla prima virgola: TERM,DEFINIZIONE CON VIRGOLE
        if "," not in row:
            continue
        term, definition = row.split(",", 1)
        term = term.strip()
        definition = definition.strip()
        if term and definition:
            entries.append({"term": term, "definition": definition})
    
    print(f"📚 Caricate {len(entries)} voci dal glossario TOON.")
    return entries


GLOSSARIO_ENTRIES = parse_toon_glossario("glossario.toon")

# =======================
#  FUZZY MATCHING
# =======================

def similarity(a: str, b: str) -> float:
    """Calcola similarità tra due stringhe (0.0 - 1.0)"""
    return SequenceMatcher(None, a.lower(), b.lower()).ratio()

def trova_voci_rilevanti(question: str, k: int = 3):
    """
    Mini-RAG migliorato con fuzzy matching:
    - Match esatti: score alto
    - Match parziali/fuzzy: score medio
    - Parole nella definizione: score basso
    Restituisce al massimo k voci.
    """
    if not GLOSSARIO_ENTRIES:
        return []

    q = question.lower()
    q_words = set(w.strip(".,;:!?\"'()").lower() for w in q.split() if len(w) > 2)
    scored = []

    for entry in GLOSSARIO_ENTRIES:
        score = 0
        term_lower = entry["term"].lower()
        def_lower = entry["definition"].lower()

        # 1. Match esatto del termine nella domanda (score massimo)
        if term_lower in q:
            score += 10

        # 2. Fuzzy match del termine (per typo o variazioni)
        sim = similarity(term_lower, q)
        if sim > 0.6:  # soglia di similarità
            score += sim * 5

        # 3. Match parziale: parole singole del termine
        term_words = set(term_lower.split())
        overlap = q_words & term_words
        score += len(overlap) * 2

        # 4. Parole della domanda nella definizione
        for word in q_words:
            if word in def_lower:
                score += 0.5

        if score > 0:
            scored.append((score, entry))

    scored.sort(reverse=True, key=lambda x: x[0])
    top_entries = [e for _, e in scored[:k]]
    
    # Debug: mostra i punteggi
    if scored:
        print(f"\n🔍 Query: '{question}'")
        for score, entry in scored[:3]:
            print(f"   {score:.1f} - {entry['term']}")
    
    return top_entries

# =======================
#  PROMPT DI SISTEMA
# =======================

SYSTEM_PROMPT_BASE = (
    "Sei un glossario di informatica pensato per studenti di 15 anni. "
    "Spieghi termini e concetti informatici in modo semplice, con frasi brevi, "
    "poche parole difficili e alcuni esempi pratici. "
    "Non usare internet e non dire mai che stai cercando online: "
    "rispondi solo usando le tue conoscenze interne e il glossario fornito. "
    "Se non sei sicuro, dillo chiaramente. Rispondi sempre in italiano. "
    "Prima pensa in silenzio al significato del termine e confrontalo con le "
    "definizioni del glossario, poi rispondi solo con la spiegazione finale. "
    "Se ti accorgi che qualcosa che stai per dire contraddice il glossario, "
    "correggilo mentalmente prima di scrivere la risposta."
)

def get_system_prompt(mostra_ragionamento: bool):
    """Genera il prompt di sistema in base alla scelta dell'utente"""
    if mostra_ragionamento:
        return (
            SYSTEM_PROMPT_BASE
            + "\n\nPer aiutare lo studente a capire come ragioni, dividi SEMPRE la risposta in due parti:\n\n"
            "### Come ci arrivo\n"
            "- Spiega passo passo il tuo ragionamento in modo semplice, usando eventualmente il glossario.\n\n"
            "### Spiegazione semplice\n"
            "- Una spiegazione finale chiara, breve (massimo 3-4 paragrafi) e riassuntiva.\n"
        )
    else:
        return (
            SYSTEM_PROMPT_BASE
            + "\n\nFornisci una spiegazione chiara e riassuntiva, breve (massimo 3-4 paragrafi), "
            "senza mostrare il processo di ragionamento."
        )

# =======================
#  METRICHE CUSTOM (approx ROUGE / BLEU / chrF)
# =======================

def _tokenize(text: str):
    # tokenizzazione molto semplice: parole alfanumeriche
    return [t for t in re.findall(r"\w+", text.lower()) if t]

def _char_list(text: str):
    return [c for c in text.lower() if not c.isspace()]

def rouge1_f1(pred: str, ref: str) -> float:
    pred_tokens = _tokenize(pred)
    ref_tokens = _tokenize(ref)
    if not pred_tokens or not ref_tokens:
        return 0.0
    pred_set = set(pred_tokens)
    ref_set = set(ref_tokens)
    overlap = len(pred_set & ref_set)
    if overlap == 0:
        return 0.0
    precision = overlap / len(pred_set)
    recall = overlap / len(ref_set)
    if precision + recall == 0:
        return 0.0
    return 2 * precision * recall / (precision + recall)

def bleu1(pred: str, ref: str) -> float:
    pred_tokens = _tokenize(pred)
    ref_tokens = _tokenize(ref)
    if not pred_tokens or not ref_tokens:
        return 0.0
    overlap = 0
    ref_counts = {}
    for t in ref_tokens:
        ref_counts[t] = ref_counts.get(t, 0) + 1
    for t in pred_tokens:
        if ref_counts.get(t, 0) > 0:
            overlap += 1
            ref_counts[t] -= 1
    precision = overlap / len(pred_tokens)
    # brevity penalty
    from math import exp
    len_p = len(pred_tokens)
    len_r = len(ref_tokens)
    if len_p == 0:
        return 0.0
    if len_p > len_r:
        bp = 1.0
    else:
        bp = exp(1 - len_r / len_p)
    return bp * precision

def chrf_simple(pred: str, ref: str) -> float:
    pred_chars = _char_list(pred)
    ref_chars = _char_list(ref)
    if not pred_chars or not ref_chars:
        return 0.0
    pred_set = set(pred_chars)
    ref_set = set(ref_chars)
    overlap = len(pred_set & ref_set)
    if overlap == 0:
        return 0.0
    precision = overlap / len(pred_set)
    recall = overlap / len(ref_set)
    if precision + recall == 0:
        return 0.0
    return 2 * precision * recall / (precision + recall)

def log_quality_metrics(question: str, voci_glossario, answer_text: str, modello_scelto: str):
    """
    Calcola metriche naive (ROUGE1-F1, BLEU1, chrF-like) confrontando la risposta del modello
    con le definizioni del glossario. Solo log su console (admin).
    """
    if not voci_glossario:
        print("ℹ️ Nessuna voce di glossario per questa domanda. Salto il calcolo delle metriche.")
        return

    reference = " ".join(e["definition"] for e in voci_glossario)
    prediction = answer_text.strip()

    r1 = rouge1_f1(prediction, reference)
    b1 = bleu1(prediction, reference)
    cf = chrf_simple(prediction, reference)

    print("\n====== VALUTAZIONE QUALITÀ RISPOSTA (NAIVE METRICS) ======")
    print(f"Modello: {modello_scelto}")
    print(f"Domanda: {question}")
    print(f"Glossario usato per riferimento: {[e['term'] for e in voci_glossario]}")
    print("--- Metriche approx ---")
    print(f"ROUGE-1 F1 (uno-grammi): {r1:.4f}")
    print(f"BLEU-1 (uno-grammi):     {b1:.4f}")
    print(f"chrF semplice (char F1): {cf:.4f}")
    print("=========================================================\n")

# =======================
#  FUNZIONE DI RISPOSTA (STREAMING + SELF-REFLECTION IMPLICITA)
# =======================

def answer_with_self_reflection(question: str, modello_scelto: str, mostra_ragionamento: bool):
    """
    Versione STREAMING con una sola chiamata al modello.
    La 'self-reflection' è incorporata nel prompt:
    - usa il glossario come fonte principale
    - chiede al modello di correggersi mentalmente prima di rispondere

    Alla fine, in background, calcola anche metriche di qualità rispetto al glossario
    e le scrive nei log (solo per amministratore).
    """
    question = question.strip()
    if not question:
        yield "Scrivi un termine o concetto di informatica che vuoi capire meglio 😊"
        return

    model_id = AVAILABLE_MODELS.get(modello_scelto)
    if model_id is None:
        yield "Si è verificato un errore: modello non valido."
        return

    client = InferenceClient(model=model_id, token=HF_TOKEN)

    # 🔍 MINI-RAG: trova voci rilevanti nel glossario TOON
    voci = trova_voci_rilevanti(question, k=3)
    
    # Prompt di sistema
    base_prompt = get_system_prompt(mostra_ragionamento)
    
    if voci:
        contesto = "\n".join(
            f"- **{e['term']}**: {e['definition']}" for e in voci
        )
        n_voci = len(voci)
        system_with_glossary = (
            base_prompt
            + "\n\nUsa il seguente glossario come FONTE PRINCIPALE: "
              "se una definizione è presente, non devi contraddirla. "
              "Puoi aggiungere solo dettagli semplici e coerenti.\n"
              "=== VOCI SELEZIONATE DAL GLOSSARIO ===\n"
            f"{contesto}\n"
              "=== FINE VOCI GLOSSARIO ===\n"
        )
    else:
        n_voci = 0
        system_with_glossary = (
            base_prompt
            + "\n\nNOTA: Non sono state trovate voci specifiche nel glossario per questa domanda. "
              "Se il termine non è nel glossario, dillo chiaramente, e poi spiega con parole semplici."
        )

    messages = [
        {"role": "system", "content": system_with_glossary},
        {
            "role": "user",
            "content": (
                "Spiega questo termine o concetto di informatica a uno studente di 15 anni:\n"
                f"\"{question}\""
            ),
        },
    ]

    if mostra_ragionamento:
        max_tokens = 800
    else:
        max_tokens = 400

    if n_voci > 0:
        prefix = f"🔎 Ho trovato **{n_voci}** voci nel glossario rilevanti per la tua domanda.\n\n"
    else:
        prefix = ""

    partial = prefix
    
    try:
        for event in client.chat_completion(
            messages=messages,
            max_tokens=max_tokens,
            temperature=0.3,
            top_p=0.9,
            stream=True,
            model=model_id,
        ):
            delta = event.choices[0].delta.content or ""
            if not delta:
                continue
            partial += delta
            yield partial
    except Exception as e:
        print("❌ Errore nella generazione:", repr(e))
        yield f"Si è verificato un errore: {e}"
        return

    # Risposta completa (senza prefisso) per le metriche admin
    if n_voci > 0:
        answer_for_metrics = partial[len(prefix):]
    else:
        answer_for_metrics = partial

    log_quality_metrics(question, voci, answer_for_metrics, modello_scelto)

# =======================
#  INTERFACCIA GRADIO
# =======================

demo = gr.Interface(
    fn=answer_with_self_reflection,
    inputs=[
        gr.Textbox(
            label="Termine o domanda",
            placeholder="Es: Che cos'è una variabile?",
            lines=2,
        ),
        gr.Dropdown(
            label="Modello",
            choices=list(AVAILABLE_MODELS.keys()),
            value="Llama 3.2 (1B)",  # valore di default valido
        ),
        gr.Checkbox(
            label="🔍 Mostra il ragionamento (Come ci arrivo)",
            value=False,
            info="Attiva per vedere come il modello ragiona passo-passo"
        ),
    ],
    outputs=gr.Markdown(label="Spiegazione"),
    title="📚 Glossario di Informatica per Novizi",
    description=(
        "Scrivi un termine o concetto di informatica che vuoi capire meglio "
        "(cookie, CPU, variabile, HTTP, algoritmo...). Parti con la domanda: Cosa è un....? "
        "La risposta arriva a pezzetti 🙂\n\n"
        "💡 **Suggerimento:** Attiva 'Mostra il ragionamento' per vedere come gli LLM pensano!\n\n"
        f"📊 **Status glossario:** {len(GLOSSARIO_ENTRIES)} termini caricati"
    ),
    submit_btn="Spiega",
    clear_btn="Pulisci",
    allow_flagging="never",
)

if __name__ == "__main__":
    demo.launch()