import os
import json
import traceback
import logging
import gradio as gr
import numpy as np
import librosa
import torch
import asyncio
import edge_tts
import re
import shutil
import time
from datetime import datetime
from fairseq import checkpoint_utils
from fairseq.data.dictionary import Dictionary
from lib.infer_pack.models import (
    SynthesizerTrnMs256NSFsid,
    SynthesizerTrnMs256NSFsid_nono,
    SynthesizerTrnMs768NSFsid,
    SynthesizerTrnMs768NSFsid_nono,
)
from vc_infer_pipeline import VC
from config import Config

# =============================
# LOAD ENVIRONMENT VARIABLES
# =============================
from dotenv import load_dotenv
load_dotenv()

HF_TOKEN = os.getenv("HF_TOKEN")
if HF_TOKEN:
    print("🔑 Hugging Face token detected")
    os.environ["HUGGINGFACE_TOKEN"] = HF_TOKEN
else:
    print("⚠️ No HF_TOKEN found")

# =============================
# AUTO-DOWNLOAD DARI HUGGING FACE - DIPERBAIKI
# =============================
def download_required_weights():
    """Fungsi untuk download model dari Hugging Face"""
    print("=" * 50)
    print("🚀 BLUE ARCHIVE VOICE CONVERSION v2.0")
    print("=" * 50)
    
    target_dir = "weights/Bocchi-the-Rock"
    
    # Cek jika model sudah ada
    if os.path.exists(target_dir):
        print(f"📁 Checking existing models in: {target_dir}")
        model_files = []
        for root, dirs, files in os.walk(target_dir):
            for file in files:
                if file.endswith(".pth"):
                    model_files.append(os.path.join(root, file))
        
        if len(model_files) >= 8:
            print(f"✅ Models already exist: {len(model_files)} .pth files found")
            print("📊 Listing available models:")
            for m in model_files:
                print(f"   - {os.path.basename(m)}")
            return True
        else:
            print(f"⚠️ Incomplete models: {len(model_files)}/8 .pth files found")
    
    try:
        from huggingface_hub import snapshot_download
        
        repo_id = "Plana-Archive/Premium-Model"
        print(f"📥 Downloading from: {repo_id}")
        print("📁 Looking for: Bocchi the Rock! - RCV/weights/Bocchi-the-Rock")
        
        # Download dengan pattern yang lebih spesifik
        downloaded_path = snapshot_download(
            repo_id=repo_id,
            allow_patterns=[
                "Bocchi the Rock! - RCV/weights/Bocchi-the-Rock/**",
                "**/folder_info.json",
                "**/model_info.json"
            ],
            local_dir=".",
            local_dir_use_symlinks=False,
            token=HF_TOKEN,
            max_workers=2
        )
        
        print("✅ Download completed")
        
        # Pindahkan file
        source_dir = "Bocchi the Rock! - RCV/weights/Bocchi-the-Rock"
        
        if os.path.exists(source_dir):
            os.makedirs("weights", exist_ok=True)
            
            if os.path.exists(target_dir):
                print("📦 Removing old weights folder...")
                shutil.rmtree(target_dir)
            
            print(f"📂 Moving models to: {target_dir}")
            shutil.move(source_dir, target_dir)
            
            # Cek isi folder setelah dipindahkan
            print("\n📊 Verifying downloaded models:")
            for root, dirs, files in os.walk(target_dir):
                for dir_name in dirs:
                    dir_path = os.path.join(root, dir_name)
                    pth_files = [f for f in os.listdir(dir_path) if f.endswith('.pth')]
                    index_files = [f for f in os.listdir(dir_path) if f.endswith('.index')]
                    image_files = [f for f in os.listdir(dir_path) if f.endswith(('.png', '.jpg', '.jpeg'))]
                    
                    if pth_files:
                        print(f"   📁 {dir_name}:")
                        print(f"      Model: {pth_files[0] if pth_files else 'NOT FOUND'}")
                        print(f"      Index: {index_files[0] if index_files else 'NOT FOUND'}")
                        print(f"      Cover: {image_files[0] if image_files else 'NOT FOUND'}")
            
            # Hapus folder sumber
            try:
                if os.path.exists("Bocchi the Rock! - RCV"):
                    shutil.rmtree("Bocchi the Rock! - RCV")
            except:
                pass
            
            # Update model_info.json dengan nama file yang sebenarnya
            update_model_info_with_actual_files(target_dir)
            
            return True
        else:
            print("❌ Source directory not found after download!")
            return False
                
    except Exception as e:
        print(f"⚠️ Download failed: {str(e)}")
        print("\n📝 Manual setup:")
        print("1. Create folder: weights/Bocchi-the-Rock/")
        print("2. Download from: https://huggingface.co/Plana-Archive/Anime-RCV")
        print("3. Look for: Bocchi the Rock! - RCV/weights/Bocchi-the-Rock")
        print("4. Put each character in their own folder")
    
    return False

