import streamlit as st
from huggingface_hub import hf_hub_download
from llama_cpp import Llama

# 1. Configurazione Interfaccia
st.set_page_config(
    page_title="Limba 2.0", 
    page_icon="https://upload.wikimedia.org/wikipedia/commons/6/65/Flag_of_Sardinia.svg"
)

# 2. Mostra il LOGO 
# Il blocco try-except evita che l'app si blocchi se il file non è ancora presente
try:
    st.image("logo.png", width=150)
except:
    pass # Se il logo non c'è, prosegue senza mostrare errori grafici

# 3. Titolo aggiornato (Senza "Mentor")
st.title("Limba 2.0")
st.markdown("### Su mentore tuo in limba sarda")

# 2. Caricamento del Modello GGUF
@st.cache_resource
def load_model():
    try:
        # Scarica il file specifico dal tuo profilo
        model_path = hf_hub_download(
            repo_id="FPll/limba-mentor-llama3-gguf", 
            filename="Meta-Llama-3.1-8B.Q4_K_M.gguf" 
        )
        # n_threads=2 è il limite ottimale per l'hardware gratuito di HF
        return Llama(model_path=model_path, n_ctx=2048, n_threads=2)
    except Exception as e:
        st.error(f"Errore nel caricamento del modello: {e}")
        return None

with st.spinner("Sto caricando Limba 2.0..."):
    llm = load_model()

# 3. Gestione della Chat
if "messages" not in st.session_state:
    st.session_state.messages = []

# Mostra i messaggi della conversazione corrente
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# Input dell'utente
if prompt := st.chat_input("Iscrie inoghe..."):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    # Generazione della risposta dell'AI
    with st.chat_message("assistant"):
        # Formattazione prompt per Llama 3.1
        full_prompt = f"### Istruzione:\n{prompt}\n\n### Risposta:\n"
        
        if llm:
            response = llm(
                full_prompt,
                max_tokens=512,
                stop=["###", "<|end_of_text|>", "</s>"],
                echo=False
            )
            answer = response["choices"][0]["text"].strip()
            st.markdown(answer)
            st.session_state.messages.append({"role": "assistant", "content": answer})
        else:
            st.error("Il modello non è stato caricato correttamente.")