# ==========================================================================
# FILE: voice_engine.py  — موتورِ «کلون کردن صدا» (Vevo)، کاملاً روی CPU
# --------------------------------------------------------------------------
# این فایل همان کدِ اسپیسِ صدای قبلی (Sada) است، بدونِ رابطِ Gradio:
#   • همان پچ‌ها (pkg_resources، LangSegment، Amphion و transformers)
#   • همان مدل: amphion/Vevo — تبدیلِ timbre با flow matching در ۳۲ گام
#   • همان برشِ هوشمند از روی سکوت (smart_split_audio_vevo) و همان ادغام
#   • همان تعدادِ هسته: همه‌ی هسته‌های واقعیِ کانتینر (cgroup)، بدونِ GPU
# هیچ تنظیمی از مدل عوض نشده؛ کیفیتِ خروجی دقیقاً مثلِ قبل است.
#
# چرا یک پروسه‌ی جدا (و نه داخلِ همان پروسه‌ی Gradio)؟
#   در این اسپیس «ساخت آهنگ» روی ZeroGPU است. ZeroGPU برای هر کارِ GPU از
#   پروسه‌ی اصلی fork می‌گیرد. اگر تبدیلِ صدا با ۱۶ رشته‌ی CPU داخلِ همان پروسه
#   اجرا شود، fork وسطِ کارِ چندرشته‌ایِ OpenMP/MKL می‌تواند پروسه‌ی GPU را قفل
#   کند. پچ‌های سراسریِ این موتور (LlamaConfig، pkg_resources و ...) هم نباید
#   روی مدل‌های آهنگ اثر بگذارند. برای همین این موتور در پروسه‌ی خودش اجرا
#   می‌شود و spaces/ZeroGPU را اصلاً ایمپورت نمی‌کند؛ یعنی دقیقاً مثلِ اسپیسِ
#   صدای قبلی، فقط CPU و RAM مصرف می‌کند و هیچ سهمیه‌ی GPU نمی‌گیرد.
#
# اجرا: voice_clone.py این فایل را با  python -u voice_engine.py --serve <socket>
# بالا می‌آورد و کارها را از راهِ یک سوکتِ محلی (فقط داخلِ همین کانتینر) می‌فرستد.
#
# 🆕 تغییرهای این نسخه (طبقِ درخواست و از روی لاگ‌های اسپیسِ قبلی):
#   ۱) صدای مرجع حداکثر ۱۲ ثانیه (VOICE_REF_MAX_SECONDS): مدلِ Vevo صدای مرجع را جلوی «هر» تکه
#      می‌گذارد؛ مرجعِ ۳ دقیقه‌ای یعنی هر تکه‌ی ۶ ثانیه‌ای مثلِ یک صدای ۳ دقیقه‌ای حساب می‌شد
#      (همان ۲۰۰ ثانیه‌ای‌های لاگ). مدل با مرجعِ حداکثر ~۱۲ ثانیه آموزش دیده؛ پس از مرجعِ بلندتر
#      «پُرصداترین» ۱۲ ثانیه برداشته می‌شود (از روی سکوت بریده، نه وسطِ کلمه). مرجعِ کوتاه‌تر
#      دست نمی‌خورد.
#   ۲) دو کارِ هم‌زمان (VOICE_CONCURRENT_JOBS=2، فقط وقتی اکشن پرو جای خالیِ دیگری ندارد):
#      کارِ تنها همه‌ی هسته‌ها را می‌گیرد (مثلِ قبل)؛ وقتی دو کار با هم‌اند هر کدام نصف. تقسیمِ
#      هسته‌ها در «هر گامِ» مدل دوباره حساب می‌شود (نه فقط اولِ کار).
#   ۳) کارِ بی‌صاحب لغو می‌شود: اگر درخواست‌دهنده دیگر منتظر نبود (اتصال قطع شد / زمانش تمام شد)،
#      همان لحظه (در گامِ بعدیِ مدل) متوقف می‌شود و CPU را برای کارهای واقعی آزاد می‌کند. قبلاً
#      کارِ بی‌صاحب تا آخر ادامه می‌داد و کارهای بعدی پشتش می‌ماندند (BrokenPipe و ۶۴۵/۶۴۷ ثانیه‌ای‌های
#      لاگ برای فایل‌های دوتکه‌ای).
# هر تکه دقیقاً مثلِ قبل ساخته می‌شود (همان نرمال‌سازی، همان ۳۲ گام، همان مدل) و تکه‌های یک
# فایل مثلِ قبل یکی‌یکی و به ترتیب ساخته و به هم وصل می‌شوند.
# ==========================================================================
import os
import sys
import json
import time
import signal
import threading
import importlib.util
import torch
import torchaudio  # noqa: F401  (مثلِ قبل: Amphion به آن نیاز دارد)
import numpy as np
from huggingface_hub import snapshot_download, hf_hub_download
import subprocess
import uuid
import soundfile as sf
import librosa
import shutil
import re
import warnings

