import os
import glob
import json
import traceback
import logging
import gradio as gr
import numpy as np
import librosa
import torch
import asyncio
import edge_tts
import yt_dlp
import ffmpeg
import subprocess
import sys
import io
import wave
import gc
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor
from functools import lru_cache
import psutil
from fairseq import checkpoint_utils
from lib.infer_pack.models import (
    SynthesizerTrnMs256NSFsid,
    SynthesizerTrnMs256NSFsid_nono,
    SynthesizerTrnMs768NSFsid,
    SynthesizerTrnMs768NSFsid_nono,
)
from vc_infer_pipeline import VC
from config import Config

config = Config()
logging.getLogger("numba").setLevel(logging.WARNING)
limitation = os.getenv("SYSTEM") == "spaces"

# ==============================
# OPTIMIZATION SETTINGS
# ==============================
MAX_WORKERS = 1  # Karena hanya 1 core
CHUNK_SIZE = 10 * 16000  # 10 detik pada 16kHz (optimal untuk single core)
OVERLAP_SIZE = int(0.5 * 16000)  # 0.5 detik overlap untuk menghindari artifacts
MEMORY_THRESHOLD = 85  # Clear cache jika memory usage > 85%
USE_HALF_PRECISION = True  # Gunakan float16 jika memungkinkan
CACHE_MODELS = True  # Cache model yang sudah dimuat

# ==============================
# MEMORY MANAGEMENT
# ==============================
class MemoryManager:
    @staticmethod
    def check_memory():
        """Cek penggunaan memori"""
        memory = psutil.virtual_memory()
        return memory.percent
    
    @staticmethod
    def clear_cache():
        """Bersihkan cache memori"""
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        gc.collect()
    
    @staticmethod
    def auto_clear():
        """Auto-clear jika memori tinggi"""
        if MemoryManager.check_memory() > MEMORY_THRESHOLD:
            print(f"Memory high ({MemoryManager.check_memory()}%), clearing cache...")
            MemoryManager.clear_cache()

# ==============================
# CHUNK PROCESSING SYSTEM
# ==============================
class ChunkProcessor:
    def __init__(self, vc_pipeline, chunk_size=CHUNK_SIZE, overlap=OVERLAP_SIZE):
        self.vc_pipeline = vc_pipeline
        self.chunk_size = chunk_size
        self.overlap = overlap
    
    def split_into_chunks(self, audio):
        """Split audio menjadi chunks dengan overlap"""
        audio_len = len(audio)
        chunks = []
        
        if audio_len <= self.chunk_size:
            return [audio], [0]
        
        start = 0
        while start < audio_len:
            end = min(start + self.chunk_size, audio_len)
            chunk = audio[start:end]
            chunks.append(chunk)
            start += self.chunk_size - self.overlap
        
        return chunks, list(range(0, audio_len, self.chunk_size - self.overlap))
    
    def process_chunk(self, chunk_data):
        """Process single chunk"""
        chunk, index, params = chunk_data
        
        try:
            # Ekstrak parameters
            (hubert_model, net_g, sid, vc_input, times, f0_up_key, 
             f0_method, file_index, index_rate, if_f0, filter_radius,
             tgt_sr, resample_sr, rms_mix_rate, version, protect) = params
            
            # Process chunk
            chunk_audio = self.vc_pipeline(
                hubert_model,
                net_g,
                sid,
                chunk,
                vc_input,
                times,
                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,
            )
            
            return index, chunk_audio
        except Exception as e:
            print(f"Error processing chunk {index}: {e}")
            return index, None
    
    def process_parallel(self, audio, params, max_workers=MAX_WORKERS):
        """Process chunks dengan parallel execution"""
        chunks, indices = self.split_into_chunks(audio)
        
        if len(chunks) == 1:
            # Direct processing untuk audio pendek
            return self.vc_pipeline(params[0], params[1], params[2], audio, 
                                   *params[3:])
        
        # Prepare chunk data
        chunk_data = [(chunk, idx, params) for chunk, idx in zip(chunks, indices)]
        
        # Sequential processing untuk 1 core (lebih stabil)
        results = []
        for data in chunk_data:
            idx, result = self.process_chunk(data)
            if result is not None:
                results.append((idx, result))
        
        # Sort by original index
        results.sort(key=lambda x: x[0])
        
        # Merge chunks with crossfade
        return self.merge_chunks([r[1] for r in results])
    
    def merge_chunks(self, chunks):
        """Merge chunks dengan crossfade untuk menghindari artifacts"""
        if not chunks:
            return np.array([], dtype=np.float32)
        
        if len(chunks) == 1:
            return chunks[0]
        
        merged = chunks[0]
        
        for i in range(1, len(chunks)):
            current_chunk = chunks[i]
            
            if len(merged) < self.overlap or len(current_chunk) < self.overlap:
                # Jika chunk terlalu pendek, langsung concatenate
                merged = np.concatenate([merged, current_chunk])
            else:
                # Crossfade overlap region
                fade_out = merged[-self.overlap:]
                fade_in = current_chunk[:self.overlap]
                
                # Linear crossfade
                t = np.linspace(0, 1, self.overlap)
                faded = fade_out * (1 - t) + fade_in * t
                
                # Merge
                merged = np.concatenate([
                    merged[:-self.overlap],
                    faded,
                    current_chunk[self.overlap:]
                ])
        
        return merged

