import argparse
import json
import os
import re
import torch
from torch import no_grad, LongTensor
import commons
import utils
import gradio as gr
import numpy as np
import librosa
from models import SynthesizerTrn
from text import text_to_sequence
from mel_processing import spectrogram_torch

# --- CSS: SEMUA LOADING PINK, RINGAN, HALUS, TANPA SHADOW ---
custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Quicksand:wght@400;700&display=swap');

body, .gradio-container {
    font-family: 'Quicksand', sans-serif !important;
    background-color: #ffffff !important;
}

/* --- KUSTOMISASI LOADING GRADIO MENJADI PINK --- */
/* 1. Loading Bar di bagian atas saat proses (Generate) */
.progress-view {
    background-color: #ffe5ec !important; /* Warna dasar bar (pink sangat muda) */
    height: 6px !important;
}
.progress-view > div {
    background-color: #ff6b81 !important; /* Warna pink cerah yang jalan */
}

/* 2. Loading Spinner/Animasi awal aplikasi */
.gradio-container .load-6 .loader {
    border-top-color: #ff4d6d !important;
}

/* 3. Menghilangkan Shadow Hitam Default Gradio agar terlihat Glowing Pink */
.status-card {
    background: #ffffff !important;
    border: 1.5px solid #ffcad4 !important;
    border-radius: 15px;
    padding: 15px 25px;
    max-width: 380px;
    margin: 10px auto;
    box-shadow: 0 0 15px rgba(255, 182, 193, 0.4) !important; 
}

.status-title {
    color: #ff4d6d !important;
    font-weight: 800;
    font-size: 1.1em;
    margin-bottom: 8px;
}

/* Scroll Container untuk Karakter */
.scroll-container {
    max-height: 250px;
    overflow-y: auto;
    border: 2px solid #ffcad4 !important;
    border-radius: 15px;
    padding: 10px;
    background: #ffffff !important;
}

/* Item Karakter: Putih Bersih, Teks Pink */
.char-item {
    display: block !important;
    width: 100% !important;
    background: #ffffff !important;
    border: 1.5px solid #ffb6c1 !important;
    border-radius: 10px !important;
    padding: 10px !important;
    margin-bottom: 8px !important;
    text-align: center !important;
    color: #ff4d6d !important;
    font-weight: 700 !important;
    font-size: 0.95em !important;
}

.char-item:hover {
    border-color: #ff4d6d !important;
    background: #fffafa !important;
}

/* Judul Character Information: Teks Putih */
.char-info-tag {
    background: #ff4d6d !important;
    color: #ffffff !important;
    padding: 6px 20px;
    border-radius: 20px;
    font-weight: bold;
    font-size: 0.95em;
    display: inline-block;
}

/* Footer: Background Putih, Teks Pink */
.footer-box {
    text-align: center;
    padding: 20px;
    margin-top: 25px;
    border: 1.5px solid #ffb6c1 !important;
    border-radius: 20px;
    background: #ffffff !important;
}

.footer-text-main {
    color: #ff4d6d !important;
    margin: 0;
    font-size: 1.1em;
    font-weight: 900;
}

.footer-text-sub {
    color: #ff8fa3 !important;
    margin: 5px 0 0 0;
    font-size: 0.85em;
    font-weight: bold;
}

/* Tombol Generate Magic */
button.primary {
    background: #ff758f !important;
    border: none !important;
    color: white !important;
    font-weight: bold !important;
    border-radius: 12px !important;
    box-shadow: 0 4px 0 #e65a73 !important; /* Efek tombol 3D pink tanpa shadow gelap */
}

