import os
import json
from textwrap import dedent

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

# Configuration des clés API
os.environ["GEMINI_API_KEY"] = os.environ.get("GEMINI_API_KEY")
os.environ["SERPER_API_KEY"] = os.environ.get("SERPER_API_KEY")

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

# Initialisation de l'outil de recherche
search_tool = SerperDevTool()

# Définition des agents
researcher = Agent(
    role='Chercheur de Sujets',
    goal=dedent("""Trouver les informations les plus pertinentes et précises sur {topic}
    en utilisant l'API SerpApi."""),
    backstory=dedent("""Un chercheur expert spécialisé dans la collecte d'informations sur divers sujets.
    Capable d'utiliser l'API SerpApi pour des recherches précises et efficaces."""),
    tools=[search_tool],
    llm=llm,
    verbose=True,
    allow_delegation=False
)

writer = Agent(
    role='Rédacteur de Flashcards',
    goal=dedent("""Créer des flashcards claires et concises en format question-réponse
    basées sur les informations fournies par le Chercheur."""),
    backstory=dedent("""Un expert en pédagogie et en création de matériel d'apprentissage.
    Capable de transformer des informations complexes en flashcards simples et mémorisables."""),
    llm=llm,
    verbose=True,
    allow_delegation=False
)

def extract_json_from_result(result_text: str) -> list:
    """Extrait le JSON des résultats de la crew."""
    try:
        # Trouve le début et la fin du JSON dans le texte
        json_start = result_text.find('[')
        json_end = result_text.rfind(']') + 1
        if json_start == -1 or json_end == 0:
            raise ValueError("JSON non trouvé dans le résultat")
        
        json_str = result_text[json_start:json_end]
        return json.loads(json_str)
    except (json.JSONDecodeError, ValueError) as e:
        raise ValueError(f"Erreur lors de l'extraction du JSON : {str(e)}")

# Définition des tâches
def research_task(topic: str) -> Task:
    return Task(
        description=dedent(f"""Effectuer une recherche approfondie sur le sujet '{topic}'.
        Compiler les informations les plus pertinentes et les plus récentes."""),
        expected_output="Une liste d'informations pertinentes sur le sujet.",
        agent=researcher
    )

def flashcard_creation_task(research_task: Task) -> Task:
    return Task(
        description=dedent("""Transformer les informations fournies par le Chercheur
        en une série de flashcards au format JSON. Chaque flashcard doit avoir une question
        d'un côté et une réponse concise de l'autre."""),
        expected_output="Une liste de flashcards au format JSON.",
        agent=writer,
        context=[research_task]
    )

# Interface Streamlit
st.title("🤖 Générateur de Flashcards avec CrewAI")

topic = st.text_input("Entrez le sujet des flashcards:", "Intelligence Artificielle")

if st.button("Générer les Flashcards"):
    if not topic:
        st.error("Veuillez entrer un sujet.")
    else:
        with st.spinner('Création des flashcards en cours...'):
            try:
                # Création des tâches
                research = research_task(topic)
                flashcard_creation = flashcard_creation_task(research)

                # Création de la crew
                crew = Crew(
                    agents=[researcher, writer],
                    tasks=[research, flashcard_creation],
                    process=Process.sequential,
                    verbose=True
                )

                # Exécution de la crew
                result = crew.kickoff()
                
                # Extraction du JSON depuis le dernier résultat de tâche
                if result.tasks_output and len(result.tasks_output) > 0:
                    last_task_output = result.tasks_output[-1].raw
                    flashcards = extract_json_from_result(last_task_output)
                    
                    # Affichage des flashcards
                    st.success("Flashcards générées avec succès !")
                    
                    # Création d'onglets pour différents modes d'affichage
                    tab1, tab2 = st.tabs(["Mode Étude", "Mode JSON"])
                    
                    with tab1:
                        for i, card in enumerate(flashcards, 1):
                            with st.expander(f"Flashcard {i}: {card['question']}", expanded=False):
                                st.write("**Réponse:**", card['answer'])
                    
                    with tab2:
                        st.json(flashcards)
                
                else:
                    st.error("Aucun résultat généré par la crew.")
                    
            except Exception as e:
                st.error(f"Une erreur est survenue : {str(e)}")
                st.write("Résultat brut pour débogage :")
                st.write(result)