import gradio as gr
import subprocess
import os
import shutil
import zipfile
import math
import glob
import tempfile
import torch

# =========================================================
# ЕРӨНХИЙ ТОХИРГОО
# =========================================================
WORK_DIR = "output_chunks"
MUSIC_DIR = "music"
FONT_PATH = "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf"

# =========================================================
# TAB 1: ВИДЕО ТАСТАГЧ
# =========================================================


def get_duration(path: str) -> float:
    cmd = [
        "ffprobe", "-v", "error",
        "-show_entries", "format=duration",
        "-of", "default=noprint_wrappers=1:nokey=1",
        path,
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    return float(result.stdout.strip())


def time_to_seconds(t: str):
    t = (t or "").strip()
    if not t:
        return None
    parts = [int(p) for p in t.split(":")]
    while len(parts) < 3:
        parts.insert(0, 0)
    h, m, s = parts
    return h * 3600 + m * 60 + s


def list_music_genres():
    if not os.path.isdir(MUSIC_DIR):
        return ["(байхгүй)"]
    genres = sorted(d for d in os.listdir(MUSIC_DIR) if os.path.isdir(os.path.join(MUSIC_DIR, d)))
    return genres if genres else ["(байхгүй)"]


def pick_music_track(genre: str):
    if not genre or genre == "(байхгүй)":
        return None
    folder = os.path.join(MUSIC_DIR, genre)
    tracks = glob.glob(os.path.join(folder, "*.mp3")) + glob.glob(os.path.join(folder, "*.wav"))
    return tracks[0] if tracks else None


def run_ffmpeg(cmd, error_prefix="FFmpeg алдаа"):
    process = subprocess.run(cmd, capture_output=True, text=True)
    if process.returncode != 0:
        raise gr.Error(f"{error_prefix}:\n{process.stderr[-800:]}")


def split_video(
    video_path,
    start_time, end_time,
    split_mode,
    minutes, seconds,
    num_parts,
    watermark_text,
    music_genre, music_volume,
    filename_prefix,
    make_thumbnails,
    progress=gr.Progress(),
):
    if video_path is None:
        raise gr.Error("Эхлээд видео файл оруулна уу.")

    if os.path.exists(WORK_DIR):
        shutil.rmtree(WORK_DIR)
    os.makedirs(WORK_DIR)

    current = video_path
    base_name = filename_prefix.strip() if filename_prefix and filename_prefix.strip() else \
        os.path.splitext(os.path.basename(video_path))[0]

    start_s = time_to_seconds(start_time)
    end_s = time_to_seconds(end_time)
    if start_s is not None or end_s is not None:
        progress(0.05, desc="Сонгосон хэсгийг тайрч байна...")
        trimmed = os.path.join(WORK_DIR, "trimmed.mp4")
        cmd = ["ffmpeg", "-y"]
        if start_s is not None:
            cmd += ["-ss", str(start_s)]
        cmd += ["-i", current]
        if end_s is not None:
            duration = end_s - (start_s or 0)
            if duration <= 0:
                raise gr.Error("Төгсгөл цаг нь эхлэл цагаас хойш байх ёстой.")
            cmd += ["-t", str(duration)]
        cmd += ["-c", "copy", "-map", "0", trimmed]
        run_ffmpeg(cmd, "Тайрахад алдаа гарлаа")
        current = trimmed

    music_track = pick_music_track(music_genre)
    needs_filters = bool((watermark_text or "").strip()) or music_track is not None

    if needs_filters:
        progress(0.2, desc="Watermark/ая нэмж байна...")
        processed = os.path.join(WORK_DIR, "processed.mp4")
        cmd = ["ffmpeg", "-y", "-i", current]

        filter_complex_parts = []
        map_video = "0:v"
        map_audio = "0:a"

        if (watermark_text or "").strip():
            text = watermark_text.strip().replace("'", "\\'").replace(":", "\\:")
            filter_complex_parts.append(
                f"[0:v]drawtext=fontfile={FONT_PATH}:text='{text}':"
                f"fontcolor=white@0.8:fontsize=28:x=w-tw-20:y=h-th-20:"
                f"box=1:boxcolor=black@0.4:boxborderw=8[vout]"
            )
            map_video = "[vout]"

        if music_track is not None:
            cmd += ["-i", music_track]
            vol = max(0.0, min(1.0, float(music_volume) / 100.0))
            filter_complex_parts.append(f"[1:a]volume={vol},aloop=loop=-1:size=2000000000[music]")
            filter_complex_parts.append(f"[0:a][music]amix=inputs=2:duration=first:dropout_transition=2[aout]")
            map_audio = "[aout]"

        cmd += ["-filter_complex", ";".join(filter_complex_parts)]
        cmd += ["-map", map_video, "-map", map_audio]
        cmd += ["-c:v", "libx264", "-preset", "veryfast", "-c:a", "aac", "-shortest", processed]
        run_ffmpeg(cmd, "Watermark/ая нэмэхэд алдаа гарлаа")
        current = processed

    progress(0.5, desc="Видеог хуваалж байна...")
    total_duration = get_duration(current)

    if split_mode == "Тэнцүү N хэсэгт":
        n = max(1, int(num_parts))
        chunk_seconds = math.ceil(total_duration / n)
    else:
        chunk_seconds = int(minutes) * 60 + int(seconds)
        if chunk_seconds <= 0:
            raise gr.Error("Хуваах хугацаа 0-ээс их байх ёстой.")

    output_pattern = os.path.join(WORK_DIR, f"{base_name}_%03d.mp4")
    cmd = [
        "ffmpeg", "-y", "-i", current,
        "-c", "copy", "-map", "0",
        "-segment_time", str(chunk_seconds),
        "-f", "segment", "-reset_timestamps", "1",
        output_pattern,
    ]
    process = subprocess.run(cmd, capture_output=True, text=True)
    if process.returncode != 0:
        cmd_reencode = [
            "ffmpeg", "-y", "-i", current,
            "-map", "0",
            "-segment_time", str(chunk_seconds),
            "-f", "segment", "-reset_timestamps", "1",
            output_pattern,
        ]
        run_ffmpeg(cmd_reencode, "Хуваахад алдаа гарлаа")

    progress(0.85, desc="Файлуудыг цуглуулж байна...")
    files = sorted(glob.glob(os.path.join(WORK_DIR, f"{base_name}_*.mp4")))
    if not files:
        raise gr.Error("Хуваалт амжилтгүй боллоо - гаралтын файл олдсонгүй.")

    thumb_files = []
    if make_thumbnails:
        progress(0.92, desc="Thumbnail үүсгэж байна...")
        for f in files:
            thumb_path = os.path.splitext(f)[0] + "_thumb.jpg"
            cmd = ["ffmpeg", "-y", "-ss", "1", "-i", f, "-vframes", "1", thumb_path]
            subprocess.run(cmd, capture_output=True, text=True)
            if os.path.exists(thumb_path):
                thumb_files.append(thumb_path)

    zip_path = os.path.join(WORK_DIR, f"{base_name}_parts.zip")
    with zipfile.ZipFile(zip_path, "w") as zf:
        for f in files + thumb_files:
            zf.write(f, os.path.basename(f))

    progress(1.0, desc="Дууслаа!")
    status = f"✅ {len(files)} хэсэгт хуваалаа (тус бүр ~{chunk_seconds // 60} мин {chunk_seconds % 60} сек)"
    if music_track:
        status += f" | Ая: {os.path.basename(music_track)}"
    return status, files + thumb_files, zip_path


# =========================================================
# TAB 2: ДУУ ХООЛОЙ КЛОН (OmniVoice, CPU, lazy-load)
# =========================================================

_voice_model = None


def get_voice_model():
    """Загварыг зөвхөн энэ tab-ыг эхлээд ашиглах үед л ачаална - splitter tab-ын хурдад нөлөөлөхгүй."""
    global _voice_model
    if _voice_model is None:
        from omnivoice import OmniVoice
        _voice_model = OmniVoice.from_pretrained(
            "k2-fsa/OmniVoice",
            device_map="cpu",
            dtype=torch.float32,
        )
        _voice_model.eval()
    return _voice_model


def clone_voice(text, ref_audio, ref_text, language, progress=gr.Progress()):
    if not text or not text.strip():
        raise gr.Error("Хэлэх текстээ оруулна уу.")
    if ref_audio is None:
        raise gr.Error("Дуу хоолойн жишээ (3-30 сек, ӨӨРИЙН эсвэл ашиглах эрхтэй хоолой) оруулна уу.")

    progress(0.05, desc="Загвар ачааллаж байна (эхний удаад л удаан)...")
    model = get_voice_model()

    progress(0.3, desc="Дуу хоолой үүсгэж байна (CPU тул хүлээнэ үү)...")
    kwargs = {"text": text.strip(), "ref_audio": ref_audio}
    if ref_text and ref_text.strip():
        kwargs["ref_text"] = ref_text.strip()
    if language and language.strip():
        kwargs["language"] = language.strip()

    with torch.no_grad():
        audio = model.generate(**kwargs)

    progress(0.9, desc="Файл бэлдэж байна...")
    out_path = tempfile.mktemp(suffix=".wav")
    result = audio[0]
    if isinstance(result, torch.Tensor):
        import torchaudio
        torchaudio.save(out_path, result if result.dim() == 2 else result.unsqueeze(0), 24000)
    else:
        import soundfile as sf
        sf.write(out_path, result, 24000)

    progress(1.0, desc="Дууслаа!")
    return out_path


# =========================================================
# TAB 3: БАГЦ ДУБЛАЖ (Segment Dubbing) - тогтвортой НЭГ хоолойтой,
# timestamp-аар синк хийж, өнгө/хурд/дуу чанга 3 горимтой
# =========================================================

import re
import numpy as np

SEGMENT_LINE_RE = re.compile(
    r"^\s*(\d{1,2}:\d{2}(?::\d{2})?)\s*-\s*(\d{1,2}:\d{2}(?::\d{2})?)\s*:\s*(.+?)\s*$"
)


def _parse_ts(ts: str) -> float:
    parts = [int(p) for p in ts.split(":")]
    while len(parts) < 3:
        parts.insert(0, 0)
    h, m, s = parts
    return h * 3600 + m * 60 + s


def parse_segments(raw_text: str):
    """
    Мөр бүрийг задална. Дэмждэг формат:
      00:00-00:03: Сайн байна уу
      00:03-00:07: Энэ бол жишээ өгүүлбэр
    Timestamp байхгүй бол зүгээр текст мөр гэж үзнэ (duration sync хийхгүй).
    Буцаах: list[dict(start, end, duration, text)]
    """
    segments = []
    for raw_line in raw_text.splitlines():
        line = raw_line.strip()
        if not line:
            continue
        m = SEGMENT_LINE_RE.match(line)
        if m:
            start_s = _parse_ts(m.group(1))
            end_s = _parse_ts(m.group(2))
            text = m.group(3).strip()
            if end_s <= start_s:
                raise gr.Error(f"Буруу мөр (эцэс нь эхлэлээс өмнө байна): {line}")
            segments.append({
                "start": start_s, "end": end_s,
                "duration": end_s - start_s, "text": text,
            })
        else:
            segments.append({"start": None, "end": None, "duration": None, "text": line})
    if not segments:
        raise gr.Error("Дор хаяж нэг мөр текст оруулна уу.")
    return segments


def postprocess_audio(y: np.ndarray, sr: int, speed: float, pitch_semitones: float, gain_db: float) -> np.ndarray:
    """3 горим: хурд (speed), өнгө (pitch), дуу чанга (gain) - үүсгэсэн ЯГ ТЭР дуу хоолойг өөрчлөхгүйгээр тохируулна."""
    import librosa

    y = np.asarray(y, dtype=np.float32)
    if y.ndim > 1:
        y = y.mean(axis=0)

    if speed and abs(speed - 1.0) > 1e-3:
        y = librosa.effects.time_stretch(y, rate=float(speed))

    if pitch_semitones and abs(pitch_semitones) > 1e-3:
        y = librosa.effects.pitch_shift(y, sr=sr, n_steps=float(pitch_semitones))

    if gain_db and abs(gain_db) > 1e-3:
        y = y * (10.0 ** (gain_db / 20.0))

    peak = np.max(np.abs(y)) if y.size else 0.0
    if peak > 0.99:
        y = y / peak * 0.99

    return y.astype(np.float32)


def _generate_one(model, text, ref_audio, ref_text, language, duration):
    kwargs = {"text": text.strip(), "ref_audio": ref_audio}
    # ref_audio/ref_text-г БҮХ сегментэд яг ижилхэн дамжуулж байгаа тул
    # хоолой хооронд солигдохгүй, тогтвортой байна.
    if ref_text and ref_text.strip():
        kwargs["ref_text"] = ref_text.strip()
    if language and language.strip():
        kwargs["language"] = language.strip()
    if duration:
        kwargs["duration"] = float(duration)  # timestamp-тай тааруулж синк хийнэ

    with torch.no_grad():
        audio = model.generate(**kwargs)

    result = audio[0]
    if isinstance(result, torch.Tensor):
        arr = result.squeeze().cpu().numpy()
    else:
        arr = np.asarray(result)
    return arr


def batch_dub(
    segments_text, ref_audio, ref_text, language,
    speed, pitch_semitones, gain_db,
    fill_gaps, progress=gr.Progress(),
):
    if not segments_text or not segments_text.strip():
        raise gr.Error("Сегментүүдээ оруулна уу.")
    if ref_audio is None:
        raise gr.Error("Дуу хоолойн жишээ (нэг л удаа) оруулна уу - энэ бүх сегментэд адилхан ашиглагдана.")

    segments = parse_segments(segments_text)

    progress(0.02, desc="Загвар ачааллаж байна...")
    model = get_voice_model()
    sr = 24000

    piece_dir = tempfile.mkdtemp(prefix="dub_segments_")
    timeline_pieces = []
    per_segment_files = []
    prev_end = None

    total = len(segments)
    for i, seg in enumerate(segments):
        progress(0.05 + 0.85 * (i / total), desc=f"Сегмент {i + 1}/{total} үүсгэж байна...")

        # timestamp-той бол сегментүүдийн хоорондох завсрыг чимээгүй хэсгээр дүүргэнэ
        if fill_gaps and prev_end is not None and seg["start"] is not None:
            gap = seg["start"] - prev_end
            if gap > 0.02:
                timeline_pieces.append(np.zeros(int(gap * sr), dtype=np.float32))

        raw = _generate_one(model, seg["text"], ref_audio, ref_text, language, seg["duration"])
        processed = postprocess_audio(raw, sr, speed, pitch_semitones, gain_db)
        timeline_pieces.append(processed)

        seg_path = os.path.join(piece_dir, f"segment_{i + 1:03d}.wav")
        import soundfile as sf
        sf.write(seg_path, processed, sr)
        per_segment_files.append(seg_path)

        if seg["end"] is not None:
            prev_end = seg["end"]

    progress(0.92, desc="Бүх сегментийг нэгтгэж байна...")
    final_audio = np.concatenate(timeline_pieces) if timeline_pieces else np.zeros(0, dtype=np.float32)
    final_path = os.path.join(piece_dir, "dubbed_full.wav")
    import soundfile as sf
    sf.write(final_path, final_audio, sr)

    zip_path = os.path.join(piece_dir, "segments.zip")
    with zipfile.ZipFile(zip_path, "w") as zf:
        for f in per_segment_files:
            zf.write(f, os.path.basename(f))

    progress(1.0, desc="Дууслаа!")
    status = f"✅ {total} сегмент нэг тогтвортой хоолойгоор үүслээ (нийт {len(final_audio) / sr:.1f} сек)."
    return status, final_path, zip_path


# =========================================================
# НЭГТГЭСЭН ИНТЕРФЭЙС
# =========================================================

with gr.Blocks(title="КиноЭзэн Студи") as demo:
    gr.Markdown("# 🎬 КиноЭзэн Студи")
    gr.Markdown("Нэг сайт дотор: видео тастагч + дуу хоолой клон (dubbing).")

    with gr.Tabs():
        with gr.Tab("✂️ Видео Тастагч"):
            video_input = gr.Video(label="Видео оруулах")

            with gr.Accordion("Хугацааны диапазон (заавал биш)", open=False):
                gr.Markdown("Формат: MM:SS эсвэл HH:MM:SS")
                with gr.Row():
                    start_input = gr.Textbox(label="Эхлэх цаг", placeholder="05:00")
                    end_input = gr.Textbox(label="Дуусах цаг", placeholder="25:00")

            split_mode = gr.Radio(
                ["Тогтмол хугацаагаар", "Тэнцүү N хэсэгт"],
                value="Тогтмол хугацаагаар", label="Хуваах горим",
            )
            with gr.Row():
                minutes_input = gr.Number(label="Минут (хэсэг тус бүр)", value=10, precision=0, minimum=0)
                seconds_input = gr.Number(label="Секунд", value=0, precision=0, minimum=0, maximum=59)
                num_parts_input = gr.Number(label="Хэдэн хэсэгт (тэнцүү горимд)", value=3, precision=0, minimum=1)

            with gr.Accordion("Watermark (заавал биш)", open=False):
                watermark_input = gr.Textbox(label="Watermark текст", placeholder="КиноЭзэн")

            with gr.Accordion("Дэвсгэр ая (заавал биш)", open=False):
                gr.Markdown("`music/<жанр>/` хавтаст royalty-free mp3/wav байршуулсан бол энд сонгоно.")
                music_genre_input = gr.Dropdown(choices=list_music_genres(), value="(байхгүй)", label="Жанр")
                refresh_btn = gr.Button("Жанр шинэчлэх", size="sm")
                music_volume_input = gr.Slider(0, 100, value=15, label="Аяны түвшин (%)")

            with gr.Row():
                filename_prefix_input = gr.Textbox(label="Файлын нэрийн эхлэл (заавал биш)")
                thumbnail_checkbox = gr.Checkbox(label="Thumbnail үүсгэх", value=False)

            split_btn = gr.Button("Хуваах", variant="primary")

            status_output = gr.Textbox(label="Төлөв", interactive=False)
            files_output = gr.File(label="Тусдаа файлууд", file_count="multiple")
            zip_output = gr.File(label="ZIP татаж авах (бүгд)")

            refresh_btn.click(fn=lambda: gr.update(choices=list_music_genres()), outputs=music_genre_input)

            split_btn.click(
                fn=split_video,
                inputs=[
                    video_input, start_input, end_input, split_mode,
                    minutes_input, seconds_input, num_parts_input,
                    watermark_input, music_genre_input, music_volume_input,
                    filename_prefix_input, thumbnail_checkbox,
                ],
                outputs=[status_output, files_output, zip_output],
            )

        with gr.Tab("🎙️ Дуу Хоолой Клон"):
            gr.Markdown(
                "⚠️ **Зөвхөн өөрийн эсвэл ашиглах эрхтэй дуу хоолойгоо клонд ашиглана уу.** "
                "CPU дээр ажилладаг тул quota-гүй, гэхч GPU-той харьцуулахад удаан (~30-90 сек/өгүүлбэр)."
            )
            with gr.Row():
                with gr.Column():
                    text_input = gr.Textbox(label="Хэлэх текст (Монгол)", lines=4)
                    ref_audio_input = gr.Audio(label="Дуу хоолойн жишээ (3-30 сек)", type="filepath")
                    ref_text_input = gr.Textbox(label="Жишээ аудионы бичвэр (заавал биш)", lines=2)
                    language_input = gr.Textbox(label="Хэл", value="Mongolian")
                    generate_btn = gr.Button("Дуу хоолой үүсгэх", variant="primary")
                with gr.Column():
                    audio_output = gr.Audio(label="Гарсан дуу хоолой")

            generate_btn.click(
                fn=clone_voice,
                inputs=[text_input, ref_audio_input, ref_text_input, language_input],
                outputs=audio_output,
            )

        with gr.Tab("🎬 Багц Дублаж"):
            gr.Markdown(
                "**Видео бүхэлд нь НЭГ л тогтвортой хоолойгоор дублаж хийнэ.** "
                "Reference audio-г зөвхөн нэг удаа сонгоно - бүх сегментэд яг адилхан ашиглагдана, "
                "тиймээс хоолой сегмент хооронд солигдохгүй.\n\n"
                "**Сегментийн формат** (мөр тус бүр нэг репликтой):\n"
                "```\n00:00-00:03: Сайн байна уу\n00:03-00:07: Энэ бол жишээ өгүүлбэр\n```\n"
                "Timestamp өгвөл (Claude API-аас гарсан орчуулгын хугацаатай тааруулж) загвар "
                "яг тэр л урттай хоолой гаргаж, синк алдагдахгүй. Timestamp-гүй бол зүгээр мөр мөрөөр "
                "текст оруулж болно (тухайн үед speed slider-ээр хурдыг тохируулна)."
            )
            with gr.Row():
                with gr.Column():
                    segments_input = gr.Textbox(
                        label="Сегментүүд",
                        lines=10,
                        placeholder="00:00-00:03: Сайн байна уу\n00:03-00:07: Энэ бол жишээ өгүүлбэр",
                    )
                    batch_ref_audio_input = gr.Audio(
                        label="Дуу хоолойн жишээ (НЭГ Л УДАА сонгоно - бүх сегментэд адилхан хэрэглэгдэнэ)",
                        type="filepath",
                    )
                    batch_ref_text_input = gr.Textbox(label="Жишээ аудионы бичвэр (заавал биш)", lines=2)
                    batch_language_input = gr.Textbox(label="Хэл", value="Mongolian")

                    gr.Markdown("**3 горим (post-processing):**")
                    speed_slider = gr.Slider(0.7, 1.4, value=1.0, step=0.01, label="⏩ Хурд (1.0 = хэвийн)")
                    pitch_slider = gr.Slider(-6, 6, value=0.0, step=0.5, label="🎚️ Өнгө: нарийн ⟷ бүдүүн (semitone)")
                    gain_slider = gr.Slider(-10, 10, value=0.0, step=0.5, label="🔊 Дуу чанга/султгах (dB)")
                    fill_gaps_checkbox = gr.Checkbox(
                        label="Сегментүүдийн хоорондох завсрыг чимээгүйгээр дүүргэх (timestamp-той үед)",
                        value=True,
                    )

                    batch_generate_btn = gr.Button("Багц дублаж эхлүүлэх", variant="primary")

                with gr.Column():
                    batch_status_output = gr.Textbox(label="Төлөв", interactive=False)
                    batch_audio_output = gr.Audio(label="Бүтэн дубласан аудио")
                    batch_zip_output = gr.File(label="Сегмент бүрийн ZIP (тус тусад нь)")

            batch_generate_btn.click(
                fn=batch_dub,
                inputs=[
                    segments_input, batch_ref_audio_input, batch_ref_text_input, batch_language_input,
                    speed_slider, pitch_slider, gain_slider, fill_gaps_checkbox,
                ],
                outputs=[batch_status_output, batch_audio_output, batch_zip_output],
            )

if __name__ == "__main__":
    demo.launch()