import os
import sys
import importlib.util
import torch
import torchaudio
import numpy as np
from huggingface_hub import snapshot_download, hf_hub_download
import subprocess
import uuid
import soundfile as sf
import spaces
import librosa
import shutil
import re
import gradio as gr

# --- 1. نصب و تنظیمات اولیه (بدون تغییر) ---

def install_espeak():
    try:
        result = subprocess.run(["which", "espeak-ng"], capture_output=True, text=True)
        if result.returncode != 0:
            print("Installing espeak-ng...")
            subprocess.run(["apt-get", "update"], check=True)
            subprocess.run(["apt-get", "install", "-y", "espeak-ng", "espeak-ng-data"], check=True)
    except Exception as e:
        print(f"Error installing espeak-ng: {e}")

install_espeak()

def patch_langsegment_init():
    try:
        spec = importlib.util.find_spec("LangSegment")
        if spec is None or spec.origin is None: return
        init_path = os.path.join(os.path.dirname(spec.origin), '__init__.py')
        with open(init_path, 'r') as f: lines = f.readlines()
        modified = False
        new_lines =[]
        target_line_prefix = "from .LangSegment import"
        for line in lines:
            if line.strip().startswith(target_line_prefix) and ('setLangfilters' in line or 'getLangfilters' in line):
                mod_line = line.replace(',setLangfilters', '').replace(',getLangfilters', '')
                mod_line = mod_line.replace('setLangfilters,', '').replace('getLangfilters,', '').rstrip(',')
                new_lines.append(mod_line + '\n')
                modified = True
            else:
                new_lines.append(line)
        if modified:
            with open(init_path, 'w') as f: f.writelines(new_lines)
            try:
                import LangSegment
                importlib.reload(LangSegment)
            except: pass
    except: pass

patch_langsegment_init()

if not os.path.exists("Amphion"):
    print("Cloning Amphion repository...")
    subprocess.run(["git", "clone", "https://github.com/open-mmlab/Amphion.git"])

amphion_path = os.path.abspath("Amphion")
if amphion_path not in sys.path:
    sys.path.append(amphion_path)

# --- پچ کردن تمام کدهای Amphion برای سازگاری با نسخه جدید transformers ---
def patch_amphion_llama_config():
    try:
        for root, dirs, files in os.walk(amphion_path):
            for file in files:
                if file.endswith(".py"):
                    file_path = os.path.join(root, file)
                    try:
                        with open(file_path, "r", encoding="utf-8") as f:
                            content = f.read()
                        
                        new_content = re.sub(
                            r'LlamaConfig\(\s*0\s*,\s*256\s*,\s*1024\s*,\s*1\s*,\s*1\s*\)',
                            r'LlamaConfig(vocab_size=0, hidden_size=256, intermediate_size=1024, num_hidden_layers=1, num_attention_heads=1)',
                            content
                        )
                        
                        if new_content != content:
                            with open(file_path, "w", encoding="utf-8") as f:
                                f.write(new_content)
                            print(f"Patched {file} successfully.")
                    except Exception:
                        pass
    except Exception as e:
        print(f"Failed to patch Amphion: {e}")

patch_amphion_llama_config()

# --- رفع مشکل rope_theta با تزریق (Monkey Patch) در LlamaConfig ---
import transformers
from transformers.models.llama.configuration_llama import LlamaConfig

_original_llama_init = LlamaConfig.__init__
def _patched_llama_init(self, *args, **kwargs):
    kwargs.setdefault('rope_theta', 10000.0)
    kwargs.setdefault('attention_bias', False)
    kwargs.setdefault('rope_scaling', None)
    kwargs.setdefault('max_position_embeddings', 4096)
    
    _original_llama_init(self, *args, **kwargs)
    
    if not hasattr(self, 'rope_theta'):
        self.rope_theta = 10000.0
    if not hasattr(self, 'attention_bias'):
        self.attention_bias = False

LlamaConfig.__init__ = _patched_llama_init

os.makedirs("wav", exist_ok=True)
os.makedirs("ckpts/Vevo", exist_ok=True)

try:
    from models.vc.vevo.vevo_utils import VevoInferencePipeline
except ImportError:
    sys.path.append(os.path.join(amphion_path))
    from models.vc.vevo.vevo_utils import VevoInferencePipeline

# --- توابع ذخیره و تنظیمات مدل ---

def save_audio_pcm16(waveform, output_path, sample_rate=24000):
    try:
        if isinstance(waveform, torch.Tensor):
            waveform = waveform.detach().cpu()
            if waveform.dim() == 2 and waveform.shape[0] == 1:
                waveform = waveform.squeeze(0)
            waveform = waveform.numpy()
        sf.write(output_path, waveform, sample_rate, subtype='PCM_16')
    except Exception as e:
        print(f"Save error: {e}")

downloaded_resources = {
    "configs": False,
    "tokenizer_vq8192": False,
    "fmt_Vq8192ToMels": False,
    "vocoder": False
}

