import streamlit as st
from ragpipeline import (RAGPipeline,Retriever,ChatBot)
import tempfile
import pandas as pd
from textutils import ParagraphDocumentProcessor, DocumentProcessor,WholeTextDocumentProcessor
from HFChatbot import HFBot
import os

###########################################
# MAIN CON STREAMLIT
###########################################
def main():
    UPLOAD_DIR = "/tmp/"
    os.makedirs(UPLOAD_DIR, exist_ok=True)
    
    codice_tabella = f"<table><tr><td>💡AURA:</td><td> AI-Driven Unified Regulatory Audit</td></tr></table>"
    st.markdown(codice_tabella, unsafe_allow_html=True)

    if "faiss_builder" not in st.session_state:
        ragpipeline = RAGPipeline(numero_frammenti=10)
        st.session_state["faiss_builder"] = ragpipeline
    else:
        ragpipeline = st.session_state["faiss_builder"]

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

    if "indice_creato" not in st.session_state:
        st.session_state["indice_creato"] = False

    modelliLLM = [
        'Almawave/Velvet-2B',
        'Almawave/Velvet-14B',
        'mistralai/Mistral-7B-Instruct-v0.1',
    ]
    modelliOllama = [
        
        'deepseek-r1:1.5b',
        'qwen2.5:7b',
        'deepseek-r1:7b',
        'llama3.2:3b',
        'Almawave/Velvet:2B',
        'Almawave/Velvet:14b',
        'mistral:latest',
        "vaiton/minerva",
        'qwen2.5:0.5b',
        'qwen3:4b',
        'minerva',
        'nemo',
        'deepseek-r1:14b',
        'qwen3:14b',
        'phi4-mini-reasoning',
        'gemma3:12b',
        
        ]

    modello_scelto = st.selectbox("Seleziona un modello:", modelliLLM)
    st.write(f"Hai selezionato: {modello_scelto}")

    st.title("Suddivisione in paragrafi")
    docprocessor_options = {
        "ParagraphDocumentProcessor": ParagraphDocumentProcessor(),
        "WholeText": WholeTextDocumentProcessor()
    }
    selected_docprocessor = st.selectbox("Divisione in paragrafi", docprocessor_options.keys())
    docprocessor = docprocessor_options[selected_docprocessor]
    st.write(f"Hai selezionato: **{selected_docprocessor}**")

    ragpipeline.docprocessor = docprocessor
    if not st.session_state["indice_creato"]:
        st.subheader("Carica l'atto principale (Determinazione)")

        main_pdf = st.file_uploader("Carica 1 file PDF - Determinazione", type=["pdf"], key="main_pdf")
        if main_pdf:
            st.session_state["main_pdf_nome"] = main_pdf.name
        

        st.subheader("Carica eventuali Allegati PDF multipli")
        other_pdfs = st.file_uploader("Carica allegati (puoi caricare più PDF)", 
                                    type=["pdf"], 
                                    accept_multiple_files=True,
                                    key="allegati_pdf")
        if st.button("Crea indice FAISS"):
            if main_pdf is not None:



                try:
                    save_path = os.path.join(UPLOAD_DIR, main_pdf.name)
                    with open(save_path, "wb") as f:
                        f.write(main_pdf.read())
                    st.success(f"✅  Determinazione caricata con successo! File salvato in: `{save_path}`")
                    ragpipeline.aggiungi_file_pdf(save_path)
                except Exception as e:
                    st.error(f"❌ Errore nel salvataggio del file: {e}")

            else:
                st.warning("Nessun PDF di Determinazione caricato.")

            if other_pdfs:
                for uploaded_file in other_pdfs:
                    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf", dir=UPLOAD_DIR) as tmp_file:
                        tmp_file.write(uploaded_file.read())
                        tmp_path = tmp_file.name
                        ragpipeline.aggiungi_file_pdf(tmp_path)
                        st.subheader(f"Caricato file {tmp_file.name}")
                st.success("Allegati caricati con successo!")
            else:
                st.info("Nessun allegato caricato.")

            ragpipeline.crea_indice()
            st.success("Indice FAISS generato e caricato.")
            st.session_state["indice_creato"] = True
            frammenti_recuperati =ragpipeline.attributi_frammenti
            for frammento_recuperato in frammenti_recuperati:
                RAGPipeline.dump_excel(dizionario=frammento_recuperato, filename="frammentiChatbot.xlsx")

    if st.session_state["indice_creato"]:
        with st.form(key="domanda_form"):
            domanda = st.text_area("Inserisci la domanda", key="domanda_input")
            istruzione = st.text_area("Inserisci le istruzioni", key="istruzione_input")
            submit_button = st.form_submit_button("Analizza atto")

        if submit_button:
            if domanda.strip().upper() == "FINE":
                st.stop()

            #cb = ChatBot(model_name="flaollama", model_orig=modello_scelto)

            cb  = HFBot(model_name=modello_scelto)

            ret = Retriever(
                indice=ragpipeline.indice,
                sentence_transformer_model=ragpipeline.sentence_transformer_model,
                query=domanda + istruzione,
                documenti=ragpipeline.documenti,
                frammenti_indicizzati=ragpipeline.frammenti_indicizzati,
                attributi_frammenti=ragpipeline.attributi_frammenti
            )

            ret.esegui_query(top_k=3)

            risposta = cb.generate(
                query=domanda,
                relevant_docs=ret.passaggi_rilevanti,
                attributi_frammenti_rilevanti=ret.attributi_rilevanti,
                istruzioni=istruzione
            )

            st.session_state.storico_domande.append((modello_scelto, domanda, istruzione, risposta))

            st.markdown(
                f"<p><strong>Domanda:</strong> {domanda}  <br/>"
                f"<strong>Istruzioni:</strong> <em>{istruzione}</em><br/><br/>"
                f"<strong>Risposta:</strong><em> {risposta}</em></p>",
                unsafe_allow_html=True
            )
            id_frammenti_recuperati =  ":".join(sorted(set(elemento['id'] for elemento in ret.attributi_rilevanti)))
            dump  = {
                 'timestamp': ragpipeline.timestamp,
                "modello": cb.model_orig, 
                "documenti": st.session_state.get("main_pdf_nome", "non disponibile"),
                "file_recuperati": "",
                "file_gold": "",
                "frammenti_recuperati":id_frammenti_recuperati,
                "frammenti_gold": "", 
                "domanda":domanda,
                "istruzioni":istruzione, 
                "risposta_gold": " ", 
                "risposta":cb.pulisci_risposta(risposta)}

            RAGPipeline.dump_excel(dizionario=dump,filename="dumpChatbot.xlsx")


    if st.session_state.storico_domande:
        st.markdown("---")
        st.subheader("Storico delle domande analizzate")
        for idx, (mymod, q, inst, resp) in enumerate(st.session_state.storico_domande, 1):
            st.markdown(
                f"""
                **{idx}. Domanda:** {q}<br/>
                <em>Modello: </em>{mymod}<br/>
                <em>Istruzioni:</em> {inst}<br/><br/>
                <strong>Risposta:</strong> {resp}<br/>
                """,
                unsafe_allow_html=True
            )

if __name__ == "__main__":
    main()