# ==========================================================================
# FILE: voice_clone.py  — بخشِ «کلون کردن صدا» داخلِ اسپیسِ ترکیبی
# --------------------------------------------------------------------------
# این فایل در پروسه‌ی اصلیِ Gradio اجرا می‌شود و سه کار دارد:
#   ۱) موتورِ صدا (voice_engine.py) را در یک پروسه‌ی جدا روشن نگه می‌دارد
#      (اگر به هر دلیلی بسته شد، خودکار دوباره روشن می‌شود).
#   ۲) تبِ «کلون کردن صدا» را می‌سازد: دقیقاً همان رابطِ اسپیسِ صدای قبلی
#      (دو فایل صوتی + دکمه‌ی شروع + خروجی).
#   ۳) دو مسیرِ API را ثبت می‌کند (برای اکشن پرو و هر کلاینتِ دیگر):
#        /predict_voice_conversion   ← همان نامِ قبلی؛ ورودی: [صدای منبع، صدای مرجع]
#                                      خروجی: یک فایل WAV (۲۴ کیلوهرتز، مونو)
#        /voice_status               ← وضعیتِ فقط بخشِ صدا (بدونِ صف، فوری):
#                                      آماده است؟ چند کارِ صدا در حال اجرا/منتظر؟ چند جای
#                                      هم‌زمان (slots)؟ و شماره‌ی تابعِ صدا (fn_index).
#
# چرا /voice_status لازم است؟ «/gradio_api/queue/status» گرادیو صفِ «همه‌ی»
# بخش‌ها را با هم می‌شمارد؛ در اسپیسِ ترکیبی یک آهنگِ منتظر باعث می‌شد اکشن پرو
# بخشِ صدا را هم «مشغول» ببیند. این مسیر فقط صفِ تبدیلِ صدا را گزارش می‌دهد.
#
# صدا فقط روی CPU (پروسه‌ی جدا، بدونِ ZeroGPU)؛ آهنگ و آواتار روی GPU (مثلِ قبل).
# 🆕 هم‌زمانی: تا «دو» درخواستِ صدا با هم (VOICE_CONCURRENCY_LIMIT=2). اکشن پرو کارِ دوم را فقط
#   وقتی می‌فرستد که هیچ اسپیسِ صدای خالیِ دیگری نمانده؛ کارِ تنها همه‌ی هسته‌ها را می‌گیرد
#   (توضیح بالای voice_engine.py).
# 🆕 کارِ بی‌صاحب لغو می‌شود: اگر درخواست‌دهنده رفت (اتصالش قطع شد)، تبدیل همان لحظه متوقف
#   می‌شود و CPU و جایش برای درخواستِ بعدی آزاد می‌شود.
# ==========================================================================
import os
import sys
import json
import time
import secrets
import threading
import subprocess

import gradio as gr

APP_DIR = os.path.dirname(os.path.abspath(__file__))

# نامِ مسیرِ API تبدیلِ صدا -- عمداً همان نامِ اسپیسِ صدای قبلی است تا هر
# کلاینتی که قبلاً با نام صدا می‌زد، بدونِ تغییر کار کند.
VOICE_API_NAME = "predict_voice_conversion"
VOICE_STATUS_API_NAME = "voice_status"
VOICE_CONCURRENCY_ID = "voice_clone"
# 🆕 حداکثر چند تبدیلِ هم‌زمان در همین اسپیس (۲ = «کارِ دوم هم‌زمان»؛ ۱ = رفتارِ قبلی).
VOICE_CONCURRENCY_LIMIT = max(1, min(4, int(os.environ.get("VOICE_CONCURRENCY_LIMIT", "2") or 2)))

# اگر موتورِ صدا هنوز در حالِ بالا آمدن است (اولین روشن‌شدن: دانلودِ مدل)،
# درخواست حداکثر این‌قدر منتظر می‌ماند.
VOICE_ENGINE_WAIT_SECONDS = float(os.environ.get("VOICE_ENGINE_WAIT_SECONDS", "900"))
# سقفِ زمانِ یک تبدیل (همان سقفی که اکشن پرو هم دارد).
VOICE_JOB_TIMEOUT_SECONDS = float(os.environ.get("VOICE_JOB_TIMEOUT_SECONDS", "3600"))


