import streamlit as st
import os
import tempfile
from HFBot import HFBot
from textutils import ParagraphDocumentProcessor, WholeTextDocumentProcessor, SmallFragmentDocumentProcessor
import pandas as pd
from ragpipeline import RAGPipeline, Retriever, OllamaChatbot

def main():
    modelliVelvet = [
            'Almawave/Velvet-2B',
            'Almawave/Velvet-14B',
            
    ]
    modelliLLM = [
            'Almawave/Velvet-2B',
            'Almawave/Velvet-14B',
            'mistralai/Mistral-7B-Instruct-v0.1',
            'Qwen/Qwen2.5-1.5B',
        ]
    
    modelliOllama = [
        'Almawave/Velvet:2B',
        'Almawave/Velvet:14b',
        'llama3.1:8b-instruct-q4_K_M',
        'qwen3:14b',
        'qwen3:30b-a3b'
                
    ]
    ## indica se sono sullo spaces di HF (deve essere inserita uan variabile I_AM_ON_HF)
    sono_su_hf =os.environ.get('I_AM_ON_HF', False)
    ## se sono su Hugginh Face non uso ollama
    if  not  sono_su_hf:
        modelliLLM.append("----- USARE SOLO CON OLLAMA -----")
        for mollama in modelliOllama:
            modelliLLM.append(mollama)
    
    UPLOAD_DIR="/tmp/"
    if "indice_creato" not in st.session_state:
        st.session_state["indice_creato"] = False
    if "faiss_builder" not in st.session_state:
        ragpipeline = RAGPipeline( )
    codice_tabella = f"<table><tr><td>💡AURA:</td><td> AI-Utilizzata per la Regolarità Amministrativa</td></tr></table>"
    st.markdown(codice_tabella, unsafe_allow_html=True)
    st.title("Cosa è AURA?")
    st.write("""
            Questo strumento, attualmente in fase sperimentale, è stato sviluppato per eseguire controlli di 
            regolarità amministrativa ai sensi dell’art. 147-bis del D.Lgs. 267/2000, 
            con riferimento agli atti relativi al PNRR.
            È in continua evoluzione. Per testarne il funzionamento, è sufficiente caricare un file PDF contenente 
            una determinazione dirigenziale.
            AURA analizzerà il documento e fornirà risposte basate su una check-list predefinita di domande.
             <p>
             AURA è un sistema RAG <em>Retrieval Augmented Generation</em> che dato un atto amministrativo ed eventuali allegati
             ed una o più domande contenute in una <em>check list</em> di regolarità amministrativa, ricerca nei documenti
             i frammenti rilevanti per la domanda; questi assieme ad alcune istruzioni (<em>In-context learning</em>) 
             vengono inviati  ad un LLM <em>Large Language Model</em> 
             al fine  di generare una risposta corretta e coerente con i frammenti rilevanti.
             </p>
             <p>
             Questa versione di AURA utilizza Velvet:2B di Almawave, rilasciato sotto 
             <a href="https://www.apache.org/licenses/LICENSE-2.0">Licenza Apache 2.0</a>  come LLM.
             </p>
            
             """ , unsafe_allow_html=True)

    st.warning("Attenzione questo tool è sperimentale. AURA può sbagliare")
    if not sono_su_hf: 
        modello_scelto = st.selectbox("Seleziona un modello:", modelliLLM, index=0)
    else:
        modello_scelto = st.selectbox("Seleziona un modello:", modelliVelvet, index=0)
    st.write(f"Hai selezionato: {modello_scelto}")
    
    if not sono_su_hf:
        st.title("Generazione testo")
        generatoriLLM =  {
            'Hugging Face Transformers': "HF",
            'Ollama in locale' :"OLLAMA"
        }
        selected_generator= st.selectbox("Scegli lo strumento per interagire con LLM", generatoriLLM.keys())
        chiave_LLM=generatoriLLM[selected_generator]            
        st.write(f"Hai selezionato {selected_generator}")
    else:
        chiave_LLM="HF"

        
    st.title("Suddivisione in paragrafi")
    docprocessor_options = {
        "Small Fragments (più veloce ma poco preciso)": SmallFragmentDocumentProcessor(),
        "ParagraphDocumentProcessor (più lento e leggermente più preciso)": ParagraphDocumentProcessor(),
        "WholeText (viene generato un solo grande frammento, può confondere gli LLM)": WholeTextDocumentProcessor(),
    }
    selected_docprocessor = st.selectbox("Divisione in paragrafi", docprocessor_options.keys())
    ragpipeline.docprocessor = docprocessor_options[selected_docprocessor]
    st.write(f"Hai selezionato: **{selected_docprocessor}**")
    
    if not st.session_state["indice_creato"]:
        st.subheader("Carica Allegati PDF multipli", help="es: determinazione ed allegati parte integranrte")
        other_pdfs = st.file_uploader("Carica allegati (puoi caricare più PDF)", 
                                  type=["pdf"], 
                                  accept_multiple_files=True,
                                  key="allegati_pdf",
                                  )

        if st.button("Esegui analisi"):
            if other_pdfs is not None:
                file_i=1
                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.success(f"Caricato file  N. {file_i}")
                    file_i=file_i+1
                ragpipeline.crea_indice( )
                st.success("Indice FAISS generato e caricato.")
                st.session_state["indice_creato"] = True

            else:
                st.warning("Nessun PDF di Determinazione caricato.")
        if st.session_state["indice_creato"]:
            st.write("Eseguo controlli")
            csv_file ="domandePNRR.xlsx"
            if chiave_LLM ==  "HF":
                    LLM=HFBot(model_name=modello_scelto)
            elif chiave_LLM == "OLLAMA":
                    LLM=OllamaChatbot(model_name="flaollama",model_orig=modello_scelto)
            else:
                    LLM=HFBot(model_name=modello_scelto)

            df = pd.read_excel(csv_file)
            if "domanda" in df.columns:
                domande = df["domanda"].tolist()
                istruzioni = df["istruzioni"].tolist();
                #print(f"ISTRUZIONI {istruzioni}")
                st.success(f"Caricate {len(domande)} domande!")
                if ragpipeline.indice is not None :
                    #cb = ChatBot(model_name="flaollama", model_orig=modello_scelto)
                    st.write("**Risposte automatiche per ciascuna domanda **:")
                    i = 0
                    for domanda in domande:
                        istruzione = istruzioni[i]
                        ret = Retriever (indice=ragpipeline.indice,
                                         sentence_transformer_model=ragpipeline.sentence_transformer_model,
                                         query=domanda ,
                                         documenti=ragpipeline.documenti,
                                         frammenti_indicizzati= ragpipeline.frammenti_indicizzati,
                                         attributi_frammenti= ragpipeline.attributi_frammenti
                                         )
                        ret.esegui_query( top_k=3  )
                        risposta = LLM.generate(query=domanda,
                                               relevant_docs=ret.passaggi_rilevanti,attributi_frammenti_rilevanti=ret.attributi_rilevanti,
                                               istruzioni=istruzione)
 
                        m_domanda = st.chat_message("user")
                        m_domanda.write(domanda)

                        m_risposta = st.chat_message("assistant")
                        m_risposta.write(risposta)
 
                        i = i+1



            

if __name__ == "__main__":
    main()