import streamlit as st
from utils.audit.response_llm import generate_response_via_langchain
from textwrap import dedent
#st.set_page_config(page_title="Agents recommandés", page_icon="", layout="wide")
def remove_images_from_content(content):
    filtered_content = {}
    
    for page, data in content.items():
        # Create a new dictionary excluding the "images" key
        filtered_data = {key: value for key, value in data.items() if key != "images"}
        filtered_content[page] = filtered_data
    
    return filtered_content

def recommended_agent_main():
    st.title("Agents recommandés")

    if "audit" not in st.session_state or "audit" is None:
        st.error("Veuillez d'abord effectuer un audit pour obtenir des recommandations d'agents.")
        return
    
    audit = st.session_state.audit_simplified
    content = st.session_state.audit["content"]

    if "response_llm" not in st.session_state:
        st.session_state.response_llm = ""


    #filter content, delete images if type is pdf
    if audit["type de fichier"] == "pdf":
        content = remove_images_from_content(content)
    #delete audio if type is audio and keep transcript
    elif audit["type de fichier"] == "audio":
        content = content["transcription"]

    ressources = content

    prompt = '''Tu es designer en intelligence artificielle (IA) spécialisé dans la création d'agents IA autonomes et performants.
        A partir de ressources fournies par l'utilisateur (texte, documents, images, audio), tu es chargé de suggérer la création d'agents autonomes pour mettre en pratique les informations contenues dans les ressources fournies.

        Tu proposes deux solutions :

        Sol. A : 1 seul agent IA dont tu suggéreras :
        * Nom
        * Rôle
        * Objectifs
        * Outils utilisés par l'agent
        * Tâches réalisées par l'agents
        * Compétences de l'agent (backstory)

        Sol. B : 1 équipe d'agents tu suggéreras :
        * Le nombre d'agents
        * Pour chacune d'eux [Nom, Rôle, Objectifs, Outils utilisés par l'agent, Tâches réalisées par l'agents, Compétences de l'agent (backstory)]

        Une fois ce travail réalisé, tu proposes une série de 3 missions avec objectifs SMART pour chacun des agents  Sol. A et Sol. B en présentation les résultats dans un tableau contenant :
        Nom de l’agent
        Objectifs à atteindre
        '''
    
    #display prompt and modify it 
    prompt_modified = st.text_area("Prompt", prompt, height=300)
    prompt_modified = dedent(prompt_modified)
    if st.button("Générer les recommandations"):
        resource_prompt = f'''Ressources fournies par l'utilisateur :{ressources}'''
        prompt_modified = f"{prompt_modified}\n{resource_prompt}"
        st.session_state.response_llm = st.write_stream(generate_response_via_langchain(query=prompt_modified,stream=True))
    elif st.session_state.response_llm:
        st.info("la dernière réponse générée est affichée ci-dessous")
        st.write(st.session_state.response_llm)

recommended_agent_main()