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

# Token para modelos privados si se requiere
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'.")

# 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  # DECORADOR para que ZeroGPU ejecute esta función en CPU/GPU remota
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=20)
    result_chunks = [chunks[i] for i in indices[0]]

    palabras_clave = pregunta.lower().split()
    filtrados = [c for c in result_chunks if any(p in c.lower() for p in palabras_clave)]
    contexto_final = "\n\n".join(filtrados[:3]) if filtrados else "\n\n".join(result_chunks[:3])

    prompt = f"""
Actúa como un asesor legal colombiano, especializado en el Código de Tránsito, Código de Policía y Código Penal. 
Analiza el siguiente contexto legal y responde la pregunta de forma clara, completa y profesional.

Instrucciones:
- Usa un lenguaje formal y preciso.
- Fundamenta tu respuesta solo en el contexto proporcionado.
- Cita artículos, normas o sanciones si están presentes en el texto.
- Finaliza con una recomendación preventiva si aplica.
- Si no encuentras información suficiente, responde:
**"No encontré información suficiente en la ley para responder a esta pregunta."**
- No inventes información ni cites leyes que no estén en el contexto.

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()