# ==============================
# MODEL CACHE SYSTEM
# ==============================
class ModelCache:
    _instance = None
    _models = {}
    _hubert = None
    
    def __new__(cls):
        if cls._instance is None:
            cls._instance = super(ModelCache, cls).__new__(cls)
        return cls._instance
    
    @classmethod
    def get_hubert(cls):
        if cls._hubert is None:
            cls._hubert = cls.load_hubert()
        return cls._hubert
    
    @classmethod
    def load_hubert(cls):
        """Load hubert model dengan optimasi"""
        print("Loading HuBERT model...")
        models, _, _ = checkpoint_utils.load_model_ensemble_and_task(
            ["hubert_base.pt"],
            suffix="",
        )
        hubert_model = models[0]
        hubert_model = hubert_model.to(config.device)
        
        if USE_HALF_PRECISION and config.is_half:
            hubert_model = hubert_model.half()
        else:
            hubert_model = hubert_model.float()
        
        hubert_model.eval()
        return hubert_model
    
    @classmethod
    def get_model(cls, model_path):
        """Get model from cache atau load baru"""
        if model_path in cls._models and CACHE_MODELS:
            print(f"Using cached model: {model_path}")
            return cls._models[model_path]
        
        return None
    
    @classmethod
    def cache_model(cls, model_path, model_data):
        """Cache model"""
        if CACHE_MODELS:
            cls._models[model_path] = model_data
    
    @classmethod
    def clear_model_cache(cls):
        """Clear model cache"""
        cls._models.clear()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        gc.collect()

