import streamlit as st
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# Cargar modelo
model_name = "Qwen/Qwen3-1.7B"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="auto",
    trust_remote_code=True
)

# Cargar prompt base
with open("SeguroCargo.txt", "r", encoding="utf-8") as f:
    prompt_base = f.read().strip()

# Función para generar respuesta con "thinking" opcional
def generar_respuesta(pregunta):
    contexto = f"{prompt_base}\nUsuario: {pregunta}\nAsistente:"
    messages = [{"role": "user", "content": contexto}]

    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False
    )

    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
    generated_ids = model.generate(
        **model_inputs,
        max_new_tokens=1024,
        do_sample=True,
        temperature=0.7,
        top_p=0.8
    )

    output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()

    try:
        index = len(output_ids) - output_ids[::-1].index(151668)
    except ValueError:
        index = 0

    content = tokenizer.decode(output_ids[index:], skip_special_tokens=True).strip("\n")
    return content.strip()

# Configuración de Streamlit
st.set_page_config(page_title="Asistente SeguroCargo2", page_icon="📦")
st.title("📦 Bienvenido a SeguroCargo2")
st.markdown("Hola, yo soy **SEGU**, asistente virtual de SeguroCargo sin el modo de pensamiento!")
st.markdown("Consulta sobre tus **envíos nacionales e internacionales** de manera rápida y profesional.")

if "messages" not in st.session_state:
    st.session_state.messages = []

# Mostrar historial de mensajes
for msg in st.session_state.messages:
    role, text = msg
    with st.chat_message(role):
        st.markdown(text)

# Entrada del usuario
user_input = st.chat_input("Escribe tu mensaje...")

if user_input:
    st.session_state.messages.append(("user", user_input))
    with st.chat_message("user"):
        st.markdown(user_input)

    # Generar y mostrar respuesta
    respuesta = generar_respuesta(user_input)
    st.session_state.messages.append(("assistant", respuesta))
    with st.chat_message("assistant"):
        st.markdown(respuesta)