import os
import sys
import subprocess
import time
import random
import gradio as gr
import torch
import numpy as np
from huggingface_hub import hf_hub_download

# --- 1. Environment Setup & Custom Compilation ---
REPO_DIR = "ac-ai-models"
if not os.path.exists(REPO_DIR):
    print("Cloning repository and setting up custom modules...")
    subprocess.run(["git", "clone", "https://github.com/c4ir-rw/ac-ai-models.git"])
    
    subprocess.run(["pip", "install", "-r", f"{REPO_DIR}/DeepKIN-AgAI/requirements.txt"])
    
    build_dir = f"{REPO_DIR}/DeepKIN-AgAI/monotonic_align"
    subprocess.run(["python", "setup.py", "build_ext", "--inplace"], cwd=build_dir)
    subprocess.run(["pip", "install", "./", "--no-build-isolation"], cwd=build_dir)

sys.path.append(f"{REPO_DIR}/DeepKIN-AgAI")

# --- 2. Load Model & Device Status ---
from deepkin.data.kinya_norm import text_to_sequence
from deepkin.models.flex_tts import FlexKinyaTTS
from deepkin.modules.tts_commons import intersperse

is_gpu = torch.cuda.is_available()
device = torch.device('cuda:0' if is_gpu else 'cpu')
hardware_status = f"🚀 **Hardware Status:** GPU (CUDA) Active" if is_gpu else "🐢 **Hardware Status:** CPU Active"

print("Downloading model checkpoint...")
model_path = hf_hub_download(repo_id="C4IR-RW/kinya-flex-tts", filename="kinya_flex_tts_base_trained.pt")

print("Loading model into memory...")
kinya_tts = FlexKinyaTTS.from_pretrained(device, model_path)
kinya_tts.eval()

# --- 3. Helper Functions ---
def synthesize_audio(text, speaker_name, speed):
    if not text.strip():
        return None, "⚠️ Please enter some text."

    start_time = time.time()
    
    speaker_map = {"Female 1": 0, "Female 2": 1, "Male": 2}
    speaker_id = speaker_map[speaker_name]

    text_id_sequence = intersperse(text_to_sequence(text, norm=True), 0)

    x = torch.LongTensor(text_id_sequence).unsqueeze(0).to(device)
    x_lengths = torch.LongTensor([x.size(1)]).to(device)
    sid_t = torch.LongTensor([speaker_id]).to(device)

    with torch.no_grad():
        audio_data = kinya_tts.flex_tts.infer(
            x, x_lengths, sid_t, noise_scale=0.667, length_scale=(1.0 / speed)
        )[0].cpu().numpy()

    audio_np = audio_data.squeeze()
    
    end_time = time.time()
    time_taken = round(end_time - start_time, 2)
    time_message = f"⏱️ *Generated in {time_taken} seconds*"
    
    return (24000, audio_np), time_message

def get_random_sentence():
    sentences = [
        "Ikiremwamuntu cyose kivukana umudendezo kandi kingana mu cyubahiro n'uburenganzira.",
        "Hinga neza, wite ku myaka yawe kugira ngo ubone umusaruro mwiza.",
        "Koresha ifumbire nziza kugira ngo ubutaka bwawe burumbuke.",
        "Iki gihugu cyacu ni cyiza cyane, kigizwe n'imisozi igihumbi.",
        "Uhinzi mwiza amenya igihe cyo gutera n'igihe cyo gusarura.",
        "Amazi ni isoko y'ubuzima, tuyabungabunge kugira ngo atagirira nabi ibihingwa.",
        "Gufata neza umusaruro nyuma y'isarura ni ingenzi kugira ngo udapfa ubusa.",
        "Ubwishingizi bw'ibihingwa bugufasha kwirinda igihombo igihe habaye ibiza."
    ]
    return random.choice(sentences)

# --- 4. Build Custom UI ---
custom_theme = gr.themes.Soft(
    primary_hue="emerald",
    neutral_hue="slate",
    font=[gr.themes.GoogleFont("Inter"), "sans-serif"]
)

custom_css = """
.header-text { text-align: center; padding-bottom: 10px; }
.time-tracker { text-align: right; color: #64748b; margin-top: 5px; }
"""

with gr.Blocks(theme=custom_theme, css=custom_css, title="Kinyarwanda TTS") as demo:
    
    # Header
    with gr.Column(elem_classes="header-text"):
        gr.Markdown("# 🇷🇼 Kinyarwanda AI Voice Generator")
        gr.Markdown("Transform text into natural-sounding speech, powered by models developed for Rwanda's agricultural advisory hotline.")
        gr.Markdown(hardware_status)

    gr.HTML("<hr>")

    # Main Interface
    with gr.Row():
        # Left Column: Inputs
        with gr.Column(scale=5):
            with gr.Group():
                with gr.Row():
                    input_text = gr.Textbox(
                        label="Kinyarwanda Text", 
                        placeholder="Type a sentence or click the random button...", 
                        lines=4,
                        # show_copy_button=True
                    )
                with gr.Row():
                    random_btn = gr.Button("🎲 Generate Random Sentence", size="sm")
                    clear_btn = gr.ClearButton([input_text], value="🗑️ Clear", size="sm")
            
            speaker_choice = gr.Radio(
                choices=["Female 1", "Female 2", "Male"], 
                label="Voice Profile", 
                value="Male",
            )
            
            speed_slider = gr.Slider(
                minimum=0.5,
                maximum=2.0,
                value=1.0,
                step=0.1,
                label="Speaking Speed",
                info="Lower value means slower speed, higher means faster (e.g., 0.8 is slightly slower than normal)."
            )
            
            generate_btn = gr.Button("🔊 Synthesize Speech", variant="primary", size="lg")

        # Right Column: Outputs
        with gr.Column(scale=4):
            with gr.Group():
                output_audio = gr.Audio(label="Audio Output", type="numpy", interactive=False)
                output_time = gr.Markdown(label="", elem_classes="time-tracker")

    gr.HTML("<br><hr>")

    # Collapsible Acknowledgments
    with gr.Accordion("📚 About the Model & Credits", open=False):
        gr.Markdown("""
        * **Model Owners:** [Center for the Fourth Industrial Revolution Rwanda (C4IR)](https://c4ir.rw/) & KiNLP
        * **Support & Finance:** Supported by GIZ and financed by BMZ.
        * **License:** Creative Commons Attribution 4.0 International License (CC-BY 4.0).
        * **Dataset:** Trained on the [Kinyarwanda Agricultural TTS Dataset](https://huggingface.co/datasets/C4IR-RW/kinya-ag-tts).
        * **Architecture:** Multi-speaker TTS model with MB-iSTFT-VITS2.
        """)

    # --- 5. Event Routing ---
    # Wire up the random sentence generator
    random_btn.click(
        fn=get_random_sentence,
        inputs=[],
        outputs=[input_text]
    )

    # Wire up the main generation button
    generate_btn.click(
        fn=synthesize_audio,
        inputs=[input_text, speaker_choice, speed_slider],
        outputs=[output_audio, output_time]
    )

if __name__ == "__main__":
    demo.launch()