button.primary:active {
    box-shadow: none !important;
    transform: translateY(4px);
}
"""

def get_text(text, hps, is_symbol):
    text_norm = text_to_sequence(text, hps.symbols, [] if is_symbol else hps.data.text_cleaners)
    if hps.data.add_blank:
        text_norm = commons.intersperse(text_norm, 0)
    return LongTensor(text_norm)

def create_tts_fn(model, hps, speaker_ids):
    def tts_fn(text, speaker, speed, is_symbol):
        if not text.strip():
            raise gr.Error("Mohon Isi Text Dulu!")
        speaker_id = speaker_ids[speaker]
        stn_tst = get_text(text, hps, is_symbol)
        with no_grad():
            x_tst = stn_tst.unsqueeze(0).to(device)
            x_tst_lengths = LongTensor([stn_tst.size(0)]).to(device)
            sid = LongTensor([speaker_id]).to(device)
            audio = model.infer(x_tst, x_tst_lengths, sid=sid, noise_scale=.667, noise_scale_w=0.8,
                                length_scale=1.0 / speed)[0][0, 0].data.cpu().float().numpy()
        return (hps.data.sampling_rate, audio)
    return tts_fn

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--device', type=str, default='cpu')
    args = parser.parse_args()
    device = torch.device(args.device)
    
    models_tts = []
    total_characters = 0
    config_file = "saved_model/info.json"
    if not os.path.exists(config_file): config_file = "model_assets/configs.json"

    with open(config_file, "r", encoding="utf-8") as f:
        models_info = json.load(f)
        
    for i, info in models_info.items():
        folder_id = "20" if i == "0" else i
        hps = utils.get_hparams_from_file(f"saved_model/{folder_id}/config.json")
        model = SynthesizerTrn(len(hps.symbols), hps.data.filter_length // 2 + 1, hps.train.segment_size // hps.data.hop_length, n_speakers=hps.data.n_speakers, **hps.model)
        utils.load_checkpoint(f"saved_model/{folder_id}/model.pth", model, None)
        model.eval().to(device)
        speakers = [name for name in hps.speakers if name != "None"]
        speaker_ids = {name: i for i, name in enumerate(hps.speakers) if name != "None"}
        total_characters += len(speakers)
        models_tts.append((info["title"], info.get("cover", ""), speakers, info["example"], create_tts_fn(model, hps, speaker_ids)))

    with gr.Blocks(css=custom_css, title="Mahou Shoujo TTS") as app:
        gr.HTML("""<div style="text-align: center; padding: 15px;"><h2 style="color: #ff4d6d; margin:0;">🩷 Mahou Shoujo TTS 🩷</h2></div>""")

        # 1. System Status
        gr.HTML(f"""
        <div class="status-card">
            <div class="status-title">📑 System Status 📑</div>
            <div style="color: #666; font-size: 0.95em; line-height: 1.6;">
                • Model: <b style="color:#ff4d6d">LOADED</b><br>
                • Characters: <b style="color:#ff4d6d">{total_characters}</b>
                <div style="border-top: 1px solid #ffe5ec; margin-top: 10px; padding-top: 5px; color: #ff8fa3; font-size: 0.85em; font-weight: 600;">
                    v3.0 • Ready for generation (dilarang copy)
                </div>
            </div>
        </div>
        """)

        with gr.Tabs() as tabs:
            with gr.TabItem("📑 Magic Generate 📑", id="gen_tab"):
                for name, cover, speakers, example, tts_fn in models_tts:
                    with gr.Row():
                        with gr.Column(scale=1):
                            if cover: gr.Image(f"saved_model/20/{cover}", show_label=False)
                        with gr.Column(scale=2):
                            txt_input = gr.TextArea(label="💬 Masukan Teks (Jepang/Romaji)", placeholder="Masukkan teks...", value=example, lines=3)
                            spk_input = gr.Dropdown(label="🏷️ Pilih Karakter", choices=speakers, value=speakers[0])
                            spd_input = gr.Slider(label="⚡ Kecepatan Suara", minimum=0.5, maximum=2.0, value=1.0, step=0.1)
                            btn_gen = gr.Button("🪄 Berubah!", variant="primary")
                            aud_out = gr.Audio(label="Magic Voice")
                            btn_gen.click(tts_fn, [txt_input, spk_input, spd_input, gr.State(False)], aud_out)

        # 2. Character Information
        gr.HTML("""<div style="text-align: center; margin-top: 25px; margin-bottom: 10px;"><span class="char-info-tag">Character Information</span></div>""")
        
        with gr.Column(elem_classes="scroll-container"):
            for char_name in speakers:
                c_btn = gr.Button(f"🎭 JP_{char_name}", elem_classes="char-item")
                c_btn.click(fn=lambda x=char_name: (gr.Tabs.update(selected="gen_tab"), x), outputs=[tabs, spk_input])

        # 3. Footer
        gr.HTML("""
        <div class="footer-box">
            <p class="footer-text-main">CREATED BY PLANA CHAN</p>
            <p class="footer-text-sub">Mahou Shoujo TTS v3.0 • Powered by VITS</p>
        </div>
        """)

    app.launch()