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=50)
    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])
    contexto_final = "\n\n".join(result_chunks[:10])
    print("Primeros chunks:")
    print("\n\n".join(result_chunks[:3]))

    prompt = f"""
    Eres un abogado colombiano especializado en el Código de Tránsito, Código de Policía y Código Penal. 
    Debes responder la siguiente pregunta utilizando únicamente el contexto legal que se te proporciona.

    Instrucciones:
    - Fundamenta tu respuesta solamente en el texto legal disponible. No inventes leyes ni artículos.
    - Puedes parafrasear el contenido si ayuda a la comprensión, pero no añadas información externa.
    - Si hay una cita textual o artículo relevante, inclúyelo de forma explícita.
    - Si no encuentras una base suficiente en el texto, responde claramente:
    "No encontré información suficiente en los documentos para responder esta pregunta."
    - Usa un lenguaje claro, sencillo y profesional, dirigido a una persona sin formación jurídica.
    - Finaliza con una breve recomendación si es pertinente.

    CONTEXTO LEGAL:
    {contexto_final}

    PREGUNTA:
    {pregunta}
    
    RESPUESTA:
    """

    resultado = llm(
        prompt,
        max_new_tokens=500,
        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()