class RequesterGone(RuntimeError):
    """درخواست‌دهنده دیگر منتظر نیست (اتصالش قطع شد)؛ تبدیل لغو شد. خطای واقعی نیست."""


class VoiceEngine:
    """مدیرِ پروسه‌ی جداگانه‌ی موتورِ صدا (روشن‌کردن، نظارت، ارسالِ کار)."""

    def __init__(self):
        tag = f"{os.getpid()}"
        self.sock_path = f"/tmp/alpha_voice_engine_{tag}.sock"
        self.state_path = f"/tmp/alpha_voice_engine_{tag}.json"
        self.authkey = secrets.token_bytes(32)
        self.proc = None
        self.lock = threading.Lock()
        self.started = False
        self.restarts = 0
        self.last_exit_code = None
        self.active_jobs = 0          # کارهایی که همین حالا داخلِ تابعِ صدا هستند
        self.jobs_lock = threading.Lock()
        self._stop = False

    # ------------------------------------------------------------------
    def _spawn(self):
        env = dict(os.environ)
        env["VOICE_ENGINE_AUTHKEY"] = self.authkey.hex()
        env["PYTHONUNBUFFERED"] = "1"
        env["VOICE_CONCURRENT_JOBS"] = str(VOICE_CONCURRENCY_LIMIT)
        for path in (self.sock_path, self.state_path):
            try:
                os.remove(path)
            except FileNotFoundError:
                pass
        # (VOICE_ENGINE_SCRIPT فقط برای آزمایش است؛ در اسپیس ست نکنید.)
        script = os.environ.get("VOICE_ENGINE_SCRIPT") or os.path.join(APP_DIR, "voice_engine.py")
        self.proc = subprocess.Popen(
            [sys.executable, "-u", script, "--serve", self.sock_path, "--state", self.state_path],
            cwd=APP_DIR, env=env,
        )
        print(f"🎙️ [کلون صدا] موتورِ صدا روی CPU در حالِ بالا آمدن است (pid={self.proc.pid})...", flush=True)

    def start(self):
        with self.lock:
            if self.started:
                return
            self.started = True
            self._spawn()
        threading.Thread(target=self._watch, daemon=True, name="voice-engine-watch").start()

    def _watch(self):
        fails = 0
        while not self._stop:
            proc = self.proc
            code = proc.wait()
            if self._stop:
                return
            self.last_exit_code = code
            fails += 1
            self.restarts += 1
            wait_s = min(60, 2 ** min(fails, 6))
            print(f"⚠️ [کلون صدا] موتورِ صدا بسته شد (کد {code})؛ {wait_s} ثانیه دیگر دوباره روشن می‌شود.",
                  flush=True)
            time.sleep(wait_s)
            with self.lock:
                self._spawn()
            # اگر موتور چند دقیقه سالم ماند، شمارنده‌ی خطا صفر می‌شود.
            for _ in range(300):
                if self._stop or self.proc.poll() is not None:
                    break
                time.sleep(1)
            else:
                fails = 0

    # ------------------------------------------------------------------
    def _state(self):
        try:
            with open(self.state_path, "r", encoding="utf-8") as f:
                return json.load(f)
        except Exception:
            return {}

    def is_ready(self):
        proc = self.proc
        if proc is None or proc.poll() is not None:
            return False
        st = self._state()
        return bool(st.get("ready")) and st.get("pid") == proc.pid and os.path.exists(self.sock_path)

    def is_warm(self):
        return bool(self._state().get("warm")) and self.is_ready()

    def wait_ready(self, timeout, progress=None):
        deadline = time.time() + max(0.0, timeout)
        shown = False
        while time.time() < deadline:
            if self.is_ready():
                return True
            if progress is not None and not shown:
                shown = True
                try:
                    progress(0.0, desc="در حال آماده‌سازی مدل صدا روی CPU...")
                except Exception:
                    pass
            time.sleep(1.0)
        return self.is_ready()

    def slots(self):
        """چند تبدیلِ هم‌زمان (۲ = «کارِ دوم هم‌زمان»)؛ کمترینِ تنظیمِ این‌جا و موتور."""
        try:
            engine = int(self._state().get("concurrency") or VOICE_CONCURRENCY_LIMIT)
        except Exception:
            engine = VOICE_CONCURRENCY_LIMIT
        return max(1, min(VOICE_CONCURRENCY_LIMIT, engine))

    # ------------------------------------------------------------------
    def convert(self, source_path, ref_path, progress=None, alive=None):
        """یک تبدیل: مسیرِ فایلِ خروجی را برمی‌گرداند یا خطا (با متنِ خودِ موتور).
        alive: تابعی که می‌گوید درخواست‌دهنده هنوز منتظر است یا نه (اگر نه، کار لغو می‌شود)."""
        return self._request({"op": "convert", "source": source_path, "ref": ref_path}, progress, alive)

    def _request(self, payload, progress=None, alive=None):
        from multiprocessing.connection import Client

        if not self.wait_ready(VOICE_ENGINE_WAIT_SECONDS, progress=progress):
            raise RuntimeError("موتور تبدیل صدا هنوز آماده نشده است؛ چند دقیقه‌ی دیگر دوباره تلاش کنید.")
        try:
            conn = Client(self.sock_path, family="AF_UNIX", authkey=self.authkey)
        except Exception as e:
            raise RuntimeError(f"اتصال به موتور تبدیل صدا ناموفق بود ({type(e).__name__}).")
        try:
            conn.send(payload)
            deadline = time.time() + VOICE_JOB_TIMEOUT_SECONDS
            while True:
                remaining = deadline - time.time()
                if remaining <= 0:
                    raise RuntimeError("تبدیل صدا بیش از حد طول کشید.")
                if alive is not None and not _safe_alive(alive):
                    # درخواست‌دهنده رفته؛ موتور با بسته شدنِ همین اتصال (finally) کار را لغو می‌کند.
                    try:
                        conn.send({"op": "cancel"})
                    except Exception:
                        pass
                    raise RequesterGone("درخواست‌دهنده دیگر منتظر نبود؛ تبدیل لغو شد.")
                if not conn.poll(min(remaining, 1.0)):
                    if self.proc is None or self.proc.poll() is not None:
                        raise RuntimeError("موتور تبدیل صدا وسطِ کار بسته شد؛ دوباره تلاش کنید.")
                    continue
                msg = conn.recv()
                op = (msg or {}).get("op")
                if op == "progress":
                    if progress is not None:
                        try:
                            progress(float(msg.get("value") or 0.0), desc=str(msg.get("desc") or ""))
                        except Exception:
                            pass
                    continue
                if op == "done":
                    return msg.get("path")
                if op == "error":
                    raise RuntimeError(str(msg.get("message") or "خطای نامشخص"))
        except EOFError:
            raise RuntimeError("موتور تبدیل صدا وسطِ کار بسته شد؛ دوباره تلاش کنید.")
        finally:
            try:
                conn.close()
            except Exception:
                pass


