import subprocess
import gradio as gr
from TTS.api import TTS
import os
import time
import torch
import json

# Initialisation du modèle TTS
tts = TTS("tts_models/multilingual/multi-dataset/xtts_v2", gpu=False)

# Dossier de base pour les projets
base_output_folder = "projects"
os.makedirs(base_output_folder, exist_ok=True)

# Fonction pour créer ou charger un projet
def create_or_load_project(project_name, speaker, agree):
    if not agree:
        raise gr.Error("❗ Veuillez accepter les conditions d'utilisation.")
    
    project_folder = os.path.join(base_output_folder, project_name)
    os.makedirs(project_folder, exist_ok=True)
    
    project_data = {
        "name": project_name,
        "speaker": speaker,
        "sections": []
    }
    
    project_file = os.path.join(project_folder, "project_data.json")
    if os.path.exists(project_file):
        with open(project_file, "r") as f:
            project_data = json.load(f)
    else:
        with open(project_file, "w") as f:
            json.dump(project_data, f)
    
    return project_data

# Fonction pour ajouter ou mettre à jour une section
def update_section(project_data, section_name, section_text):
    for section in project_data["sections"]:
        if section["name"] == section_name:
            section["text"] = section_text
            return project_data
    
    project_data["sections"].append({
        "name": section_name,
        "text": section_text,
        "audio_path": None
    })
    return project_data

# Fonction pour générer l'audio d'une section
def generate_section_audio(project_data, section_name):
    project_folder = os.path.join(base_output_folder, project_data["name"])
    for section in project_data["sections"]:
        if section["name"] == section_name:
            output_path = os.path.join(project_folder, f"{section_name}.wav")
            tts.tts_to_file(
                text=section["text"],
                file_path=output_path,
                speaker_wav=[os.path.join("examples", f) for f in os.listdir("examples") if f.startswith(project_data["speaker"]) and f.endswith(".wav")],
                language="fr"
            )
            section["audio_path"] = output_path
            return project_data, output_path
    
    raise gr.Error(f"❌ Section '{section_name}' non trouvée.")

# Interface Gradio
with gr.Blocks() as demo:
    gr.Markdown("# 🎙️ Projet Margaux TTS")
    
    # État du projet
    project_state = gr.State(value={})
    
    # Étape 1: Création ou chargement du projet
    with gr.Tab("🆕 Création/Chargement du Projet"):
        project_name = gr.Textbox(label="📝 Nom du projet")
        speaker = gr.Dropdown(label="🔊 Voix", choices=["Margaux"], value="Margaux")
        agree = gr.Checkbox(label="✅ J'accepte les conditions d'utilisation")
        create_btn = gr.Button("🚀 Créer/Charger le Projet")
    
    # Étape 2: Gestion des sections
    with gr.Tab("📋 Gestion des Sections"):
        sections_list = gr.List(label="📜 Sections du projet")
        section_name = gr.Textbox(label="📝 Nom de la section")
        section_text = gr.Textbox(label="✍️ Texte de la section", lines=5)
        add_section_btn = gr.Button("➕ Ajouter/Mettre à jour la section")
        generate_section_btn = gr.Button("🎤 Générer l'audio de la section")
        audio_output = gr.Audio(label="🔊 Audio généré")
    
    # Fonctions de callback
    def load_project(project_name, speaker, agree):
        project_data = create_or_load_project(project_name, speaker, agree)
        sections = [section["name"] for section in project_data["sections"]]
        return project_data, gr.List.update(value=sections)
    
    def add_section(project_data, section_name, section_text):
        updated_project = update_section(project_data, section_name, section_text)
        sections = [section["name"] for section in updated_project["sections"]]
        return updated_project, gr.List.update(value=sections)
    
    def generate_audio(project_data, section_name):
        updated_project, audio_path = generate_section_audio(project_data, section_name)
        return updated_project, audio_path
    
    # Connexion des événements
    create_btn.click(load_project, inputs=[project_name, speaker, agree], outputs=[project_state, sections_list])
    add_section_btn.click(add_section, inputs=[project_state, section_name, section_text], outputs=[project_state, sections_list])
    generate_section_btn.click(generate_audio, inputs=[project_state, section_name], outputs=[project_state, audio_output])

# Lancement de l'interface
demo.launch()