# ==============================
# OPTIMIZED VC FUNCTION
# ==============================
def create_vc_fn(model_title, tgt_sr, net_g, vc, if_f0, version, file_index):
    def vc_fn(
        vc_audio_mode,
        vc_input, 
        vc_upload,
        tts_text,
        tts_voice,
        f0_up_key,
        f0_method,
        index_rate,
        filter_radius,
        resample_sr,
        rms_mix_rate,
        protect,
    ):
        try:
            # Auto clear memory sebelum proses
            MemoryManager.auto_clear()
            
            # Load audio
            if vc_audio_mode == "Input path" or "Youtube" and vc_input != "":
                # Gunakan librosa dengan optimasi
                audio, sr = librosa.load(vc_input, sr=16000, mono=True)
            elif vc_audio_mode == "Upload audio":
                if vc_upload is None:
                    return "You need to upload an audio", None
                sampling_rate, audio = vc_upload
                
                # Batasi durasi jika di space
                if limitation:
                    duration = audio.shape[0] / sampling_rate
                    if duration > 360:
                        return "Please upload an audio file that is less than 1 minute.", None
                
                # Konversi audio
                audio = (audio / np.iinfo(audio.dtype).max).astype(np.float32)
                if len(audio.shape) > 1:
                    audio = librosa.to_mono(audio.transpose(1, 0))
                if sampling_rate != 16000:
                    audio = librosa.resample(audio, orig_sr=sampling_rate, target_sr=16000)
            elif vc_audio_mode == "TTS Audio":
                if limitation and len(tts_text) > 600:
                    return "Text is too long", None
                if tts_text is None or tts_voice is None:
                    return "You need to enter text and select a voice", None
                
                # Generate TTS
                asyncio.run(edge_tts.Communicate(tts_text, "-".join(tts_voice.split('-')[:-1])).save("tts.mp3"))
                audio, sr = librosa.load("tts.mp3", sr=16000, mono=True)
                vc_input = "tts.mp3"
            else:
                return "Invalid audio mode", None
            
            # Persiapkan parameters
            hubert_model = ModelCache.get_hubert()
            times = [0, 0, 0]
            f0_up_key = int(f0_up_key)
            
            # Parameters untuk chunk processing
            params = (
                hubert_model,
                net_g,
                0,
                vc_input,
                times,
                f0_up_key,
                f0_method,
                file_index,
                index_rate,
                if_f0,
                filter_radius,
                tgt_sr,
                resample_sr,
                rms_mix_rate,
                version,
                protect
            )
            
            # Gunakan chunk processor untuk audio panjang
            if len(audio) > CHUNK_SIZE * 2:  # Hanya chunk jika > 20 detik
                print(f"Processing {len(audio)/16000:.2f}s audio in chunks...")
                processor = ChunkProcessor(vc.pipeline)
                audio_opt = processor.process_parallel(audio, params, max_workers=MAX_WORKERS)
            else:
                # Direct processing untuk audio pendek
                audio_opt = vc.pipeline(
                    hubert_model,
                    net_g,
                    0,
                    audio,
                    vc_input,
                    times,
                    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,
                )
            
            info = f"[{datetime.now().strftime('%Y-%m-%d %H:%M')}]: npy: {times[0]:.2f}s, f0: {times[1]:.2f}s, infer: {times[2]:.2f}s"
            print(f"{model_title} | {info}")
            
            # Clear memory setelah proses
            MemoryManager.auto_clear()
            
            return info, (tgt_sr, audio_opt)
        except Exception as e:
            info = f"Error: {str(e)}\n{traceback.format_exc()}"
            print(info)
            MemoryManager.auto_clear()
            return info, (None, None)
    return vc_fn

# ==============================
# OPTIMIZED MODEL LOADING
# ==============================
def load_model():
    categories = []
    with open("weights/folder_info.json", "r", encoding="utf-8") as f:
        folder_info = json.load(f)
    
    for category_name, category_info in folder_info.items():
        if not category_info['enable']:
            continue
        
        category_title = category_info['title']
        category_folder = category_info['folder_path']
        models = []
        
        with open(f"weights/{category_folder}/model_info.json", "r", encoding="utf-8") as f:
            models_info = json.load(f)
        
        for character_name, info in models_info.items():
            if not info['enable']:
                continue
            
            model_title = info['title']
            model_name = info['model_path']
            model_author = info.get("author", None)
            model_cover = f"weights/{category_folder}/{character_name}/{info['cover']}"
            model_index = f"weights/{category_folder}/{character_name}/{info['feature_retrieval_library']}"
            
            # Cek cache dulu
            model_path = f"weights/{category_folder}/{character_name}/{model_name}"
            cached_model = ModelCache.get_model(model_path)
            
            if cached_model:
                cpt, tgt_sr, if_f0, version, net_g = cached_model
            else:
                # Load model dengan memory optimasi
                cpt = torch.load(model_path, map_location="cpu", weights_only=False)
                tgt_sr = cpt["config"][-1]
                cpt["config"][-3] = cpt["weight"]["emb_g.weight"].shape[0]  # n_spk
                if_f0 = cpt.get("f0", 1)
                version = cpt.get("version", "v1")
                
                # Load model architecture
                if version == "v1":
                    if if_f0 == 1:
                        net_g = SynthesizerTrnMs256NSFsid(*cpt["config"], is_half=config.is_half)
                    else:
                        net_g = SynthesizerTrnMs256NSFsid_nono(*cpt["config"])
                elif version == "v2":
                    if if_f0 == 1:
                        net_g = SynthesizerTrnMs768NSFsid(*cpt["config"], is_half=config.is_half)
                    else:
                        net_g = SynthesizerTrnMs768NSFsid_nono(*cpt["config"])
                
                # Load weights
                del net_g.enc_q
                net_g.load_state_dict(cpt["weight"], strict=False)
                net_g.eval().to(config.device)
                
                if USE_HALF_PRECISION and config.is_half:
                    net_g = net_g.half()
                else:
                    net_g = net_g.float()
                
                # Cache model
                ModelCache.cache_model(model_path, (cpt, tgt_sr, if_f0, version, net_g))
            
            # Create VC
            vc = VC(tgt_sr, config)
            models.append((character_name, model_title, model_author, model_cover, 
                          version.upper(), create_vc_fn(model_title, tgt_sr, net_g, vc, if_f0, version, model_index)))
            
            print(f"Model loaded: {character_name} | {version.upper()}")
        
        categories.append([category_title, category_folder, models])
    
    return categories

