import gradio as gr
import pickle
import os
import json
from transformers import AutoTokenizer, AutoModel, pipeline
import torch
import faiss
import numpy as np
from spaces import GPU  # IMPORTANTE para ZeroGPU

hf_token = os.getenv("HF_KEY")

# Cargar índice FAISS y los chunks
if os.path.exists("index.pkl"):
    with open("index.pkl", "rb") as f:
        index, chunks = pickle.load(f)
else:
    raise FileNotFoundError("No se encontró el archivo 'index.pkl'.")

# Cargar diccionario de sinónimos
with open("sinonimos.json", "r", encoding="utf-8") as f:
    diccionario_sinonimos = json.load(f)

# Función para expandir palabras clave con sinónimos
def expandir_con_sinonimos(palabras, diccionario):
    resultado = set(palabras)
    for palabra in palabras:
        for clave, sinonimos in diccionario.items():
            if palabra == clave or palabra in sinonimos:
                resultado.update([clave] + sinonimos)
    return list(resultado)

# Modelo de embeddings
model_id = "jinaai/jina-embeddings-v2-base-es"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModel.from_pretrained(model_id)

def generar_embedding(texto):
    inputs = tokenizer(texto, return_tensors="pt", padding=True, truncation=True)
    with torch.no_grad():
        outputs = model(**inputs)
        last_hidden = outputs.last_hidden_state
        mask = inputs["attention_mask"].unsqueeze(-1).expand(last_hidden.size()).float()
        summed = torch.sum(last_hidden * mask, 1)
        counted = torch.clamp(mask.sum(1), min=1e-9)
        mean_pooled = summed / counted
        return mean_pooled.numpy()

# LLM para generar respuesta final
llm = pipeline(
    "text-generation",
    model="meta-llama/Llama-3.2-3B-Instruct",
    token=hf_token,
    trust_remote_code=True
)

@GPU
def responder(pregunta):
    if not pregunta:
        return "Por favor ingresa una pregunta."

    pregunta_embedding = generar_embedding(pregunta)
    distances, indices = index.search(pregunta_embedding.reshape(1, -1), k=10)  
    result_chunks = [chunks[i] for i in indices[0]]

    palabras_clave = pregunta.lower().split()
    palabras_expandidas = expandir_con_sinonimos(palabras_clave, diccionario_sinonimos)

        def contar_coincidencias(chunk, palabras):
        return sum(1 for p in palabras if p in chunk.lower())

    filtrados = sorted(
        result_chunks,
        key=lambda c: contar_coincidencias(c, palabras_expandidas),
        reverse=True
    )

    contexto_final = "\n\n".join(filtrados[:3]) if filtrados else "\n\n".join(result_chunks[:3])

    prompt = f"""
Eres un asistente legal especializado en la legislación colombiana, con conocimientos en el Código de Tránsito, el Código Nacional de Policía y el Código Penal. Tu tarea es responder de forma clara, breve y profesional la siguiente pregunta, utilizando únicamente el contexto legal que se te proporciona. Puedes interpretar el significado de los artículos si su relación con la pregunta es evidente, pero no debes inventar leyes, sanciones ni artículos que no estén presentes en el texto.

No utilices conocimiento externo, no cites fuentes adicionales, ni incluyas enlaces o advertencias innecesarias. Si el contexto no ofrece información suficiente para responder, dilo de forma directa y concluye. Siempre que sea posible, sintetiza tu respuesta para que sea comprensible incluso para personas sin formación jurídica.

CONTEXTO LEGAL:
{contexto_final}

PREGUNTA:
{pregunta}

RESPUESTA:
"""

    resultado = llm(
        prompt,
        max_new_tokens=350,
        temperature=0.4,
        top_p=0.9,
        repetition_penalty=1.2
    )[0]["generated_text"]

    if "RESPUESTA:" in resultado:
        solo_respuesta = resultado.split("RESPUESTA:")[-1].strip()
    else:
        solo_respuesta = resultado.strip()

    return solo_respuesta

# Interfaz de usuario con Gradio
demo = gr.Interface(
    fn=responder,
    inputs=gr.Textbox(label="Escribe tu pregunta"),
    outputs=gr.Textbox(label="Respuesta generada"),
    title="Asistente Legal Colombiano",
    description="Consulta el Código de Tránsito, Código de Policía y Código Penal colombiano."
)

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