ENGINE = VoiceEngine()


def _safe_alive(alive):
    try:
        return bool(alive())
    except Exception:
        return True


def _request_alive(request):
    """آیا کسی هنوز منتظرِ همین درخواست است؟ وقتی اتصالِ درخواست‌دهنده (اکشن پرو یا صفحه)
    قطع شود، Gradio رویدادش را «مرده» علامت می‌زند و دیگر هیچ نتیجه‌ای به او نمی‌رسد؛ پس
    ادامه‌ی تبدیل فقط CPU و جای کارِ بعدی را هدر می‌داد. اگر معلوم نشد، «زنده» فرض می‌شود."""
    session_hash = getattr(request, "session_hash", None)
    queue = getattr(_VOICE.get("demo"), "_queue", None)
    if not session_hash or queue is None:
        return lambda: True

    def _alive():
        found = False
        for job_set in list(getattr(queue, "active_jobs", None) or []):
            for event in list(job_set or []):
                if getattr(event, "session_hash", None) != session_hash:
                    continue
                cid = getattr(event, "concurrency_id", None) or getattr(getattr(event, "fn", None), "concurrency_id", None)
                if cid is not None and cid != VOICE_CONCURRENCY_ID:
                    continue
                found = True
                if getattr(event, "alive", True):
                    return True
        return not found
    return _alive


def start_engine():
    """در ابتدای برنامه صدا زده می‌شود تا مدلِ صدا هم‌زمان با مدل‌های آهنگ بالا بیاید."""
    ENGINE.start()