# ==============================
# OPTIMIZED FUNCTIONS
# ==============================
def cut_vocal_and_inst(url, audio_provider, split_model):
    """Optimized audio splitting dengan progress feedback"""
    if not url:
        raise gr.Error("URL Required!")
    
    if not os.path.exists("dl_audio"):
        os.mkdir("dl_audio")
    
    try:
        if audio_provider == "Youtube":
            ydl_opts = {
                'format': 'bestaudio/best',
                'postprocessors': [{
                    'key': 'FFmpegExtractAudio',
                    'preferredcodec': 'wav',
                }],
                "outtmpl": 'dl_audio/youtube_audio',
                'quiet': True,
                'no_warnings': True,
            }
            
            with yt_dlp.YoutubeDL(ydl_opts) as ydl:
                ydl.download([url])
            
            audio_path = "dl_audio/youtube_audio.wav"
        
        # Optimize demucs command
        model_name = "htdemucs" if split_model == "htdemucs" else "mdx_extra_q"
        output_dir = f"output/{model_name}"
        
        # Gunakan subprocess dengan Popen untuk better control
        cmd = ["demucs", "--two-stems=vocals", "-n", model_name, 
               audio_path, "-o", "output", "--quiet"]
        
        process = subprocess.Popen(cmd, stdout=subprocess.PIPE, 
                                  stderr=subprocess.PIPE, text=True)
        
        # Monitor progress
        for line in process.stdout:
            if "Progress" in line:
                print(line.strip())
        
        process.wait()
        
        vocal_path = f"{output_dir}/youtube_audio/vocals.wav"
        inst_path = f"{output_dir}/youtube_audio/no_vocals.wav"
        
        return vocal_path, inst_path, audio_path, vocal_path
    
    except Exception as e:
        raise gr.Error(f"Error processing audio: {str(e)}")

def combine_vocal_and_inst(audio_data, audio_volume, split_model):
    """Optimized audio combining"""
    if not os.path.exists("output/result"):
        os.makedirs("output/result")
    
    vocal_path = "output/result/output.wav"
    output_path = "output/result/combine.mp3"
    
    model_name = "htdemucs" if split_model == "htdemucs" else "mdx_extra_q"
    inst_path = f"output/{model_name}/youtube_audio/no_vocals.wav"
    
    # Write vocal file
    with wave.open(vocal_path, "w") as wave_file:
        wave_file.setnchannels(1)
        wave_file.setsampwidth(2)
        wave_file.setframerate(audio_data[0])
        wave_file.writeframes(audio_data[1].tobytes())
    
    # Optimize ffmpeg command
    cmd = [
        'ffmpeg', '-y',
        '-i', inst_path,
        '-i', vocal_path,
        '-filter_complex', f'[1:a]volume={audio_volume}dB[v];[0:a][v]amix=inputs=2:duration=longest',
        '-b:a', '320k',
        '-c:a', 'libmp3lame',
        '-threads', '1',  # Gunakan single thread
        '-loglevel', 'error',
        output_path
    ]
    
    subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    return output_path