def setup_configs():
    if downloaded_resources["configs"]: return
    target_config_path = "models/vc/vevo/config"
    os.makedirs(target_config_path, exist_ok=True)
    
    config_files =["Vq8192ToMels.json", "Vocoder.json", "hubert_large_l18_mean_std.npz"]
    source_config_path = os.path.join(amphion_path, "models/vc/vevo/config")

    for file in config_files:
        target_file_path = f"{target_config_path}/{file}"
        source_file_path = os.path.join(source_config_path, file)
        
        if not os.path.exists(target_file_path):
            if os.path.exists(source_file_path):
                shutil.copy(source_file_path, target_file_path)
            else:
                try:
                    file_data = hf_hub_download(repo_id="amphion/Vevo", filename=f"config/{file}", repo_type="model")
                    if os.path.exists(target_file_path): os.remove(target_file_path)
                    subprocess.run(["cp", file_data, target_file_path])
                except Exception as e: 
                    print(f"Error downloading config {file}: {e}")
                    
    downloaded_resources["configs"] = True

setup_configs()
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
inference_pipelines = {}

downloaded_content_style_tokenizer_path = None
downloaded_fmt_path = None
downloaded_vocoder_path = None

def preload_all_resources():
    setup_configs()
    global downloaded_content_style_tokenizer_path, downloaded_fmt_path, downloaded_vocoder_path
    if not downloaded_resources["tokenizer_vq8192"]:
        downloaded_content_style_tokenizer_path = snapshot_download(repo_id="amphion/Vevo", repo_type="model", cache_dir="./ckpts/Vevo", allow_patterns=["tokenizer/vq8192/"])
        downloaded_resources["tokenizer_vq8192"] = True
    if not downloaded_resources["fmt_Vq8192ToMels"]:
        downloaded_fmt_path = snapshot_download(repo_id="amphion/Vevo", repo_type="model", cache_dir="./ckpts/Vevo", allow_patterns=["acoustic_modeling/Vq8192ToMels/"])
        downloaded_resources["fmt_Vq8192ToMels"] = True
    if not downloaded_resources["vocoder"]:
        downloaded_vocoder_path = snapshot_download(repo_id="amphion/Vevo", repo_type="model", cache_dir="./ckpts/Vevo", allow_patterns=["acoustic_modeling/Vocoder/*"])
        downloaded_resources["vocoder"] = True

preload_all_resources()

def get_pipeline():
    if "timbre" in inference_pipelines: return inference_pipelines["timbre"]
    tokenizer_path = os.path.join(downloaded_content_style_tokenizer_path, "tokenizer/vq8192")
    fmt_ckpt = os.path.join(downloaded_fmt_path, "acoustic_modeling/Vq8192ToMels")
    vocoder_ckpt = os.path.join(downloaded_vocoder_path, "acoustic_modeling/Vocoder")
    pipeline = VevoInferencePipeline(
        content_style_tokenizer_ckpt_path=tokenizer_path,
        fmt_cfg_path="./models/vc/vevo/config/Vq8192ToMels.json",
        fmt_ckpt_path=fmt_ckpt,
        vocoder_cfg_path="./models/vc/vevo/config/Vocoder.json",
        vocoder_ckpt_path=vocoder_ckpt,
        device=device,
    )
    inference_pipelines["timbre"] = pipeline
    return pipeline


# --- تابع جدید برش هوشمند با قابلیت افزایش داینامیک بازه جستجو ---
def smart_split_audio_vevo(y, sr, min_segment=4):
    total_samples = len(y)
    current_sample = 0
    chunks = []
    
    # اگر طول کل فایل کمتر از حد مجاز اولیه (8 ثانیه) باشد، نیازی به برش نیست
    if total_samples <= 8 * sr:
        return [y]
        
    while current_sample < total_samples:
        # اگر حجم نمونه‌های باقی‌مانده کمتر از حداقل زمان قطعه (4 ثانیه) باشد، بقیه فایل را یکجا بردار
        if total_samples - current_sample <= min_segment * sr:
            chunks.append(y[current_sample:])
            break
            
        # بازه‌های گام‌به‌گام برای جستجوی سکوت: ابتدا ۸، سپس ۱۲، ۱۵ و نهایتاً ۲۰ ثانیه
        steps = [8, 12, 15, 20]
        silence_threshold = 0.015  # آستانه شدت انرژی سیگنال برای پذیرش به عنوان سکوت واقعی
        
        best_split_point = None
        best_rms_val = float('inf')
        
        for max_sec in steps:
            start_search = current_sample + int(min_segment * sr)
            end_search = current_sample + int(max_sec * sr)
            
            if start_search >= total_samples:
                best_split_point = total_samples
                break
            
            if end_search > total_samples:
                end_search = total_samples
                
            search_region = y[start_search:end_search]
            if len(search_region) == 0:
                best_split_point = total_samples
                break
                
            # محاسبه میزان انرژی (RMS) قطعه
            rms = librosa.feature.rms(y=search_region, frame_length=2048, hop_length=512)[0]
            if len(rms) == 0:
                continue
                
            min_rms_idx = np.argmin(rms)
            min_rms_val = rms[min_rms_idx]
            candidate_point = start_search + (min_rms_idx * 512)
            
            # ذخیره بهترین نقطه کم‌صدا (در صورتی که سکوت مطلق پیدا نشود، کم‌صدا‌ترین نقطه انتخاب می‌شود)
            if min_rms_val < best_rms_val:
                best_rms_val = min_rms_val
                best_split_point = candidate_point
                
            # اگر شدت صدای نقطه پیدا شده پایین‌تر از آستانه سکوت بود، فرآیند جستجو را متوقف کرده و برش بزن
            if min_rms_val < silence_threshold:
                break
        
        if best_split_point is None:
            best_split_point = min(current_sample + int(8 * sr), total_samples)
            
        chunk_audio = y[current_sample:best_split_point]
        chunks.append(chunk_audio)
        current_sample = best_split_point
        
    # بررسی قطعه آخر: اگر بیش از حد کوتاه باشد، آن را با بخش قبلی ادغام می‌کنیم
    if len(chunks) > 1 and len(chunks[-1]) < int(1.5 * sr):
        last_chunk = chunks.pop()
        combined_len = len(chunks[-1]) + len(last_chunk)
        # تا سقف حداکثر ۲۰ ثانیه ادغام قطعه آخر مجاز است
        if combined_len <= int(20 * sr):
            chunks[-1] = np.concatenate([chunks[-1], last_chunk])
        else:
            chunks.append(last_chunk)
            
    return chunks