def update_model_info_with_actual_files(target_dir):
    """Update model_info.json dengan nama file yang sebenarnya ada"""
    model_info_path = os.path.join(target_dir, "model_info.json")
    
    if not os.path.exists(model_info_path):
        print("⚠️ model_info.json not found, creating default...")
        # Buat model_info.json berdasarkan file yang ada
        model_info = {}
        for char_dir in os.listdir(target_dir):
            char_path = os.path.join(target_dir, char_dir)
            if os.path.isdir(char_path):
                pth_files = [f for f in os.listdir(char_path) if f.endswith('.pth')]
                index_files = [f for f in os.listdir(char_path) if f.endswith('.index')]
                image_files = [f for f in os.listdir(char_path) if f.endswith(('.png', '.jpg', '.jpeg'))]
                
                if pth_files:
                    model_info[char_dir] = {
                        "enable": True,
                        "model_path": pth_files[0],
                        "title": f"Bocchi the Rock! - {char_dir.replace('-', ' ')}",
                        "cover": image_files[0] if image_files else "cover.png",
                        "feature_retrieval_library": index_files[0] if index_files else f"{char_dir}.index",
                        "author": "Plana-Archive"
                    }
        
        with open(model_info_path, "w", encoding="utf-8") as f:
            json.dump(model_info, f, indent=2, ensure_ascii=False)
        print(f"✅ Created model_info.json with {len(model_info)} characters")
    else:
        print("📄 Found model_info.json, checking for file mismatches...")
        try:
            with open(model_info_path, "r", encoding="utf-8") as f:
                model_info = json.load(f)
            
            updated = False
            for char_name, info in model_info.items():
                if not info.get('enable', True):
                    continue
                
                char_path = os.path.join(target_dir, char_name)
                if os.path.exists(char_path):
                    # Cek apakah file model ada
                    expected_model = info.get('model_path')
                    actual_models = [f for f in os.listdir(char_path) if f.endswith('.pth')]
                    
                    if expected_model not in actual_models and actual_models:
                        print(f"   🔄 Updating {char_name}: {expected_model} → {actual_models[0]}")
                        info['model_path'] = actual_models[0]
                        updated = True
                    
                    # Cek index file
                    expected_index = info.get('feature_retrieval_library')
                    actual_indices = [f for f in os.listdir(char_path) if f.endswith('.index')]
                    
                    if expected_index not in actual_indices and actual_indices:
                        print(f"   🔄 Updating {char_name} index: {expected_index} → {actual_indices[0]}")
                        info['feature_retrieval_library'] = actual_indices[0]
                        updated = True
            
            if updated:
                with open(model_info_path, "w", encoding="utf-8") as f:
                    json.dump(model_info, f, indent=2, ensure_ascii=False)
                print("✅ Updated model_info.json with actual file names")
        except Exception as e:
            print(f"⚠️ Error updating model_info.json: {str(e)}")

# Jalankan download
download_required_weights()

# Inisialisasi konfigurasi
config = Config()
logging.getLogger("numba").setLevel(logging.WARNING)
logging.getLogger("fairseq").setLevel(logging.WARNING)

# Cache untuk model
model_cache = {}
hubert_loaded = False
hubert_model = None

# Mode audio
spaces = True
if spaces:
    audio_mode = ["Upload audio", "TTS Audio"]
else:
    audio_mode = ["Input path", "Upload audio", "TTS Audio"]

# Metode F0 extraction
f0method_mode = ["pm", "harvest"]
if os.path.isfile("rmvpe.pt"):
    f0method_mode.insert(2, "rmvpe")

def clean_title(title):
    """Membersihkan judul model"""
    title = re.sub(r'^Blue Archive\s*-\s*', '', title, flags=re.IGNORECASE)
    title = re.sub(r'^Bocchi the Rock!\s*-\s*', '', title, flags=re.IGNORECASE)
    return re.sub(r'\s*-\s*\d+\s*epochs', '', title, flags=re.IGNORECASE)

def _load_audio_input(vc_audio_mode, vc_input, vc_upload, tts_text, spaces_limit=20):
    """Memuat audio dari berbagai sumber"""
    temp_file = None
    try:
        if vc_audio_mode == "Input path" and vc_input:
            audio, sr = librosa.load(vc_input, sr=16000, mono=True)
            return audio.astype(np.float32), 16000, None
            
        elif vc_audio_mode == "Upload audio":
            if vc_upload is None:
                raise ValueError("Please upload an audio file!")
            sampling_rate, audio = vc_upload
            
            if audio.dtype != np.float32:
                audio = audio.astype(np.float32) / np.iinfo(audio.dtype).max
            
            if len(audio.shape) > 1:
                audio = np.mean(audio, axis=0)
                
            if sampling_rate != 16000:
                audio = librosa.resample(audio, orig_sr=sampling_rate, target_sr=16000, res_type='kaiser_fast')
                
            return audio.astype(np.float32), 16000, None
            
        elif vc_audio_mode == "TTS Audio":
            if not tts_text or tts_text.strip() == "":
                raise ValueError("Please enter text for TTS!")
            
            temp_file = f"tts_temp_{int(time.time())}.wav"
            
            async def tts_task():
                return await edge_tts.Communicate(tts_text, "ja-JP-NanamiNeural").save(temp_file)
            
            try:
                asyncio.run(asyncio.wait_for(tts_task(), timeout=15))
            except asyncio.TimeoutError:
                raise ValueError("TTS timeout!")
            
            audio, sr = librosa.load(temp_file, sr=16000, mono=True)
            return audio.astype(np.float32), 16000, temp_file
            
    except Exception as e:
        if temp_file and os.path.exists(temp_file):
            os.remove(temp_file)
        raise e
        
    raise ValueError("Invalid audio mode")