# ==============================
# MAIN APPLICATION
# ==============================
if __name__ == '__main__':
    # Setup audio mode
    audio_mode = []
    f0method_mode = []
    f0method_info = ""
    
    if limitation is True:
        audio_mode = ["Upload audio", "TTS Audio"]
        f0method_mode = ["pm", "crepe", "harvest"]
        f0method_info = "PM is fast, rmvpe is middle, Crepe or harvest is good but it was extremely slow (Default: PM)"
    else:
        audio_mode = ["Upload audio", "Youtube", "TTS Audio"]
        f0method_mode = ["pm", "crepe", "harvest"]
        f0method_info = "PM is fast, rmvpe is middle. Crepe or harvest is good but it was extremely slow (Default: PM)"
    
    if os.path.isfile("rmvpe.pt"):
        f0method_mode.insert(2, "rmvpe")
    
    # Pre-load hubert model
    print("Initializing HuBERT model...")
    ModelCache.get_hubert()
    
    # Load models
    print("Loading RVC models...")
    categories = load_model()
    
    # Get TTS voices
    tts_voice_list = asyncio.get_event_loop().run_until_complete(edge_tts.list_voices())
    voices = [f"{v['ShortName']}-{v['Gender']}" for v in tts_voice_list]
    
    # Create Gradio app
    with gr.Blocks(theme=gr.themes.Base(), css="""
        .gradio-container {max-width: 1200px !important;}
        .tab-nav {scroll-behavior: smooth;}
    """) as app:
        gr.Markdown("""
        # ▶️ RVC Youtuber Indonesia 👳🏿‍♂️
        [⚡] **Optimized Version** - Chunk Processing Enabled
        """)
        
        for (folder_title, folder, models) in categories:
            with gr.TabItem(folder_title):
                if not models:
                    gr.Markdown("# <center> No Model Loaded.")
                    continue
                
                for (name, title, author, cover, model_version, vc_fn) in models:
                    with gr.TabItem(name):
                        with gr.Row():
                            gr.Markdown(f"""
                            <div align="center">
                            <div><b>{title}</b></div>
                            <div>RVC {model_version} Model</div>
                            {f'<div>Model author: {author}</div>' if author else ""}
                            {f'<img style="width:auto;height:300px;" src="file/{cover}">' if cover else ""}
                            </div>
                            """)
                        
                        with gr.Row():
                            with gr.Column():
                                vc_audio_mode = gr.Dropdown(
                                    label="Input voice", 
                                    choices=audio_mode, 
                                    value="Upload audio"
                                )
                                
                                # Audio inputs
                                vc_input = gr.Textbox(label="Input audio path", visible=False)
                                vc_upload = gr.Audio(label="Upload audio file", visible=True)
                                
                                # Youtube
                                vc_link = gr.Textbox(label="Youtube URL", visible=False)
                                vc_split_model = gr.Dropdown(
                                    label="Splitter Model", 
                                    choices=["htdemucs", "mdx_extra_q"], 
                                    value="htdemucs", 
                                    visible=False
                                )
                                vc_split = gr.Button("Split Audio", variant="primary", visible=False)
                                vc_vocal_preview = gr.Audio(label="Vocal Preview", visible=False)
                                vc_inst_preview = gr.Audio(label="Instrumental Preview", visible=False)
                                vc_audio_preview = gr.Audio(label="Audio Preview", visible=False)
                                
                                # TTS
                                tts_text = gr.Textbox(visible=False, label="TTS text")
                                tts_voice = gr.Dropdown(
                                    label="Edge-tts speaker", 
                                    choices=voices, 
                                    visible=False, 
                                    value="en-US-AnaNeural-Female"
                                )
                            
                            with gr.Column():
                                vc_transform0 = gr.Number(
                                    label="Transpose", 
                                    value=0,
                                    info='Type "12" for male to female, "-12" for female to male'
                                )
                                
                                f0method0 = gr.Radio(
                                    label="Pitch extraction algorithm",
                                    info=f0method_info,
                                    choices=f0method_mode,
                                    value="pm"
                                )
                                
                                index_rate1 = gr.Slider(
                                    minimum=0, maximum=1,
                                    label="Retrieval feature ratio",
                                    value=0.4,
                                    info="Too high = robotic"
                                )
                                
                                filter_radius0 = gr.Slider(
                                    minimum=0, maximum=7,
                                    label="Median Filtering",
                                    value=1, step=1,
                                    info="Reduce breathiness"
                                )
                                
                                resample_sr0 = gr.Slider(
                                    minimum=0, maximum=48000,
                                    label="Resample output",
                                    value=0, step=1,
                                    info="0 for no resampling"
                                )
                                
                                rms_mix_rate0 = gr.Slider(
                                    minimum=0, maximum=1,
                                    label="Volume Envelope",
                                    value=1,
                                    info="1 = use output envelope"
                                )
                                
                                protect0 = gr.Slider(
                                    minimum=0, maximum=0.5,
                                    label="Voice Protection",
                                    value=0.23, step=0.01,
                                    info="Protect voiceless sounds"
                                )
                            
                            with gr.Column():
                                vc_log = gr.Textbox(label="Output Information")
                                vc_output = gr.Audio(label="Output Audio")
                                vc_convert = gr.Button("Convert", variant="primary")
                                
                                vc_volume = gr.Slider(
                                    minimum=0, maximum=10,
                                    label="Vocal volume", value=4, step=1,
                                    visible=False
                                )
                                
                                vc_combined_output = gr.Audio(
                                    label="Output Combined Audio", 
                                    visible=False
                                )
                                vc_combine = gr.Button("Combine", variant="primary", visible=False)
                        
                        # Connect events
                        vc_convert.click(
                            fn=vc_fn, 
                            inputs=[
                                vc_audio_mode, vc_input, vc_upload,
                                tts_text, tts_voice, vc_transform0,
                                f0method0, index_rate1, filter_radius0,
                                resample_sr0, rms_mix_rate0, protect0,
                            ], 
                            outputs=[vc_log, vc_output]
                        ).then(
                            fn=lambda: MemoryManager.auto_clear(),
                            outputs=[]
                        )
                        
                        vc_split.click(
                            fn=cut_vocal_and_inst, 
                            inputs=[vc_link, gr.Dropdown(value="Youtube", visible=False), vc_split_model], 
                            outputs=[vc_vocal_preview, vc_inst_preview, vc_audio_preview, vc_input]
                        )
                        
                        vc_combine.click(
                            fn=combine_vocal_and_inst,
                            inputs=[vc_output, vc_volume, vc_split_model],
                            outputs=[vc_combined_output]
                        )
                        
                        vc_audio_mode.change(
                            fn=lambda mode: (
                                gr.Textbox.update(visible=(mode == "Input path")),
                                gr.Audio.update(visible=(mode == "Upload audio")),
                                gr.Textbox.update(visible=(mode == "Youtube")),
                                gr.Dropdown.update(visible=(mode == "Youtube")),
                                gr.Button.update(visible=(mode == "Youtube")),
                                gr.Audio.update(visible=(mode == "Youtube")),
                                gr.Audio.update(visible=(mode == "Youtube")),
                                gr.Audio.update(visible=(mode == "Youtube")),
                                gr.Slider.update(visible=(mode == "Youtube")),
                                gr.Audio.update(visible=(mode == "Youtube")),
                                gr.Button.update(visible=(mode == "Youtube")),
                                gr.Textbox.update(visible=(mode == "TTS Audio")),
                                gr.Dropdown.update(visible=(mode == "TTS Audio"))
                            ),
                            inputs=[vc_audio_mode],
                            outputs=[
                                vc_input, vc_upload, vc_link, vc_split_model,
                                vc_split, vc_vocal_preview, vc_inst_preview,
                                vc_audio_preview, vc_volume, vc_combined_output,
                                vc_combine, tts_text, tts_voice
                            ]
                        )
    
    # Launch app
    # Definisikan parameter launch

    
launch_kwargs = {
    'enable_queue': True,      
    'max_threads': 1,         
    'share': config.colab if limitation else True,  
    'server_name': "0.0.0.0"  
}

app.launch(**launch_kwargs)