# ==========================================================================
#  تابعِ تبدیلِ صدا -- همان امضا و همان پیام‌های اسپیسِ صدای قبلی
#  (بدونِ @spaces.GPU، یعنی کاملاً روی CPU)
# ==========================================================================
def predict_voice_conversion(source_audio_path, ref_audio_path, request: gr.Request = None, progress=gr.Progress()):
    if source_audio_path is None or ref_audio_path is None:
        raise gr.Error("لطفاً هر دو فایل صوتی (اصلی و مرجع) را وارد کنید.")

    with ENGINE.jobs_lock:
        ENGINE.active_jobs += 1
    try:
        return ENGINE.convert(source_audio_path, ref_audio_path, progress=progress, alive=_request_alive(request))
    except gr.Error:
        raise
    except RequesterGone as e:
        # کسی منتظرِ این نتیجه نیست (پیام به هیچ‌کس نمی‌رسد)؛ فقط یک خطِ کوتاه در لاگ، بدونِ traceback.
        print("⏹️ [کلون صدا] درخواست‌دهنده رفت؛ تبدیل لغو و CPU برای کارِ بعدی آزاد شد.", flush=True)
        raise gr.Error(str(e), print_exception=False)
    except Exception as e:
        raise gr.Error(f"خطایی در حین پردازش رخ داد: {str(e)}")
    finally:
        with ENGINE.jobs_lock:
            ENGINE.active_jobs = max(0, ENGINE.active_jobs - 1)


# ==========================================================================
#  وضعیتِ فقط بخشِ صدا (برای اکشن پرو): /gradio_api/run/voice_status
# ==========================================================================
_VOICE = {"demo": None, "fn_index": None}


def _voice_queue_counts():
    """(منتظر، در حال اجرا) فقط برای تابعِ صدا، از صفِ خودِ گرادیو."""
    demo = _VOICE.get("demo")
    try:
        queue = getattr(demo, "_queue", None)
        event_queue = queue.event_queue_per_concurrency_id.get(VOICE_CONCURRENCY_ID) if queue else None
        if event_queue is not None:
            return len(event_queue.queue), int(event_queue.current_concurrency)
    except Exception:
        pass
    running = min(ENGINE.active_jobs, VOICE_CONCURRENCY_LIMIT)
    return max(0, ENGINE.active_jobs - running), running


def voice_status() -> dict:
    """وضعیتِ لحظه‌ایِ بخشِ کلونِ صدا. بدونِ صف و فوری؛ هیچ کاری اجرا نمی‌کند."""
    waiting, running = _voice_queue_counts()
    return {
        "service": "voice",
        "api_name": VOICE_API_NAME,
        "fn_index": _VOICE.get("fn_index"),
        "ready": ENGINE.is_ready(),
        "warm": ENGINE.is_warm(),
        "running": running,
        "waiting": waiting,
        "concurrency": VOICE_CONCURRENCY_LIMIT,
        # 🆕 چند تبدیلِ هم‌زمان (اکشن پرو کارِ دوم را فقط وقتی می‌فرستد که هیچ اسپیسِ خالیِ دیگری نیست)
        "slots": ENGINE.slots(),
        "music": True,
        "avatar": True,
        "version": 3,
    }


# ==========================================================================
#  تبِ «کلون کردن صدا» -- همان رابطِ اسپیسِ صدای قبلی
# ==========================================================================
def build_voice_tab(demo):
    gr.Markdown("# سامانه تبدیل صدا (Vevo Voice Conversion) - نسخه CPU")

    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")

    dep = submit_btn.click(
        fn=predict_voice_conversion,
        inputs=[source_input, ref_input],
        outputs=audio_output,
        api_name=VOICE_API_NAME,
        concurrency_limit=VOICE_CONCURRENCY_LIMIT,
        concurrency_id=VOICE_CONCURRENCY_ID,
    )

    def _fn_index_of(event):
        try:
            idx = getattr(getattr(event, "fn", None), "_id", None)
            if idx is None and isinstance(event, dict):
                idx = event.get("id")
            return idx
        except Exception:
            return None

    _VOICE["demo"] = demo
    _VOICE["fn_index"] = _fn_index_of(dep)

    # مسیرِ وضعیت: بدونِ صف (queue=False) تا حتی وقتی صفِ آهنگ پُر است فوری جواب بدهد.
    gr.api(voice_status, api_name=VOICE_STATUS_API_NAME, queue=False)
    return {"source_input": source_input, "ref_input": ref_input,
            "submit_btn": submit_btn, "audio_output": audio_output}