def adjust_audio_speed(audio, speed):
    """Menyesuaikan kecepatan audio"""
    if speed == 1.0:
        return audio
    return librosa.effects.time_stretch(audio.astype(np.float32), rate=speed)

def preprocess_audio(audio):
    """Preprocessing audio"""
    if np.max(np.abs(audio)) > 1.0:
        audio = audio / np.max(np.abs(audio)) * 0.9
    return audio.astype(np.float32)

def create_vc_fn(model_key, tgt_sr, net_g, vc, if_f0, version, file_index):
    """Membuat fungsi konversi voice"""
    def vc_fn(
        vc_audio_mode, vc_input, vc_upload, tts_text,
        f0_up_key, f0_method, index_rate, filter_radius, 
        resample_sr, rms_mix_rate, protect, speed,
    ):
        temp_audio_file = None
        try:
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
                
            net_g.to(config.device)
            
            yield "Status: 🚀 Processing audio...", None
            
            audio, sr, temp_audio_file = _load_audio_input(vc_audio_mode, vc_input, vc_upload, tts_text)
            audio = preprocess_audio(audio)
            audio_tensor = torch.FloatTensor(audio).to(config.device)
            
            times = [0, 0, 0]
            max_chunk_size = 16000 * 30
            
            if len(audio) > max_chunk_size:
                chunks = []
                for i in range(0, len(audio), max_chunk_size):
                    chunk = audio[i:i + max_chunk_size]
                    chunk_tensor = torch.FloatTensor(chunk).to(config.device)
                    
                    chunk_opt = vc.pipeline(
                        hubert_model, net_g, 0, chunk_tensor, 
                        "chunk" if vc_input else "temp", times, 
                        int(f0_up_key), f0_method, file_index, index_rate,
                        if_f0, filter_radius, tgt_sr, resample_sr,
                        rms_mix_rate, version, protect, f0_file=None,
                    )
                    chunks.append(chunk_opt)
                
                audio_opt = np.concatenate(chunks)
            else:
                audio_opt = vc.pipeline(
                    hubert_model, net_g, 0, audio_tensor, 
                    vc_input if vc_input else "temp", times, 
                    int(f0_up_key), f0_method, file_index, index_rate,
                    if_f0, filter_radius, tgt_sr, resample_sr,
                    rms_mix_rate, version, protect, f0_file=None,
                )
            
            audio_opt = audio_opt.astype(np.float32)
            
            if speed != 1.0:
                audio_opt = adjust_audio_speed(audio_opt, speed)
            
            if np.max(np.abs(audio_opt)) > 0:
                audio_opt = (audio_opt / np.max(np.abs(audio_opt)) * 0.9).astype(np.float32)
            
            yield "Status: ✅ Conversion completed!", (tgt_sr, audio_opt)
            
        except Exception as e:
            yield f"❌ Error: {str(e)}", None
        finally:
            if temp_audio_file and os.path.exists(temp_audio_file):
                os.remove(temp_audio_file)
            
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
                
            if model_key not in model_cache:
                net_g.to('cpu')
    
    return vc_fn

