import os
import ast
import streamlit as st
from crewai import Agent, Crew, Process, Task
from crewai_tools import SerperDevTool
from langchain_google_genai import ChatGoogleGenerativeAI
from crewai import LLM
from typing import List

# Configuration des clés API
os.environ["GEMINI_API_KEY"] = "AIzaSyD6yZxfVOnh63GXBJjakAupk9aP4CZrgrQ"
os.environ["SERPER_API_KEY"] = "9b90a274d9e704ff5b21c0367f9ae1161779b573"

llm = LLM(
    model="gemini/gemini-1.5-flash",
    temperature=0.7,
    timeout=120,
    max_tokens=8000,
)

# Définitions des agents (inchangées)
project_manager = Agent(
    role="Chef de Projet",
    goal="Coordonner les autres agents pour assurer la cohérence et la qualité de l'exposé.",
    backstory="Un chef de projet expérimenté avec une expertise dans la gestion d'équipes et la coordination de projets complexes.",
    llm=llm,
    verbose=True,
    allow_delegation=True,
)

topic_analyst = Agent(
    role="Analyste de Thème",
    goal="Analyser le thème de l'exposé et générer un plan détaillé.",
    backstory="Un expert en recherche et en structuration de contenu, capable de décomposer des sujets complexes en plans clairs et concis.",
    llm=llm,
    verbose=True,
    tools=[SerperDevTool()],
)

content_writer = Agent(
    role="Rédacteur de Contenu",
    goal="Rédiger les sections de l'exposé en se basant sur le plan et les recherches.",
    backstory="Un rédacteur spécialisé dans la création de contenu clair, concis et engageant.",
    llm=llm,
    verbose=True,
    tools=[SerperDevTool()],
)

editor = Agent(
    role="Éditeur/Réviseur",
    goal="Réviser et peaufiner le contenu, assurer la cohérence du style et du ton, corriger les erreurs.",
    backstory="Un éditeur expérimenté avec un souci du détail et une excellente maîtrise de la langue.",
    llm=llm,
    verbose=True,
)

def parse_sections(crew_output) -> List[str]:
    """Parse the sections from CrewOutput object."""
    try:
        # Extraire le raw output
        raw_output = crew_output.raw
        # Nettoyer la string (enlever les guillemets externes s'ils existent)
        cleaned_output = raw_output.strip('"\'')
        # Évaluer la string comme une liste Python
        sections = eval(cleaned_output)
        return sections
    except Exception as e:
        st.error(f"Erreur lors du parsing des sections: {e}")
        st.error(f"Raw output: {crew_output.raw}")
        return []

def assign_tasks(topic: str) -> List[Task]:
    # Tâche d'analyse qui retourne une liste de sections
    analyze_topic_task = Task(
        description=f"""Analyser le thème '{topic}' et générer un plan détaillé pour l'exposé.
        IMPORTANT: Ta réponse doit être UNIQUEMENT une liste Python.
        Exemple de réponse: ['Introduction', 'Historique', 'État actuel', 'Perspectives futures', 'Conclusion']
        Ne pas ajouter d'autres explications ou texte.""",
        expected_output="Une liste Python contenant les sections de l'exposé",
        agent=topic_analyst
    )
    
    tasks = [analyze_topic_task]
    
    # Création dynamique des tâches d'écriture pour chaque section
    def create_writing_tasks(sections: List[str]) -> List[Task]:
        writing_tasks = []
        for section in sections:
            write_task = Task(
                description=f"""Rédiger la section '{section}' de l'exposé sur le thème '{topic}'.
                Concentre-toi uniquement sur cette section spécifique.
                Utilise un style académique et professionnel.""",
                expected_output=f"Contenu détaillé de la section '{section}' au format markdown",
                agent=content_writer,
                context=[analyze_topic_task]
            )
            writing_tasks.append(write_task)
        return writing_tasks
    
    # Tâche d'édition finale
    def create_editing_task(writing_tasks: List[Task]) -> Task:
        return Task(
            description="Réviser et harmoniser l'ensemble du contenu de l'exposé. Assurer la cohérence entre toutes les sections.",
            expected_output="Version finale de l'exposé au format",
            agent=editor,
            context=writing_tasks
        )
    
    return {
        'analyze_task': analyze_topic_task,
        'create_writing_tasks': create_writing_tasks,
        'create_editing_task': create_editing_task
    }

# Interface Streamlit
st.title("🤖 Générateur d'Exposés avec CrewAI")

topic = st.text_input("Entrez le thème de l'exposé:", "L'Intelligence Artificielle en 2024")

if st.button("Générer l'exposé"):
    with st.spinner("Création de l'équipe d'agents..."):
        task_creators = assign_tasks(topic)
        
        # Exécution de la tâche d'analyse
        crew_analysis = Crew(
            agents=[project_manager, topic_analyst],
            tasks=[task_creators['analyze_task']],
            process=Process.hierarchical,
            manager_llm=llm,
            verbose=True,
        )
        analysis_result = crew_analysis.kickoff()
        
        # Parsing des sections depuis le résultat
        sections = parse_sections(analysis_result)
        st.write("Sections identifiées:", sections)
        
        if sections:
            # Création des tâches d'écriture basées sur les sections
            writing_tasks = task_creators['create_writing_tasks'](sections)
            
            # Création de la tâche d'édition
            editing_task = task_creators['create_editing_task'](writing_tasks)
            
            # Toutes les tâches dans l'ordre
            all_tasks = [task_creators['analyze_task']] + writing_tasks + [editing_task]
            
            # Création et exécution du crew final
            crew = Crew(
                agents=[project_manager, topic_analyst, content_writer, editor],
                tasks=all_tasks,
                process=Process.hierarchical,
                manager_llm=project_manager,
                verbose=True,
            )

            with st.spinner("Génération de l'exposé..."):
                result = crew.kickoff()
                
                st.success("Exposé généré avec succès!")
                st.markdown(result)
        else:
            st.error("Impossible de continuer sans une liste valide de sections.")