# -*- coding: utf-8 -*-
# app.py — Human-grade Multilingual STT for Hugging Face Spaces
# خود-بهینه‌ساز + مود مخصوص فارسی + سازگاری خودکار با نسخه‌های مختلف Gradio

import os, sys, subprocess, importlib, io, time, tempfile, math, re, inspect
import numpy as np

def ensure(pkg, import_name=None, extra=None, quiet=False):
    try:
        return importlib.import_module(import_name or pkg)
    except ImportError:
        print(f"[setup] installing {pkg} ...")
        cmd = [sys.executable, "-m", "pip", "install", "--upgrade", pkg]
        if extra: cmd += extra
        if quiet: cmd += ["-q"]
        subprocess.check_call(cmd)
        return importlib.import_module(import_name or pkg)

gradio = ensure("gradio", "gradio")
fw = ensure("faster-whisper", "faster_whisper")
soundfile = ensure("soundfile", "soundfile")
# اختیاری برای کاهش نویز دقیق‌تر و ریسَمپل دقیق
try:
    librosa = ensure("librosa", "librosa"); noisereduce = ensure("noisereduce", "noisereduce", quiet=True)
except Exception:
    librosa = None; noisereduce = None

try:
    psutil = ensure("psutil", "psutil", quiet=True)
except Exception:
    psutil = None

import gradio as gr
from faster_whisper import WhisperModel
import soundfile as sf

# ------------------------- تشخیص منابع -------------------------
def detect_ram_gb():
    try:
        if psutil:
            return max(1, int(psutil.virtual_memory().total / (1024**3)))
    except Exception:
        pass
    return 8

def detect_cpu_threads():
    try:
        return max(1, min(8, os.cpu_count() or 4))
    except Exception:
        return 4

RAM_GB = detect_ram_gb()
CPU_THREADS = detect_cpu_threads()

def default_model_by_ram(ram_gb:int) -> str:
    if ram_gb < 4: return "tiny"
    if ram_gb < 8: return "base"
    if ram_gb < 12: return "small"
    return "medium"

DEFAULT_MODEL = default_model_by_ram(RAM_GB)
DEFAULT_COMPUTE = "int8_float16"
DEFAULT_CHUNK = 30
DEFAULT_BEAM = 3
DEFAULT_VAD = True

# ------------------------- تقویت صوت -------------------------
def limiter_soft(x: np.ndarray, gain: float = 10.0) -> np.ndarray:
    if x.size == 0: return x
    peak = max(np.max(np.abs(x)), 1e-6)
    y = (x / peak) * gain
    return np.tanh(y)