# مرجعی که PySoundFile نمی‌خواند (بعضی mp3ها) خودکار با audioread خوانده می‌شود؛ فقط هشدارِ
# «منسوخ شدن» در لاگ می‌ماند که هیچ خطایی نیست — بی‌صدا شد تا لاگ تمیز بماند (رفتار همان است).
warnings.filterwarnings("ignore", message="PySoundFile failed")
warnings.filterwarnings("ignore", category=FutureWarning, module="librosa")

# --- رفع خطای pkg_resources (نسخه‌های جدید setuptools آن را حذف کرده‌اند و pyworld به آن نیاز دارد) ---
try:
    import pkg_resources  # noqa: F401
except ImportError:
    import types
    from importlib import metadata as _md

    class _DistributionNotFound(Exception):
        pass

    class _Dist:
        def __init__(self, name):
            try:
                self.version = _md.version(name)
            except _md.PackageNotFoundError:
                raise _DistributionNotFound(name)
            self.project_name = name
            self.key = name.lower()

    _pr = types.ModuleType("pkg_resources")
    _pr.DistributionNotFound = _DistributionNotFound
    _pr.get_distribution = lambda name: _Dist(name)
    _pr.require = lambda *a, **k: []
    sys.modules["pkg_resources"] = _pr
    print("[FIX] pkg_resources shim installed")

# --- 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()

# =========================================================================
#  اجرای کامل روی CPU (بدون مصرف اعتبار ZeroGPU)
#  نکته: این پروسه spaces را ایمپورت نمی‌کند، پس GPUی در کار نیست؛ با این حال
#  دستگاه مثلِ قبل صریحاً CPU گذاشته می‌شود. کیفیت و تنظیمات مدل دست نخورده است.
# =========================================================================
def get_real_cpu_count():
    n = os.cpu_count() or 1
    try:
        n = min(n, len(os.sched_getaffinity(0)))
    except Exception:
        pass
    try:
        # سهمیه واقعی کانتینر (cgroup v2)
        with open("/sys/fs/cgroup/cpu.max") as f:
            quota, period = f.read().split()
            if quota != "max":
                n = min(n, max(1, int(int(quota) / int(period))))
    except Exception:
        pass
    return max(1, n)