def load_model():
    """Memuat semua model"""
    print("\n" + "=" * 50)
    print("🎵 LOADING VOICE MODELS")
    print("=" * 50)
    
    categories = []
    base_path = "weights"
    
    if not os.path.exists(base_path):
        print(f"❌ Folder '{base_path}' not found!")
        return categories
    
    # Buat folder_info.json jika tidak ada
    folder_info_path = f"{base_path}/folder_info.json"
    if not os.path.isfile(folder_info_path):
        print(f"📄 Creating {folder_info_path}...")
        default_folder_info = {
            "BocchiTheRock": {
                "title": "Bocchi the Rock! - RCV Collection",
                "folder_path": "Bocchi-the-Rock",
                "description": "Official RVC Weights for Bocchi the Rock! characters",
                "enable": True
            }
        }
        
        with open(folder_info_path, "w", encoding="utf-8") as f:
            json.dump(default_folder_info, f, indent=2, ensure_ascii=False)
        print(f"✅ Created {folder_info_path}")
    
    with open(folder_info_path, "r", encoding="utf-8") as f:
        folder_info = json.load(f)
    
    print(f"📁 Found {len(folder_info)} category(ies) in folder_info.json")
    
    for category_name, category_info in folder_info.items():
        if not category_info.get('enable', True): 
            continue
            
        category_title = category_info['title']
        category_folder = category_info['folder_path']
        description = category_info['description']
        
        category_folder = os.path.basename(category_folder)
        
        models = []
        model_info_path = f"{base_path}/{category_folder}/model_info.json"
        
        print(f"\n📂 Loading category: {category_title}")
        print(f"   Path: {model_info_path}")
        
        if os.path.exists(model_info_path):
            with open(model_info_path, "r", encoding="utf-8") as f:
                models_info = json.load(f)
            
            print(f"   Found {len(models_info)} character(s) in model_info.json")
            
            for character_name, info in models_info.items():
                if not info.get('enable', True): 
                    continue
                
                model_title = info['title']
                model_name = info['model_path']
                model_author = info.get("author", "Plana-Archive")
                
                cache_key = f"{category_folder}_{character_name}"
                
                char_dir = f"{base_path}/{category_folder}/{character_name}"
                model_path = f"{char_dir}/{model_name}"
                cover_path = f"{char_dir}/{info['cover']}"
                index_path = f"{char_dir}/{info['feature_retrieval_library']}"
                
                print(f"\n   👤 Character: {character_name}")
                print(f"      Expected model: {model_name}")
                print(f"      Expected cover: {info['cover']}")
                print(f"      Expected index: {info['feature_retrieval_library']}")
                
                # Cek file yang diperlukan
                required_files = [model_path, cover_path, index_path]
                missing_files = [f for f in required_files if not os.path.exists(f)]
                
                if missing_files:
                    print(f"      ⚠️ Missing files:")
                    for f in missing_files:
                        print(f"         - {os.path.basename(f)}")
                    
                    # Coba cari file alternatif
                    if os.path.exists(char_dir):
                        actual_files = os.listdir(char_dir)
                        print(f"      📁 Actual files in directory:")
                        for f in actual_files:
                            print(f"         - {f}")
                        
                        # Cari file .pth
                        pth_files = [f for f in actual_files if f.endswith('.pth')]
                        if pth_files and not os.path.exists(model_path):
                            print(f"      🔄 Found alternative model: {pth_files[0]}")
                            model_name = pth_files[0]
                            model_path = f"{char_dir}/{model_name}"
                        
                        # Cari file index
                        index_files = [f for f in actual_files if f.endswith('.index')]
                        if index_files and not os.path.exists(index_path):
                            print(f"      🔄 Found alternative index: {index_files[0]}")
                            index_path = f"{char_dir}/{index_files[0]}"
                        
                        # Cari cover
                        image_files = [f for f in actual_files if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
                        if image_files and not os.path.exists(cover_path):
                            print(f"      🔄 Found alternative cover: {image_files[0]}")
                            cover_path = f"{char_dir}/{image_files[0]}"
                    
                    # Cek ulang setelah mencari alternatif
                    required_files = [model_path, cover_path, index_path]
                    missing_files = [f for f in required_files if not os.path.exists(f)]
                    
                    if missing_files:
                        print(f"      ❌ Skipping {character_name} - still missing files")
                        continue
                
                # Gunakan cache jika tersedia
                if cache_key in model_cache:
                    tgt_sr, net_g, vc, if_f0, version, model_index = model_cache[cache_key]
                    print(f"      ✅ Loaded from cache")
                else:
                    try:
                        print(f"      ⏳ Loading model weights...")
                        
                        cpt = torch.load(model_path, map_location="cpu")
                        tgt_sr = cpt["config"][-1]
                        cpt["config"][-3] = cpt["weight"]["emb_g.weight"].shape[0]
                        if_f0 = cpt.get("f0", 1)
                        version = cpt.get("version", "v1")
                        
                        if version == "v1":
                            if if_f0 == 1:
                                net_g = SynthesizerTrnMs256NSFsid(*cpt["config"], is_half=config.is_half)
                            else:
                                net_g = SynthesizerTrnMs256NSFsid_nono(*cpt["config"])
                        else:
                            if if_f0 == 1:
                                net_g = SynthesizerTrnMs768NSFsid(*cpt["config"], is_half=config.is_half)
                            else:
                                net_g = SynthesizerTrnMs768NSFsid_nono(*cpt["config"])
                        
                        if hasattr(net_g, "enc_q"): 
                            del net_g.enc_q
                        net_g.load_state_dict(cpt["weight"], strict=False)
                        net_g.eval().to('cpu')
                        
                        vc = VC(tgt_sr, config)
                        model_cache[cache_key] = (tgt_sr, net_g, vc, if_f0, version, index_path)
                        
                        print(f"      ✅ Model loaded successfully (v{version}, SR: {tgt_sr})")
                        
                    except Exception as e:
                        print(f"      ❌ Error loading model: {str(e)}")
                        continue
                
                models.append((
                    character_name, 
                    model_title, 
                    model_author, 
                    cover_path, 
                    version, 
                    create_vc_fn(cache_key, tgt_sr, net_g, vc, if_f0, version, index_path)
                ))
        else:
            print(f"   ⚠️ model_info.json not found at: {model_info_path}")
        
        if models:
            categories.append([category_title, category_folder, description, models])
            print(f"\n   📊 Category '{category_title}' loaded with {len(models)} model(s)")
        else:
            print(f"\n   ⚠️ No models loaded for category '{category_title}'")
    
    total_models = sum(len(models) for _, _, _, models in categories)
    print(f"\n🎯 Total categories loaded: {len(categories)}")
    print(f"👥 Total models loaded: {total_models}")
    print("=" * 50)
    
    return categories

def load_hubert():
    """Memuat model HuBERT"""
    global hubert_model, hubert_loaded
    if hubert_loaded:
        return
        
    print("🔧 Loading HuBERT model...")
    torch.serialization.add_safe_globals([Dictionary])
    models, _, _ = checkpoint_utils.load_model_ensemble_and_task(
        ["hubert_base.pt"], 
        suffix="",
    )
    hubert_model = models[0].to(config.device)
    hubert_model = hubert_model.half() if config.is_half else hubert_model.float()
    hubert_model.eval()
    hubert_loaded = True
    print("✅ HuBERT model loaded successfully")

def change_audio_mode(vc_audio_mode):
    """Mengubah tampilan input audio"""
    is_input_path = vc_audio_mode == "Input path"
    is_upload = vc_audio_mode == "Upload audio"
    is_tts = vc_audio_mode == "TTS Audio"
    
    return (
        gr.Textbox.update(visible=is_input_path),
        gr.Checkbox.update(visible=is_upload),
        gr.Audio.update(visible=is_upload),
        gr.Textbox.update(visible=is_tts, lines=4 if is_tts else 2)
    )

def use_microphone(microphone):
    """Toggle microphone/upload source"""
    return gr.Audio.update(source="microphone" if microphone else "upload")

# CSS dengan tema PINK
css = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&family=Quicksand:wght@400;600;700&display=swap');
body, .gradio-container { background-color: #ffffff !important; font-family: 'Inter', sans-serif !important; }
footer { display: none !important; }
.arona-loading-container { display: flex; align-items: center; justify-content: center; gap: 15px; margin-top: 15px; padding: 10px; }
.loading-text-pink { font-family: 'Quicksand', sans-serif; font-size: 20px; font-weight: 700; color: #ff69b4; letter-spacing: 1px; }
.loading-gif-small { width: 100px; height: auto; border-radius: 8px; }
.header-img-container { text-align: center; padding: 10px 0; background: #ffffff !important; }
.header-img { width: 100%; max-width: 500px; border-radius: 15px; margin: 0 auto; display: block; }
.status-card { background: #ffffff; border: 1px solid #ffe4ec; border-radius: 14px; padding: 15px 10px; margin: 0 auto 15px auto; max-width: 400px; display: flex; flex-direction: column; align-items: center; }
.status-online-box { display: flex; align-items: center; gap: 8px; margin-bottom: 12px; }
.status-details-container { display: flex; width: 100%; justify-content: center; align-items: center; border-top: 1px solid #fff0f7; padding-top: 10px; }
.status-detail-item { flex: 1; display: flex; flex-direction: column; align-items: center; text-align: center; }
.status-detail-item:first-child { border-right: 1px solid #ffe4ec; }
.status-text-main { font-size: 13px !important; font-weight: 600; color: #7b4d5a; }
.status-text-sub { font-size: 11px !important; color: #b07d8b; }
.dot-online { height: 8px; width: 8px; background-color: #ff69b4; border-radius: 50%; display: inline-block; animation: blink-pink 1.5s infinite; }
@keyframes blink-pink { 0% { opacity: 1; } 50% { opacity: 0.4; } 100% { opacity: 1; } }
.gr-form .gr-block label span, .gr-box label span, .gr-panel label span { background: linear-gradient(135deg, #ff69b4 0%, #ff1493 100%) !important; color: white !important; padding: 4px 12px !important; border-radius: 8px !important; font-weight: 600 !important; box-shadow: 0 0 15px rgba(255, 105, 180, 0.4) !important; }
input[type="range"] { accent-color: #ff69b4 !important; }
.char-scroll-box { display: grid !important; grid-template-columns: repeat(2, 1fr) !important; gap: 12px !important; max-height: 280px; overflow-y: auto; padding: 15px; background: #ffffff; border: 2px solid #ffeef4; border-radius: 14px; }
.char-card { background: white; padding: 12px; border-radius: 12px; cursor: pointer; border: 1px solid #ffe4ec; border-left: 5px solid #ff69b4; transition: all 0.2s ease; display: flex; flex-direction: column; height: 65px; }
.char-card:hover { transform: translateY(-3px); box-shadow: 0 5px 15px rgba(255, 105, 180, 0.2); border-left-color: #ff1493; }
.char-name-jp { font-weight: 700; font-size: 11px !important; color: #7b4d5a; }
.char-name-en { font-size: 8.5px !important; color: #b07d8b; text-transform: uppercase; }
.speed-section { margin-top: 20px; padding: 18px; border-radius: 20px; background: linear-gradient(135deg, #fff0f7 0%, #ffffff 100%); border: 2px solid #ffe4ec; }
.speed-title { font-family: 'Quicksand', sans-serif; font-weight: 700; color: #ff69b4; text-align: center; margin-bottom: 12px; font-size: 14px; }
.generate-btn { font-family: 'Quicksand', sans-serif; font-weight: 700 !important; background: linear-gradient(135deg, #ff69b4 0%, #ff1493 100%) !important; color: white !important; border-radius: 12px !important; padding: 12px 24px !important; transition: all 0.3s ease !important; }
.generate-btn:hover { transform: scale(1.05); box-shadow: 0 5px 20px rgba(255, 20, 147, 0.3) !important; }
.footer-text { text-align: center; padding: 20px; border-top: 1px solid #f8f0f4; color: #b07d8b; font-size: 11px; }
.speed-notes-box { font-family: 'Arial'; border: 1px solid #ffd1dc; border-radius: 8px; padding: 12px; background: #fff5f8; border-left: 4px solid #ff69b4; margin-top: 10px; }
.speed-notes-title { color: #ff1493; font-size: 12px; margin: 0 0 5px 0; font-weight: bold; }
.speed-notes-content { color: #d81b60; font-size: 11px; margin: 0; }
.model-tab { background: linear-gradient(135deg, #fff8fb 0%, #ffffff 100%) !important; border-radius: 15px !important; padding: 15px !important; }
.advanced-settings { background: #f9f9f9 !important; border-radius: 10px !important; padding: 15px !important; border: 1px solid #e0e0e0 !important; }
.error-box { background: #ffebee; border: 1px solid #ffcdd2; border-radius: 8px; padding: 15px; margin: 10px 0; color: #c62828; }
.info-box { background: #fce4ec; border: 1px solid #f8bbd9; border-radius: 8px; padding: 15px; margin: 10px 0; color: #ad1457; }
"""

if __name__ == '__main__':
    # Preload HuBERT
    load_hubert()
    
    # Load models
    categories = load_model()
    total_models = sum(len(models) for _, _, _, models in categories)
    
    # UI dengan Gradio
    with gr.Blocks(css=css, theme=gr.themes.Soft(primary_hue="pink")) as app:
        gr.HTML('<div class="header-img-container"><img src="https://huggingface.co/spaces/Library-Anime/Bocchi-the-Rock/resolve/main/Bocchi-the-Rock.PNG" class="header-img"></div>')
        
        # Status card
        if total_models > 0:
            gr.HTML(f'''
                <div class="status-card">
                    <div class="status-online-box">
                        <span class="dot-online"></span>
                        <b style="color: #ff69b4; font-size: 14px;">Voice Conversion System Online</b>
                    </div>
                    <div class="status-details-container">
                        <div class="status-detail-item">
                            <span class="status-text-main">👥 {total_models} Students</span>
                            <span class="status-text-sub">Ready for Conversion</span>
                        </div>
                        <div class="status-detail-item">
                            <span class="status-text-main">📊 Total Models</span>
                            <span class="status-text-sub">Database: {total_models}</span>
                        </div>
                    </div>
                </div>
            ''')
        else:
            gr.HTML(f'''
                <div class="error-box">
                    <h3>⚠️ No Models Loaded</h3>
                    <p>Please check console logs for details.</p>
                    <p>Download from: <a href="https://huggingface.co/Plana-Archive/Anime-RCV" target="_blank">https://huggingface.co/Plana-Archive/Anime-RCV</a></p>
                </div>
            ''')
        
        # Tabs untuk setiap kategori
        if categories:
            for cat_idx, (folder_title, folder, description, models) in enumerate(categories):
                with gr.TabItem(folder_title, elem_classes="model-tab"):
                    with gr.Accordion("📑 Select Student Voice", open=True):
                        char_html = "".join([
                            f'<div class="char-card" onclick="selectModel(\'{folder_title}\', \'{name}\')">'
                            f'<span class="char-name-jp">{clean_title(title)}</span>'
                            f'<span class="char-name-en">{name}</span>'
                            f'</div>' 
                            for name, title, author, cover, version, vc_fn in models
                        ])
                        gr.HTML(f'<div class="char-scroll-box">{char_html}</div>')
                    
                    # Tabs untuk setiap model
                    with gr.Tabs():
                        for model_idx, (name, title, author, cover, model_version, vc_fn) in enumerate(models):
                            with gr.TabItem(name, id=f"model_{cat_idx}_{model_idx}"):
                                with gr.Row():
                                    # Kolom kiri: Model info
                                    with gr.Column(scale=1):
                                        gr.HTML(f'''
                                            <div style="display:flex;flex-direction:column;align-items:center;padding:20px;background:white;border-radius:20px;border:1px solid #ffeef4;">
                                                <img style="width:200px;height:260px;object-fit:cover;border-radius:15px;" src="file/{cover}">
                                                <div style="font-family:'Quicksand',sans-serif;font-weight:700;font-size:18px;color:#ff1493;margin-top:15px;">
                                                    {clean_title(title)}
                                                </div>
                                                <div style="font-size:11px;color:#b07d8b;margin-top:5px;">
                                                    {model_version} • {author}
                                                </div>
                                            </div>
                                        ''')
                                    
                                    # Kolom tengah: Input dan settings
                                    with gr.Column(scale=2):
                                        # Input group
                                        with gr.Group():
                                            vc_audio_mode = gr.Dropdown(
                                                label="Input Mode", 
                                                choices=audio_mode, 
                                                value="TTS Audio"
                                            )
                                            vc_input = gr.Textbox(visible=False)
                                            vc_microphone_mode = gr.Checkbox(
                                                label="Use Microphone", 
                                                value=False
                                            )
                                            vc_upload = gr.Audio(
                                                label="Upload Audio Source", 
                                                source="upload", 
                                                visible=False,
                                                type="numpy"
                                            )
                                            tts_text = gr.Textbox(
                                                label="TTS Text", 
                                                visible=True, 
                                                placeholder="Type your message here...",
                                                lines=4
                                            )
                                        
                                        # Basic settings
                                        with gr.Row():
                                            with gr.Column():
                                                vc_transform0 = gr.Slider(
                                                    minimum=-12, 
                                                    maximum=12, 
                                                    label="Pitch", 
                                                    value=12, 
                                                    step=1
                                                )
                                                f0method0 = gr.Radio(
                                                    label="Conversion Algorithm", 
                                                    choices=f0method_mode, 
                                                    value="rmvpe" if "rmvpe" in f0method_mode else "pm"
                                                )
                                            with gr.Column():
                                                with gr.Accordion("⚙️ Advanced Tuning", open=True, elem_classes="advanced-settings"):
                                                    index_rate1 = gr.Slider(
                                                        0, 1, 
                                                        label="Index Rate", 
                                                        value=0.75
                                                    )
                                                    filter_radius0 = gr.Slider(
                                                        0, 7, 
                                                        label="Filter Radius", 
                                                        value=7, 
                                                        step=1
                                                    )
                                                    resample_sr0 = gr.Slider(
                                                        0, 48000, 
                                                        label="Resample SR", 
                                                        value=0
                                                    )
                                                    rms_mix_rate0 = gr.Slider(
                                                        0, 1, 
                                                        label="Volume Mix", 
                                                        value=0.76
                                                    )
                                                    protect0 = gr.Slider(
                                                        0, 0.5, 
                                                        label="Voice Protect", 
                                                        value=0.33
                                                    )
                                        
                                        # Notes
                                        with gr.Row():
                                            with gr.Column():
                                                gr.HTML("""
                                                    <div style="font-family: 'Arial'; border: 1px solid #ffd1e0; border-radius: 8px; padding: 12px; background: #fff5f9; border-left: 4px solid #ff69b4; margin-bottom: 8px;">
                                                        <h4 style="color: #ff1493; font-size: 13px; margin: 0 0 5px 0;">📝 Notes & Guide</h4>
                                                        <p style="color: #d81b60; font-size: 11px; margin: 0 0 3px 0;"><b>Pitch:</b> Adjust voice pitch</p>
                                                        <p style="color: #d81b60; font-size: 11px; margin: 0 0 3px 0;"><b>Algorithm:</b> F0 extraction method</p>
                                                        <p style="color: #d81b60; font-size: 11px; margin: 0 0 3px 0;"><b>Retrieval:</b> Voice similarity (0-1)</p>
                                                        <p style="color: #d81b60; font-size: 11px; margin: 0 0 3px 0;"><b>Filter:</b> Noise reduction</p>
                                                        <p style="color: #d81b60; font-size: 11px; margin: 0 0 3px 0;"><b>Volume:</b> Volume stability</p>
                                                        <p style="color: #d81b60; font-size: 11px; margin: 0;"><b>Protect:</b> Protect voice</p>
                                                    </div>
                                                """)
                                            with gr.Column():
                                                gr.HTML("""
                                                    <div style="font-family: 'Arial'; border: 1px solid #ffd6e7; border-radius: 8px; padding: 12px; background: #fff0f7; border-left: 4px solid #ff69b4;">
                                                        <h4 style="color: #ff1493; font-size: 13px; margin: 0 0 5px 0;">📑 RECOMMENDED</h4>
                                                        <p style="color: #d81b60; font-size: 11px; margin: 0 0 3px 0;"><b>Pitch:</b> <span style="color: #ff1493; font-weight: bold;">+12</span></p>
                                                        <p style="color: #d81b60; font-size: 11px; margin: 0 0 3px 0;"><b>Algorithm:</b> <span style="color: #ff1493; font-weight: bold;">RMVPE</span></p>
                                                        <p style="color: #d81b60; font-size: 11px; margin: 0 0 3px 0;"><b>Retrieval:</b> <span style="color: #ff1493; font-weight: bold;">0.75</span></p>
                                                        <p style="color: #d81b60; font-size: 11px; margin: 0 0 3px 0;"><b>Filter:</b> <span style="color: #ff1493; font-weight: bold;">7</span></p>
                                                        <p style="color: #d81b60; font-size: 11px; margin: 0 0 3px 0;"><b>Volume:</b> <span style="color: #ff1493; font-weight: bold;">0.76</span></p>
                                                        <p style="color: #d81b60; font-size: 11px; margin: 0;"><b>Protect:</b> <span style="color: #ff1493; font-weight: bold;">0.33</span></p>
                                                    </div>
                                                """)
                                        
                                        # Speed section
                                        with gr.Column(elem_classes="speed-section"):
                                            gr.HTML('<div class="speed-title">⚡ VOICE SPEED CONTROL ⚡</div>')
                                            speed_slider = gr.Slider(
                                                0.5, 2.0, 
                                                value=1.0, 
                                                step=0.1, 
                                                label="Speed"
                                            )
                                            
                                            gr.HTML("""
                                                <div class="speed-notes-box">
                                                    <div class="speed-notes-title">ℹ️ Speed Guide</div>
                                                    <div class="speed-notes-content">
                                                        • <b>Left (0.5):</b> Slow down voice<br>
                                                        • <b>Center (1.0):</b> Normal speed<br>
                                                        • <b>Right (2.0):</b> Speed up voice<br>
                                                    </div>
                                                </div>
                                            """)
                                            
                                            # Loading indicator
                                            gr.HTML(
                                                '<div class="arona-loading-container">'
                                                '<div class="loading-text-pink">Ready to Generate!</div>'
                                                '<img class="loading-gif-small" src="https://huggingface.co/spaces/Library-Anime/Bocchi-the-Rock/resolve/main/Bocchi Chan.gif">'
                                                '</div>'
                                            )
                                    
                                    # Kolom kanan: Output
                                    with gr.Column(scale=1):
                                        vc_log = gr.Textbox(
                                            label="Process Logs", 
                                            interactive=False,
                                            lines=4
                                        )
                                        vc_output = gr.Audio(
                                            label="Result Audio", 
                                            interactive=False,
                                            type="numpy"
                                        )
                                        vc_convert = gr.Button(
                                            "🎸 GENERATE VOICE 🎸", 
                                            variant="primary", 
                                            elem_classes="generate-btn",
                                            size="lg"
                                        )
                                
                                # Connect button click
                                vc_convert.click(
                                    fn=vc_fn, 
                                    inputs=[
                                        vc_audio_mode, vc_input, vc_upload, tts_text, 
                                        vc_transform0, f0method0, index_rate1, filter_radius0, 
                                        resample_sr0, rms_mix_rate0, protect0, speed_slider
                                    ], 
                                    outputs=[vc_log, vc_output]
                                )
                                
                                # Connect audio mode change
                                vc_audio_mode.change(
                                    fn=change_audio_mode, 
                                    inputs=[vc_audio_mode], 
                                    outputs=[vc_input, vc_microphone_mode, vc_upload, tts_text]
                                )
                                
                                # Connect microphone toggle
                                vc_microphone_mode.change(
                                    fn=use_microphone, 
                                    inputs=vc_microphone_mode, 
                                    outputs=vc_upload
                                )
        
        # Footer
        gr.HTML(
            '<div class="footer-text">'
            '<div>DESIGNED BY ☘️Mutsumi-Chan☘️</div>'
            '<div style="font-weight:700; color:#b07d8b;">Bocchi the Rock - RCV v1.0 • Pink Edition</div>'
            '</div>'
        )
        
        # JavaScript untuk model selection
        app.load(
            None, None, None, 
            js="""
            () => { 
                window.selectModel = (cat, mod) => { 
                    const tabs = document.querySelectorAll('.tabs .tab-nav button'); 
                    for (let t of tabs) { 
                        if (t.textContent.trim() === cat) { 
                            t.click(); 
                            setTimeout(() => { 
                                const mTabs = document.querySelectorAll('.tabs .tab-nav button'); 
                                for (let mt of mTabs) { 
                                    if (mt.textContent.trim() === mod) {
                                        mt.click(); 
                                        window.scrollTo({top: 0, behavior: 'smooth'});
                                    }
                                } 
                            }, 100); 
                            break; 
                        } 
                    } 
                } 
            }
            """
        )
    
    # Launch app
    print("\n" + "=" * 50)
    print("🌐 STARTING WEB INTERFACE")
    print("=" * 50)
    
    app.queue(max_size=3).launch(
        share=False,
        server_name="0.0.0.0" if os.getenv('SPACE_ID') else "127.0.0.1",
        server_port=7860,
        quiet=False,
        show_error=True
    )