# =========================================================================
#  سیستم استنتاج مستقیم بر روی Zero-GPU با رابط کاربری Gradio
# =========================================================================

@spaces.GPU(duration=120)
def predict_voice_conversion(source_audio_path, ref_audio_path):
    if source_audio_path is None or ref_audio_path is None:
        raise gr.Error("لطفاً هر دو فایل صوتی (اصلی و مرجع) را وارد کنید.")

    task_id = str(uuid.uuid4())
    out_path = f"wav/{task_id}_out.wav"
    SR = 24000

    try:
        # پردازش صوت اصلی
        content_data, _ = librosa.load(source_audio_path, sr=SR)
        
        # پردازش صوت مرجع
        ref_data, _ = librosa.load(ref_audio_path, sr=SR)
        ref_tensor = torch.FloatTensor(ref_data).unsqueeze(0)
        ref_tensor = ref_tensor / (torch.max(torch.abs(ref_tensor)) + 1e-6) * 0.95
        
        temp_r = f"wav/{task_id}_temp_r.wav"
        save_audio_pcm16(ref_tensor, temp_r, SR)

        # اعمال تقسیم‌کننده هوشمند چندمرحله‌ای
        chunks = smart_split_audio_vevo(content_data, SR, min_segment=4)

        generated_chunks = []
        pipeline = get_pipeline()

        # شبیه‌سازی قطعه به قطعه
        for idx, chunk in enumerate(chunks):
            chunk_tensor = torch.FloatTensor(chunk).unsqueeze(0)
            chunk_tensor = chunk_tensor / (torch.max(torch.abs(chunk_tensor)) + 1e-6) * 0.95
            
            temp_c_chunk = f"wav/{task_id}_temp_c_{idx}.wav"
            save_audio_pcm16(chunk_tensor, temp_c_chunk, SR)

            gen = pipeline.inference_fm(
                src_wav_path=temp_c_chunk,
                timbre_ref_wav_path=temp_r,
                flow_matching_steps=32,
            )

            if torch.isnan(gen).any(): 
                gen = torch.nan_to_num(gen, nan=0.0)
            gen_np = gen.detach().cpu().squeeze().numpy()
            
            generated_chunks.append(gen_np)

            if os.path.exists(temp_c_chunk): 
                os.remove(temp_c_chunk)

        # ادغام قطعات خروجی
        final_gen = np.concatenate(generated_chunks, axis=0)

        # ذخیره نهایی
        save_audio_pcm16(final_gen, out_path, SR)

        if os.path.exists(temp_r): 
            os.remove(temp_r)

        return out_path

    except Exception as e:
        raise gr.Error(f"خطایی در حین پردازش رخ داد: {str(e)}")

# طراحی ساختار وب با استفاده از Gradio
with gr.Blocks() as demo:
    gr.Markdown("# سامانه تبدیل صدا (Vevo Voice Conversion) - نسخه Zero-GPU PRO")
    
    with gr.Row():
        with gr.Column():
            source_input = gr.Audio(type="filepath", label="فایل صوتی اصلی (Source Audio)")
            ref_input = gr.Audio(type="filepath", label="فایل صوتی مرجع/هدف (Reference Timbre)")
            submit_btn = gr.Button("شروع فرآیند تبدیل صدا", variant="primary")
        
        with gr.Column():
            audio_output = gr.Audio(label="خروجی صدای شبیه‌سازی شده", type="filepath")

    submit_btn.click(
        fn=predict_voice_conversion,
        inputs=[source_input, ref_input],
        outputs=audio_output
    )

if __name__ == "__main__":
    demo.queue(max_size=20).launch(server_name="0.0.0.0", server_port=7860, theme=gr.themes.Soft())