CPU_THREADS = get_real_cpu_count()
# حداکثر چند درخواست هم‌زمان (۲ = «کارِ دوم هم‌زمان»؛ ۱ = رفتارِ قبلی). اکشن پرو کارِ دوم را فقط
# وقتی می‌فرستد که هیچ اسپیسِ صدای خالیِ دیگری نمانده باشد.
VOICE_CONCURRENT_JOBS = max(1, min(4, int(os.environ.get("VOICE_CONCURRENT_JOBS", "2") or 2)))
if CPU_THREADS < 2 * VOICE_CONCURRENT_JOBS:
    VOICE_CONCURRENT_JOBS = max(1, CPU_THREADS // 2)
torch.set_num_threads(CPU_THREADS)
try:
    torch.set_num_interop_threads(max(1, min(4, CPU_THREADS // 4)))
except Exception:
    pass
os.environ["OMP_NUM_THREADS"] = str(CPU_THREADS)
os.environ["MKL_NUM_THREADS"] = str(CPU_THREADS)
print(f"[CPU MODE] Using {CPU_THREADS} CPU threads (os.cpu_count={os.cpu_count()}); "
      f"up to {VOICE_CONCURRENT_JOBS} job(s) at once (a lone job gets all threads)")

device = 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()

_pipeline_lock = threading.Lock()

def get_pipeline():
    with _pipeline_lock:
        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,
        )
        _install_step_hooks(pipeline)
        inference_pipelines["timbre"] = pipeline
        return pipeline


# =========================================================================
#  🆕 تقسیمِ هسته‌ها بینِ کارهای هم‌زمان + لغوِ کارِ بی‌صاحب، در «هر گامِ» مدل
# -------------------------------------------------------------------------
# روی ماژول‌های سنگینِ مدل (هر گامِ flow matching، لایه‌های HuBERT، vocoder ...) یک
# forward-pre-hook گذاشته می‌شود که در همان رشته‌ی کار اجرا می‌شود و دو کار می‌کند:
#   • اگر کار لغو شده (درخواست‌دهنده دیگر منتظر نیست) همان لحظه متوقفش می‌کند؛
#   • تعدادِ رشته‌های همین کار را با «همه‌ی هسته‌ها ÷ تعدادِ کارهای در جریان» یکی می‌کند
#     (torch.set_num_threads برای هر رشته جداست). کارِ تنها = همه‌ی هسته‌ها، مثلِ قبل.
# هیچ عددی از خودِ مدل عوض نمی‌شود؛ فقط این‌که هر لحظه چند هسته کار کنند.
# =========================================================================
_tls = threading.local()
_jobs_lock = threading.Lock()
_active_jobs = 0


class _JobCancelled(Exception):
    pass


class _Job:
    def __init__(self):
        self.cancelled = threading.Event()
        self.done = threading.Event()
        self.reason = ""
        self.threads = 0
        self.max_parallel = 1

    def cancel(self, reason):
        if self.done.is_set() or self.cancelled.is_set():
            return
        self.reason = reason
        self.cancelled.set()

    def check(self):
        if self.cancelled.is_set():
            raise _JobCancelled(self.reason or "cancelled")


def _sync_threads(job):
    """رشته‌های کارِ همین رشته را با تعدادِ کارهای در جریان هماهنگ می‌کند."""
    active = max(1, _active_jobs)
    if active > job.max_parallel:
        job.max_parallel = active
    want = max(1, CPU_THREADS // active)
    if job.threads != want:
        torch.set_num_threads(want)
        job.threads = want


def _step_hook(module, args):
    job = getattr(_tls, "job", None)
    if job is not None:
        job.check()
        _sync_threads(job)
    return None


def _install_step_hooks(pipeline):
    """هوک روی ماژول‌هایی که در هر تکه بارها صدا زده می‌شوند (و لایه‌های داخلِ ModuleListهایشان)."""
    names = ("fmt_model", "hubert_model", "vocoder_model", "mel_model", "content_style_tokenizer")
    hooked, roots = 0, []
    for name in names:
        root = getattr(pipeline, name, None)
        if not isinstance(root, torch.nn.Module):
            continue
        roots.append(name)
        targets = {id(root): root}
        estimator = getattr(root, "diff_estimator", None)
        if isinstance(estimator, torch.nn.Module):
            targets[id(estimator)] = estimator          # هر گامِ flow matching
        for sub in root.modules():
            if isinstance(sub, torch.nn.ModuleList):
                for layer in sub:
                    if isinstance(layer, torch.nn.Module):
                        targets[id(layer)] = layer      # لایه‌ها / بلوک‌ها
        for mod in targets.values():
            mod.register_forward_pre_hook(_step_hook)
            hooked += 1
    print(f"[CPU MODE] per-step thread/cancel hooks on {hooked} modules ({', '.join(roots) or 'none'})", flush=True)


# --- تابع جدید برش هوشمند با قابلیت افزایش داینامیک بازه جستجو ---
# ⚠️ اکشن پرو هم دقیقاً همین الگوریتم را (با همین عددها) پیش از ارسال اجرا
# می‌کند تا تکه‌ها را هم‌زمان به چند اسپیس بفرستد؛ اگر این‌جا عوض شد، آن‌جا
# (proxy_worker.py → _voice_split_points) هم باید عوض شود.
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


# =========================================================================
#  همان تابعِ اصلیِ قبلی (predict_voice_conversion)، فقط به‌جای gr.Progress
#  یک تابعِ گزارشِ پیشرفت می‌گیرد. خطاها مثلِ قبل بالا می‌روند و پروسه‌ی
#  Gradio همان پیامِ قبلی را («خطایی در حین پردازش رخ داد: ...») نشان می‌دهد.
#  تکه‌ها مثلِ قبل یکی‌یکی و به ترتیب ساخته و به هم وصل می‌شوند.
# =========================================================================
SR = 24000

# 🆕 سقفِ طولِ صدای مرجع (صدای مدل). توضیح: بالای فایل، بندِ ۱. ۰ = بدونِ سقف (رفتارِ قبلی).
REF_MAX_SECONDS = float(os.environ.get("VOICE_REF_MAX_SECONDS", "12"))
_REF_TOLERANCE_SECONDS = 0.5      # تا ۱۲٫۵ ثانیه دست نمی‌خورد (اکشن پرو خودش ۱۲ ثانیه می‌فرستد)
_REF_FRAME = 480                  # ۲۰ میلی‌ثانیه در ۲۴ کیلوهرتز (همان گامِ فریم‌های مدل)


def pick_reference_segment(y, sr=SR, max_seconds=REF_MAX_SECONDS):
    """صدای مرجعِ بلند ← «پُرصداترین» تکه‌ی حداکثر max_seconds ثانیه‌ای (بیشترین گفتار)، که درست
    پیش از شروعِ گفتار شروع و روی یک مکث تمام می‌شود (نه وسطِ کلمه). مرجعِ کوتاه‌تر از سقف بدونِ
    هیچ تغییری برمی‌گردد. (اکشن پرو همین را پیش از ارسال اجرا می‌کند: _voice_pick_reference.)"""
    y = np.asarray(y, dtype=np.float32)
    n = len(y)
    if max_seconds <= 0 or n <= int((max_seconds + _REF_TOLERANCE_SECONDS) * sr):
        return y
    frame = max(1, int(_REF_FRAME * sr / 24000))
    nf = n // frame
    win = max(1, int(max_seconds * sr / frame))
    if nf <= win:
        return y
    rms = np.sqrt(np.mean(np.square(y[: nf * frame].reshape(nf, frame), dtype=np.float64), axis=1))
    thr = max(0.008, 0.12 * float(np.percentile(rms, 95)))      # زیرِ این = مکث/سکوت
    active = (rms > thr).astype(np.int64)
    csum = np.concatenate([[0], np.cumsum(active)])
    step = max(1, int(0.5 * sr / frame))
    starts = np.arange(0, nf - win + 1, step)
    counts = csum[starts + win] - csum[starts]
    best = int(starts[int(np.argmax(counts))])                  # بیشترین گفتار؛ در تساوی، زودترین
    sec = max(1, int(sr / frame))
    # شروع: آخرین فریمِ ساکت در ±۱ ثانیه‌ی دورِ بهترین شروع (یعنی درست پیش از گفتار)
    h_lo, h_hi = max(0, best - sec), min(nf - win, best + sec) + 1
    s0 = h_lo + _last_quiet(rms[h_lo:h_hi], thr)
    # پایان: آخرین فریمِ ساکت در ۳ ثانیه‌ی آخرِ پنجره (تا جای ممکن بلند، روی مکث)
    t_lo, t_hi = max(s0 + 1, s0 + win - 3 * sec), min(nf, s0 + win)
    e0 = (t_lo + _last_quiet(rms[t_lo:t_hi], thr)) if t_hi > t_lo else t_hi
    e0 = max(e0, s0 + 1)
    return y[s0 * frame: e0 * frame]


def _last_quiet(values, thr):
    """اندیسِ آخرین فریمِ ساکت (زیرِ thr)؛ اگر نبود، آخرین فریمی که هم‌اندازه‌ی کم‌صداترین است."""
    values = np.asarray(values, dtype=np.float64)
    if len(values) == 0:
        return 0
    quiet = np.nonzero(values <= thr)[0]
    if len(quiet):
        return int(quiet[-1])
    low = float(values.min())
    return int(np.nonzero(values <= low * 1.05 + 1e-6)[0][-1])


def run_voice_conversion(source_audio_path, ref_audio_path, progress=None, job=None):
    if progress is None:
        progress = lambda value, desc="": None  # noqa: E731
    if job is None:
        job = _Job()

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

    # پردازش صوت اصلی
    content_data, _ = librosa.load(source_audio_path, sr=SR)

    # پردازش صوت مرجع (🆕 مرجعِ بلندتر از سقف ← پُرصداترین تکه‌اش؛ کوتاه‌تر دست نمی‌خورد)
    ref_data, _ = librosa.load(ref_audio_path, sr=SR)
    ref_full = len(ref_data) / SR
    ref_data = pick_reference_segment(ref_data, SR)
    if len(ref_data) < int(ref_full * SR):
        print(f"[CPU MODE] reference {ref_full:.1f}s -> using its best {len(ref_data) / SR:.1f}s "
              f"(VOICE_REF_MAX_SECONDS={REF_MAX_SECONDS:g})", flush=True)
    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)

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

        generated_chunks = []
        pipeline = get_pipeline()

        # شبیه‌سازی قطعه به قطعه (مثلِ قبل، یکی‌یکی و به ترتیب)
        total_chunks = len(chunks)
        print(f"[CPU MODE] {total_chunks} chunks, {len(content_data)/SR:.1f}s audio "
              f"(ref {len(ref_data)/SR:.1f}s)", flush=True)
        for idx, chunk in enumerate(chunks):
            job.check()
            progress(idx / total_chunks, f"در حال پردازش روی CPU: قطعه {idx+1} از {total_chunks}")
            chunk_started = time.time()
            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)

            try:
                _sync_threads(job)
                with torch.inference_mode():
                    gen = pipeline.inference_fm(
                        src_wav_path=temp_c_chunk,
                        timbre_ref_wav_path=temp_r,
                        flow_matching_steps=32,
                    )
            finally:
                if os.path.exists(temp_c_chunk):
                    os.remove(temp_c_chunk)

            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 total_chunks > 1:
                print(f"   [CPU MODE] chunk {idx+1}/{total_chunks}: {len(chunk)/SR:.1f}s in "
                      f"{time.time() - chunk_started:.1f}s ({job.threads} threads)", flush=True)

        job.check()
        progress(1.0, "در حال ادغام و ذخیره خروجی")
        # ادغام قطعات خروجی (به ترتیب)
        final_gen = np.concatenate(generated_chunks, axis=0)

        # ذخیره نهایی
        save_audio_pcm16(final_gen, out_path, SR)
    finally:
        if os.path.exists(temp_r):
            os.remove(temp_r)

    return out_path


# =========================================================================
#  سرویس‌دهی به پروسه‌ی Gradio (فقط داخلِ همین کانتینر، با کلیدِ مخفی)
# =========================================================================
_job_slots = threading.BoundedSemaphore(VOICE_CONCURRENT_JOBS)
_state = {"busy": False, "warm": False, "jobs": 0, "cancelled": 0, "started_at": time.time()}


def _write_state(path, **extra):
    if not path:
        return
    data = {"pid": os.getpid(), "ready": True, "busy": _state["busy"], "warm": _state["warm"],
            "jobs": _state["jobs"], "cancelled": _state["cancelled"], "threads": CPU_THREADS,
            "concurrency": VOICE_CONCURRENT_JOBS, "active": _active_jobs, "ts": time.time()}
    data.update(extra)
    try:
        tmp = f"{path}.tmp"
        with open(tmp, "w", encoding="utf-8") as f:
            json.dump(data, f)
        os.replace(tmp, path)
    except Exception:
        pass


def _cleanup_old_outputs(max_age_seconds=2 * 3600):
    """خروجی‌های قدیمیِ پوشه‌ی wav پاک می‌شوند (Gradio هر خروجی را لحظه‌ی تحویل در
    کشِ خودش کپی می‌کند، پس بعد از چند ساعت این‌جا دیگر لازم نیستند)."""
    try:
        cutoff = time.time() - max_age_seconds
        for name in os.listdir("wav"):
            path = os.path.join("wav", name)
            try:
                if os.path.isfile(path) and os.path.getmtime(path) < cutoff:
                    os.remove(path)
            except Exception:
                pass
    except Exception:
        pass


def _watch_client(conn, job):
    """اگر درخواست‌دهنده (پروسه‌ی Gradio) اتصال را بست -- یعنی دیگر کسی منتظرِ این نتیجه نیست
    (کاربر/اکشن پرو رفت یا زمانش تمام شد) -- کار همان لحظه لغو می‌شود."""
    while not job.done.is_set() and not job.cancelled.is_set():
        try:
            ready = conn.poll(0.5)
        except (OSError, EOFError, ValueError):
            job.cancel("the requester closed the connection")
            return
        if not ready:
            continue
        try:
            msg = conn.recv()
        except (OSError, EOFError, ValueError):
            job.cancel("the requester is no longer waiting")
            return
        if (msg or {}).get("op") == "cancel":
            job.cancel("cancel requested")
            return


def _serve(sock_path, state_path):
    from multiprocessing.connection import Listener

    authkey = bytes.fromhex(os.environ.get("VOICE_ENGINE_AUTHKEY", ""))
    if not authkey:
        raise SystemExit("VOICE_ENGINE_AUTHKEY is missing")
    try:
        os.remove(sock_path)
    except FileNotFoundError:
        pass
    listener = Listener(sock_path, family="AF_UNIX", authkey=authkey)
    _write_state(state_path)
    print(f"🎙️ [Voice Engine] آماده است (pid={os.getpid()}, CPU threads={CPU_THREADS}, "
          f"تا {VOICE_CONCURRENT_JOBS} کارِ هم‌زمان).", flush=True)

    def handle(conn):
        global _active_jobs
        try:
            msg = conn.recv()
            op = (msg or {}).get("op")
            if op == "ping":
                conn.send({"op": "pong", "busy": _state["busy"], "warm": _state["warm"], "jobs": _state["jobs"],
                           "active": _active_jobs, "concurrency": VOICE_CONCURRENT_JOBS})
                return
            if op != "convert":
                conn.send({"op": "error", "message": f"unknown op: {op}"})
                return
            job = _Job()
            with _job_slots:
                with _jobs_lock:
                    _active_jobs += 1
                    _state["busy"] = True
                _write_state(state_path)
                threading.Thread(target=_watch_client, args=(conn, job), daemon=True,
                                 name="voice-client-watch").start()
                _tls.job = job
                started = time.time()
                send_lock = threading.Lock()

                def report(value, desc=""):
                    if job.cancelled.is_set():
                        return
                    try:
                        with send_lock:
                            conn.send({"op": "progress", "value": float(value), "desc": str(desc)})
                    except Exception:
                        job.cancel("the requester is no longer waiting")

                try:
                    out_path = run_voice_conversion(msg.get("source"), msg.get("ref"), progress=report, job=job)
                    _state["warm"] = True
                    job.check()
                    try:
                        with send_lock:
                            conn.send({"op": "done", "path": os.path.abspath(out_path)})
                    except Exception:
                        job.cancel("the requester left right before the result")
                        raise _JobCancelled(job.reason)
                    job.done.set()
                    together = " (هم‌زمان با کارِ دیگر)" if job.max_parallel > 1 else ""
                    print(f"✅ [Voice Engine] تبدیل در {time.time() - started:.1f} ثانیه کامل شد{together}.", flush=True)
                except _JobCancelled:
                    _state["cancelled"] += 1
                    print(f"⏹️ [Voice Engine] کار بعد از {time.time() - started:.1f} ثانیه لغو شد "
                          f"({job.reason}); CPU آزاد شد.", flush=True)
                except Exception as e:
                    if job.cancelled.is_set():
                        _state["cancelled"] += 1
                        print(f"⏹️ [Voice Engine] کار بعد از {time.time() - started:.1f} ثانیه لغو شد "
                              f"({job.reason}); CPU آزاد شد.", flush=True)
                    else:
                        print(f"❌ [Voice Engine] خطا: {type(e).__name__}: {e}", flush=True)
                        try:
                            with send_lock:
                                conn.send({"op": "error", "message": str(e)})
                        except Exception:
                            pass
                finally:
                    job.done.set()
                    _tls.job = None
                    with _jobs_lock:
                        _active_jobs = max(0, _active_jobs - 1)
                        _state["busy"] = _active_jobs > 0
                        _state["jobs"] += 1
                        jobs_done = _state["jobs"]
                    _write_state(state_path)
                    if jobs_done % 25 == 0:
                        _cleanup_old_outputs()
        except EOFError:
            pass
        except Exception as e:
            print(f"⚠️ [Voice Engine] اتصال: {type(e).__name__}: {e}", flush=True)
        finally:
            try:
                conn.close()
            except Exception:
                pass

    while True:
        try:
            conn = listener.accept()
        except Exception as e:
            print(f"⚠️ [Voice Engine] accept: {type(e).__name__}: {e}", flush=True)
            time.sleep(0.5)
            continue
        threading.Thread(target=handle, args=(conn,), daemon=True).start()


def _exit_with_parent(parent_pid):
    """اگر پروسه‌ی اصلیِ اسپیس بسته شد، این موتور هم بسته شود (پروسه‌ی یتیم نماند)."""
    while True:
        time.sleep(5)
        if os.getppid() != parent_pid:
            print("🛑 [Voice Engine] پروسه‌ی اصلی بسته شد؛ موتور صدا هم بسته می‌شود.", flush=True)
            os._exit(0)


def main():
    sock_path = sys.argv[sys.argv.index("--serve") + 1]
    state_path = sys.argv[sys.argv.index("--state") + 1] if "--state" in sys.argv else ""
    threading.Thread(target=_exit_with_parent, args=(os.getppid(),), daemon=True).start()
    signal.signal(signal.SIGTERM, lambda *_: os._exit(0))

    # گرم‌کردنِ مدل همان اول (قبلاً اولین درخواست این زمان را می‌داد). اگر نشد،
    # مثلِ قبل در اولین درخواست دوباره ساخته می‌شود.
    try:
        t0 = time.time()
        get_pipeline()
        _state["warm"] = True
        print(f"🔥 [Voice Engine] مدل Vevo در {time.time() - t0:.1f} ثانیه روی CPU آماده شد.", flush=True)
        _warm_up()
    except Exception as e:
        print(f"⚠️ [Voice Engine] گرم‌کردنِ مدل ناموفق بود ({type(e).__name__}: {e})؛ "
              f"در اولین درخواست دوباره امتحان می‌شود.", flush=True)
    _serve(sock_path, state_path)


def _warm_up():
    """🆕 یک تبدیلِ آزمایشیِ خیلی کوتاه پیش از اولین درخواست (صدای ساختگیِ ۲ ثانیه‌ای): اولین
    اجرای مدل و librosa کارهای یک‌باره‌ای دارد (کامپایل/آماده‌سازی) که قبلاً به اولین کاربر می‌افتاد
    (همان ۴۰ ثانیه‌ایِ اولِ لاگ برای یک تکه‌ی ۶ ثانیه‌ای). خروجی بلافاصله پاک می‌شود."""
    t0 = time.time()
    tag = uuid.uuid4().hex
    src, ref, out = f"wav/{tag}_warm_src.wav", f"wav/{tag}_warm_ref.wav", None
    try:
        t = np.arange(int(2.0 * SR)) / SR
        tone = (0.2 * np.sin(2 * np.pi * 180 * t) * np.hanning(len(t))).astype(np.float32)
        sf.write(src, tone, SR, subtype="PCM_16")
        sf.write(ref, tone[::-1].copy(), SR, subtype="PCM_16")
        out = run_voice_conversion(src, ref)
        print(f"🔥 [Voice Engine] گرم‌کردنِ اولیه در {time.time() - t0:.1f} ثانیه انجام شد.", flush=True)
    except Exception as e:
        print(f"⚠️ [Voice Engine] گرم‌کردنِ اولیه انجام نشد ({type(e).__name__}: {e}); مشکلی نیست.", flush=True)
    finally:
        for path in (src, ref, out):
            try:
                if path and os.path.exists(path):
                    os.remove(path)
            except Exception:
                pass


if __name__ == "__main__":
    main()