import gradio as gr
import os
import pickle
import numpy as np
import faiss
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import torch
from sentence_transformers import SentenceTransformer

# Carga del índice FAISS y los chunks de texto
with open("index.pkl", "rb") as f:
    index, chunks = pickle.load(f)

# Modelo de embeddings (debe coincidir con el usado en indexador.py)
embedder = SentenceTransformer("jinaai/jina-embeddings-v2-base-es")

# Modelo LLM (LLaMA o similar)
llm = pipeline(
    "text-generation",
    model="meta-llama/Llama-3.2-3B-Instruct",
    device_map="auto",
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    token=os.getenv("HF_KEY"),
    trust_remote_code=True
)

def responder(pregunta):
    pregunta_embedding = embedder.encode([pregunta])
    _, indices = index.search(np.array(pregunta_embedding).reshape(1, -1), k=50)
    result_chunks = [chunks[i] for i in indices[0]]

    contexto_final = "\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.
Tu tarea es analizar el siguiente contexto legal y responder la pregunta de forma clara, completa y profesional.

Instrucciones:
- Usa un lenguaje formal y comprensible.
- Cita artículos o sanciones si están presentes en el contexto.
- No inventes leyes que no están en el texto.
- Si no hay suficiente contexto, responde:
"No encontré información suficiente en la ley para responder a esta pregunta."

Contexto legal:
{contexto_final}

Pregunta:
{pregunta}

Respuesta:
"""

    output = llm(prompt, max_new_tokens=400, temperature=0.4)[0]["generated_text"]

    return output.split("Respuesta:")[-1].strip() if "Respuesta:" in output else output.strip()

# Interfaz Gradio
demo = gr.Interface(
    fn=responder,
    inputs=gr.Textbox(lines=2, 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.",
    theme="default"
)

demo.launch()