def enhance_audio(wav: np.ndarray, sr: int, boost: float = 10.0, precise: bool = False):
    if precise and (librosa is not None) and (noisereduce is not None):
        pad = min(len(wav), sr // 2)
        try:
            den = noisereduce.reduce_noise(y=wav, sr=sr, y_noise=wav[:pad] if pad > 0 else None, stationary=False)
        except Exception:
            den = wav
        return limiter_soft(den, gain=boost)
    else:
        return limiter_soft(wav, gain=boost)

def load_audio_to_float32(file_path: str):
    data, sr = sf.read(file_path, dtype="float32", always_2d=False)
    if data.ndim == 2:
        data = data.mean(axis=1)
    return data.astype(np.float32), sr

def to_pcm16_wav_file(wav: np.ndarray, sr: int) -> str:
    if (librosa is not None) and (sr != 16000):
        try:
            wav = librosa.resample(wav, orig_sr=sr, target_sr=16000); sr = 16000
        except Exception:
            pass
    tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    sf.write(tmp.name, wav, sr, subtype="PCM_16")
    tmp.close()
    return tmp.name

# ------------------------- مدل -------------------------
def build_model(size: str, compute_type: str, threads: int):
    return WhisperModel(
        size,
        device="auto",
        compute_type=compute_type,
        cpu_threads=int(threads),
        download_root=os.getenv("HF_HOME", None)
    )

# ------------------------- مود مخصوص فارسی -------------------------
FA_COMMON_FIXES = {
    "می باشد": "است",
    "میشود": "می‌شود",
    "میشه": "می‌شود",
    "ها ی": "های",
    " می خواهم": " می‌خواهم",
    " نمی توان": " نمی‌توان",
    " نمیخواهم": " نمی‌خواهم",
    " میشود": " می‌شود",
    " نمی شود": " نمی‌شود",
    " های ": "‌های ",
}

FA_COMMA_HINTS = [
    "البته","اما","ولی","یعنی","مثلا","در نتیجه","در نهایت","در کل","از طرفی","به علاوه","در واقع"
]

FA_HALFSPACE_PATTERNS = [
    (r"\b(می)\s+(شود|کنم|کنی|کند|کنیم|کنید|کنند)\b", r"\1‌\2"),
    (r"\b(نمی)\s+(شود|خواهم|خواهی|خواهد|خواهیم|خواهید|خواهند|کنم|کنی|کند|کنیم|کنید|کنند)\b", r"\1‌\2"),
    (r"\b(تر|ترین)\b", r"‌\1"),
]

def fa_normalize_spaces(text: str) -> str:
    s = text
    s = re.sub(r"\s+", " ", s).strip()
    s = s.replace(" ,", ",").replace(" .", ".").replace(" !", "!").replace(" ؟", "؟").replace(" ؛", "؛").replace(" :", ":")
    s = re.sub(r"\s+،", "،", s)
    s = re.sub(r"\s+%", "%", s)
    return s

def fa_insert_commas(text: str) -> str:
    s = text
    for hint in FA_COMMA_HINTS:
        s = re.sub(rf"(\s{hint})\s", rf"\1، ", s)
    return s

def fa_sentenceize(text: str) -> str:
    s = text.strip()
    if not s: return s
    if len(s) > 40 and not s.endswith((".", "؟", "!", "؛")):
        s += "."
    parts = re.split(r"([\.!\؟؛])", s)
    out = []
    buf = ""
    for p in parts:
        if p in [".","!","؟","؛"]:
            buf += p
            out.append(buf.strip())
            buf = ""
        else:
            if buf:
                buf += " " + p.strip()
            else:
                buf = p.strip()
    if buf.strip():
        out.append(buf.strip())
    return " ".join(out)

def fa_apply_halfspaces(text: str) -> str:
    s = " " + text + " "
    for pat, rep in FA_HALFSPACE_PATTERNS:
        s = re.sub(pat, rep, s)
    return s.strip()

def fa_apply_common_fixes(text: str) -> str:
    s = text
    for k, v in FA_COMMON_FIXES.items():
        s = s.replace(k, v)
    return s

def fa_polish(text: str) -> str:
    s = text
    s = fa_apply_common_fixes(s)
    s = fa_insert_commas(s)
    s = fa_sentenceize(s)
    s = fa_apply_halfspaces(s)
    s = fa_normalize_spaces(s)
    return s

def maybe_persian_polish(text: str, language: str) -> str:
    if (language or "").lower() in ["fa", "auto", ""]:
        return fa_polish(text)
    return text

# ------------------------- خود-بهینه‌سازی -------------------------
def autotune_params(ram_gb:int, base_model:str, compute_type:str, threads:int,
                    chunk_len:int, beam:int, target_rt:float=1.0):
    model = base_model; comp = compute_type; ch = int(chunk_len); bm = int(beam); th = int(threads)
    if ram_gb < 4:
        model = "tiny"; comp = "int8"; ch = min(ch, 20); bm = min(bm, 2); th = max(1, min(th, 4))
    elif ram_gb < 8:
        model = "base"; comp = "int8"; ch = min(ch, 25); bm = min(bm, 3); th = max(2, min(th, 6))
    elif ram_gb < 12:
        model = "small"; comp = "int8_float16"; ch = min(ch, 30); bm = min(bm, 3); th = max(3, min(th, 8))
    else:
        model = "small" if base_model in ["tiny","base"] else base_model
        comp = "int8_float16"; ch = min(max(ch, 30), 45); bm = min(max(bm, 3), 5); th = max(4, min(th, 8))
    return dict(model=model, compute=comp, chunk=ch, beam=bm, threads=th, target_rt=target_rt)

def quick_benchmark(model: WhisperModel, wav_path: str, language: str,
                    chunk_len:int, beam:int, vad:bool) -> float:
    data, sr = sf.read(wav_path, dtype="float32", always_2d=False)
    if data.ndim == 2: data = data.mean(axis=1)
    max_len = int(sr * 12)
    data = data[:max_len] if len(data) > max_len else data
    tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    sf.write(tmp.name, data, sr, subtype="PCM_16"); tmp.close()
    t0 = time.time()
    segs, info = model.transcribe(
        tmp.name,
        language=None if language=="auto" else language,
        vad_filter=vad,
        chunk_length=int(chunk_len),
        beam_size=int(beam),
        vad_parameters=dict(min_silence_duration_ms=300),
        no_speech_threshold=0.6,
        compression_ratio_threshold=2.8
    )
    t1 = time.time()
    os.unlink(tmp.name)
    audio_sec = len(data)/sr if sr>0 else 12.0
    proc_sec = max(t1 - t0, 1e-6)
    return proc_sec / max(audio_sec, 1e-6)

def refine_params_by_rt(params:dict, rt:float):
    if rt > 1.2:
        params["beam"] = max(1, params["beam"] - 1)
        params["chunk"] = max(15, params["chunk"] - 5)
        if params["model"] == "medium": params["model"] = "small"
        elif params["model"] == "small": params["model"] = "base"
        elif params["model"] == "base": params["model"] = "tiny"
    elif rt < 0.7:
        params["beam"] = min(5, params["beam"] + 1)
        params["chunk"] = min(45, params["chunk"] + 5)
    return params

# ------------------------- SRT -------------------------
def make_srt(segments):
    def fmt_time(t):
        h = int(t // 3600); t -= 3600*h
        m = int(t // 60); s = t - 60*m
        return f"{h:02d}:{m:02d}:{s:06.3f}".replace(".", ",")
    lines = []
    for i, s in enumerate(segments, 1):
        txt = s.text.strip()
        lines.append(str(i))
        lines.append(f"{fmt_time(s.start)} --> {fmt_time(s.end)}")
        lines.append(txt)
        lines.append("")
    return "\n".join(lines)

# ------------------------- پردازش اصلی (استریم) -------------------------
def transcribe_stream(
    audio_file,
    lang="auto",
    auto_tune=True,
    model_size=DEFAULT_MODEL,
    compute_type=DEFAULT_COMPUTE,
    threads=CPU_THREADS,
    boost=10,
    precise_enhance=False,
    chunk_len=DEFAULT_CHUNK,
    beam_size=DEFAULT_BEAM,
    vad_filter=DEFAULT_VAD,
    timestamps=False
):
    if audio_file is None:
        yield "لطفاً یک فایل صوتی انتخاب کن."
        return

    wav_fp = audio_file.name
    wav, sr = load_audio_to_float32(wav_fp)
    wav = enhance_audio(wav, sr, boost=float(boost), precise=bool(precise_enhance))
    wav_path = to_pcm16_wav_file(wav, sr)

    try:
        p = autotune_params(RAM_GB, model_size, compute_type, int(threads), int(chunk_len), int(beam_size))
        if not auto_tune:
            p = dict(model=model_size, compute=compute_type, chunk=int(chunk_len), beam=int(beam_size),
                     threads=int(threads), target_rt=1.0)

        yield f"آماده‌سازی مدل ({p['model']} | {p['compute']} | threads={p['threads']})..."
        model = build_model(p["model"], p["compute"], p["threads"])

        if auto_tune:
            yield "بنچمارک کوتاه برای تنظیم سرعت/دقت..."
            rt = quick_benchmark(model, wav_path, lang, p["chunk"], p["beam"], bool(vad_filter))
            p = refine_params_by_rt(p, rt)
            yield f"تنظیمات نهایی: model={p['model']} compute={p['compute']} chunk={p['chunk']} beam={p['beam']} threads={p['threads']} (RT≈{rt:.2f})"
            if p["model"] != model_size:
                yield f"تعویض مدل به {p['model']} ..."
                model = build_model(p["model"], p["compute"], p["threads"])

        yield "شروع پردازش..."
        language = None if lang == "auto" else lang
        segs, info = model.transcribe(
            wav_path,
            language=language,
            vad_filter=bool(vad_filter),
            chunk_length=int(p["chunk"]),
            beam_size=int(p["beam"]),
            vad_parameters=dict(min_silence_duration_ms=300),
            no_speech_threshold=0.6,
            compression_ratio_threshold=2.8
        )

        segments_collected = []
        full_text = []
        for seg in segs:
            segments_collected.append(seg)
            piece = seg.text.strip()
            if piece:
                full_text.append(piece)
                current = " ".join(full_text)
                if timestamps:
                    yield f"[{seg.start:.2f}–{seg.end:.2f}] {piece}\n\n---\n{current}"
                else:
                    yield current

        final_text = " ".join(full_text).strip()
        final_text = maybe_persian_polish(final_text, (lang or "auto"))

        srt_text_raw = make_srt(segments_collected)
        if (lang or "").lower() in ["fa","auto"]:
            srt_text_raw = fa_normalize_spaces(srt_text_raw)

        txt_fp = tempfile.NamedTemporaryFile(suffix=".txt", delete=False).name
        with open(txt_fp, "w", encoding="utf-8") as f:
            f.write(final_text or "")

        srt_fp = tempfile.NamedTemporaryFile(suffix=".srt", delete=False).name
        with open(srt_fp, "w", encoding="utf-8") as f:
            f.write(srt_text_raw or "")

        yield f"\n---\nپایان. از تبِ خروجی فایل‌ها را دانلود کن."

    finally:
        if os.path.exists(wav_path):
            try: os.unlink(wav_path)
            except: pass

def transcribe_and_return_files(
    audio_file, lang, auto_tune, model_size, compute_type, threads, boost, precise_enhance, chunk_len, beam_size, vad_filter
):
    if audio_file is None:
        return None, None, "فایل انتخاب نشده."
    wav_fp = audio_file.name
    wav, sr = load_audio_to_float32(wav_fp)
    wav = enhance_audio(wav, sr, boost=float(boost), precise=bool(precise_enhance))
    wav_path = to_pcm16_wav_file(wav, sr)

    try:
        p = autotune_params(RAM_GB, model_size, compute_type, int(threads), int(chunk_len), int(beam_size))
        if not auto_tune:
            p = dict(model=model_size, compute=compute_type, chunk=int(chunk_len), beam=int(beam_size),
                     threads=int(threads), target_rt=1.0)
        model = build_model(p["model"], p["compute"], p["threads"])

        segs, info = model.transcribe(
            wav_path,
            language=None if lang=="auto" else lang,
            vad_filter=bool(vad_filter),
            chunk_length=int(p["chunk"]),
            beam_size=int(p["beam"]),
            vad_parameters=dict(min_silence_duration_ms=300),
            no_speech_threshold=0.6,
            compression_ratio_threshold=2.8
        )
        texts = [s.text.strip() for s in segs if s.text.strip()]
        final_text = " ".join(texts).strip()
        final_text = maybe_persian_polish(final_text, lang)

        srt_text = make_srt(segs)
        if (lang or "").lower() in ["fa","auto"]:
            srt_text = fa_normalize_spaces(srt_text)

        txt_fp = tempfile.NamedTemporaryFile(suffix=".txt", delete=False).name
        with open(txt_fp, "w", encoding="utf-8") as f:
            f.write(final_text or "")

        srt_fp = tempfile.NamedTemporaryFile(suffix=".srt", delete=False).name
        with open(srt_fp, "w", encoding="utf-8") as f:
            f.write(srt_text or "")

        return txt_fp, srt_fp, "آماده دانلود."
    except Exception as e:
        return None, None, f"خطا: {e}"
    finally:
        if os.path.exists(wav_path):
            try: os.unlink(wav_path)
            except: pass

# ------------------------- رابط Gradio -------------------------
LANGS = [
    ("Automatic (detect)", "auto"),
    ("فارسی", "fa"),
    ("English", "en"),
    ("العربية", "ar"),
    ("Türkçe", "tr"),
    ("Français", "fr"),
    ("Deutsch", "de"),
    ("Español", "es"),
    ("Русский", "ru"),
    ("中文", "zh")
]
MODELS = ["tiny", "base", "small", "medium"]
COMPUTE_TYPES = ["int8_float16", "int8", "float16", "float32"]

with gr.Blocks(title="Human-grade STT — Fast & Self-Tuning", css="#stream{white-space:pre-wrap;}") as demo:
    gr.Markdown("### تبدیل گفتار به متن چندزبانه — سریع، خود-بهینه‌ساز، با مود مخصوص فارسی")
    with gr.Row():
        with gr.Column(scale=2):
            audio = gr.Audio(label="فایل صوتی را آپلود کن یا ضبط کن", type="filepath", sources=["upload", "microphone"])
            lang = gr.Dropdown(choices=[l[1] for l in LANGS], value="auto", label="Language")
            auto_tune = gr.Checkbox(value=True, label="خود-بهینه‌سازی (پیشنهادی)")
            model_size = gr.Dropdown(choices=MODELS, value=DEFAULT_MODEL, label="Model size")
            compute_type = gr.Dropdown(choices=COMPUTE_TYPES, value=DEFAULT_COMPUTE, label="Compute type")
            threads = gr.Slider(1, 8, step=1, value=CPU_THREADS, label="CPU threads")
            chunk_len = gr.Slider(10, 60, step=5, value=DEFAULT_CHUNK, label="Chunk length (sec)")
            beam_size = gr.Slider(1, 5, step=1, value=DEFAULT_BEAM, label="Beam size")
            vad_filter = gr.Checkbox(value=True, label="Voice activity detection")
            boost = gr.Slider(1, 12, step=1, value=10, label="Boost / Limiter")
            precise_enhance = gr.Checkbox(value=False, label="Precise denoise (CPU/RAM بیشتر)")

            run_btn = gr.Button("شروع استریم")
            build_btn = gr.Button("تولید خروجی‌های نهایی (TXT/SRT)")
        with gr.Column(scale=3):
            stream_out = gr.Textbox(label="استریم زنده متن", lines=18, elem_id="stream")
            txt_file = gr.File(label="دانلود متن (TXT)")
            srt_file = gr.File(label="دانلود زیرنویس (SRT)")
            status = gr.Markdown(f"RAM≈{RAM_GB}GB • اگر کند است: model=tiny/base، compute=int8، chunk=20–30، beam=1–2")

    run_btn.click(
        fn=transcribe_stream,
        inputs=[audio, lang, auto_tune, model_size, compute_type, threads, boost, precise_enhance, chunk_len, beam_size, vad_filter, gr.Checkbox(False, visible=False)],
        outputs=stream_out
    )
    build_btn.click(
        fn=transcribe_and_return_files,
        inputs=[audio, lang, auto_tune, model_size, compute_type, threads, boost, precise_enhance, chunk_len, beam_size, vad_filter],
        outputs=[txt_file, srt_file, status]
    )

if __name__ == "__main__":
    # سازگاری خودکار با نسخه‌های مختلف Gradio (رفع خطای concurrency_count)
    queue_sig = inspect.signature(gr.Blocks.queue)
    kwargs = {}
    if "concurrency_count" in queue_sig.parameters:
        kwargs["concurrency_count"] = 1
    if "max_size" in queue_sig.parameters:
        kwargs["max_size"] = 8

    demo.queue(**kwargs).launch(
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", 7860))
    )