speaker_diarization.py

"""Speaker Diarization Tool - Ai Đang Nói?
Phân biệt giọng nói trong video dựa trên file SRT.
Tích hợp: Lồng tiếng OmniVoice + Xuất CapCut Project.
"""
import sys, os, re, subprocess, tempfile, time, json, uuid
import soundfile as sf

# OmniVoice TTS path
TTS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "chumtts2")
VOICES_FILE = os.path.join(TTS_DIR, "voice", "voices.json")
if TTS_DIR not in sys.path:
    sys.path.insert(0, TTS_DIR)
import numpy as np
import torch
import torchaudio
from PyQt5.QtWidgets import (QApplication, QWidget, QVBoxLayout, QHBoxLayout,
    QLabel, QLineEdit, QPushButton, QTextEdit, QFileDialog, QTableWidget,
    QTableWidgetItem, QProgressBar, QHeaderView, QComboBox, QGroupBox,
    QAbstractItemView, QMessageBox, QSpinBox, QDoubleSpinBox, QDialog, QScrollArea, QFrame)
from PyQt5.QtCore import QThread, pyqtSignal, Qt
from PyQt5.QtGui import QColor, QPainter, QFont, QBrush, QPen

# Fix Windows symlink error: monkey-patch os.symlink to copy instead
import shutil
_original_symlink = os.symlink
def _safe_symlink(src, dst, target_is_directory=False):
    try:
        _original_symlink(src, dst, target_is_directory)
    except OSError:
        if os.path.isdir(str(src)):
            shutil.copytree(str(src), str(dst), dirs_exist_ok=True)
        else:
            shutil.copy2(str(src), str(dst))
os.symlink = _safe_symlink


# Note: speechbrain/sklearn no longer needed — using pyannote.audio for diarization

COLORS = [QColor('#e74c3c'), QColor('#3498db'), QColor('#2ecc71'), QColor('#f39c12'),
          QColor('#9b59b6'), QColor('#1abc9c'), QColor('#e67e22'), QColor('#34495e'),
          QColor('#c0392b'), QColor('#27ae60')]
NAMES = [f'Speaker {chr(65+i)}' for i in range(10)]
MIN_DUR = 0.3

# ── SRT Parser ──────────────────────────────────────────────
def parse_srt(path):
    with open(path, 'r', encoding='utf-8-sig') as f:
        content = f.read()
    pat = re.compile(
        r'(\d+)\s*\n(\d{2}):(\d{2}):(\d{2})[,.](\d{3})\s*-->\s*'
        r'(\d{2}):(\d{2}):(\d{2})[,.](\d{3})\s*\n((?:(?!\n\n|\d+\s*\n\d{2}:\d{2}).+\n?)+)', re.MULTILINE)
    entries = []
    for m in pat.finditer(content):
        s_ms = (int(m.group(2))*3600+int(m.group(3))*60+int(m.group(4)))*1000+int(m.group(5))
        e_ms = (int(m.group(6))*3600+int(m.group(7))*60+int(m.group(8)))*1000+int(m.group(9))
        entries.append({'index': int(m.group(1)), 'start_ms': s_ms, 'end_ms': e_ms,
                        'text': m.group(10).strip(), 'speaker': None})
    return entries

def ms_fmt(ms):
    return f"{ms//60000:02d}:{(ms%60000)//1000:02d}"

def ms_srt(ms):
    h, r = divmod(ms, 3600000)
    m, r = divmod(r, 60000)
    s, ms_r = divmod(r, 1000)
    return f"{h:02d}:{m:02d}:{s:02d},{ms_r:03d}"

def find_ffmpeg():
    for d in [os.path.dirname(os.path.abspath(__file__)),
              os.path.join(os.path.dirname(os.path.abspath(__file__)), '..'), '']:
        p = os.path.join(d, 'ffmpeg.exe') if d else 'ffmpeg'
        if d == '' or os.path.exists(p):
            return p
    return 'ffmpeg'

# ── Analysis Worker ─────────────────────────────────────────
class AnalysisWorker(QThread):
    log = pyqtSignal(str)
    prog = pyqtSignal(int)
    done = pyqtSignal(list, int)
    err = pyqtSignal(str)

    def __init__(self, video, entries, expected_k=0):
        super().__init__()
        self.video = video
        self.entries = entries
        self.expected_k = expected_k

    def run(self):
        try:
            tmp = tempfile.mkdtemp()
            wav = os.path.join(tmp, 'audio.wav')

            self.log.emit("🔊 [1/4] Trích xuất audio từ video...")
            subprocess.run([find_ffmpeg(), '-y', '-i', self.video, '-ac', '1', '-ar', '16000', '-vn', wav],
                           capture_output=True)
            if not os.path.exists(wav):
                self.err.emit("Không trích xuất được audio!"); return
            self.prog.emit(8)

            # --- BS-ROFORMER VOCAL SEPARATION ---
            self.log.emit("🎤 [2/4] Tách giọng nói bằng BS-RoFormer (chất lượng cao)...")
            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
            try:
                import yaml
                from bs_roformer import get_model_from_config, demix_track
                from ml_collections import ConfigDict

                # Load BS-RoFormer model
                roformer_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                                            "models", "roformer-model-bs-roformer-sw-by-jarredou")
                config_path = os.path.join(roformer_dir, "BS-Rofo-SW-Fixed.yaml")
                ckpt_path = os.path.join(roformer_dir, "BS-Rofo-SW-Fixed.ckpt")

                with open(config_path) as f:
                    config = ConfigDict(yaml.load(f, Loader=yaml.SafeLoader))

                roformer_model = get_model_from_config("bs_roformer", config)
                roformer_model.load_state_dict(torch.load(ckpt_path, map_location="cpu"))
                roformer_model = roformer_model.to(device)
                roformer_model.eval()
                self.prog.emit(15)

                # Load audio and prepare for BS-RoFormer (expects stereo float32)
                mix, sr_orig = sf.read(wav)
                if len(mix.shape) == 1:
                    mix = np.stack([mix, mix], axis=-1)
                mixture = torch.tensor(mix.T, dtype=torch.float32)

                self.log.emit(f"  ⏳ BS-RoFormer: {mix.shape[0] / sr_orig:.0f}s audio...")
                res, _ = demix_track(config, roformer_model, mixture, device)

                # Get vocals stem and compute instrumental = original - vocals
                vocals_np = res["vocals"].T  # [T, 2]
                instrumental_np = mix - vocals_np  # Preserve everything except voice

                # Save high-quality stereo stems to disk for CapCut export
                vid_base = os.path.splitext(self.video)[0]
                self.vocal_path = f"{vid_base}_vocal.wav"
                self.inst_path = f"{vid_base}_instrumental.wav"
                sf.write(self.vocal_path, vocals_np, sr_orig, subtype="FLOAT")
                sf.write(self.inst_path, instrumental_np, sr_orig, subtype="FLOAT")

                # Convert vocals to mono 16kHz for pyannote diarization
                vocals_tensor = torch.tensor(vocals_np.T, dtype=torch.float32)  # [2, T]
                vocals_mono = vocals_tensor.mean(dim=0, keepdim=True)  # [1, T]
                vocals_16k = torchaudio.functional.resample(vocals_mono, sr_orig, 16000)

                # Free GPU memory
                del roformer_model, mixture, res
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()

                self.log.emit("  ✅ Tách vocal thành công! (BS-RoFormer)")
                self.prog.emit(25)
            except Exception as roformer_err:
                self.log.emit(f"  ⚠️ BS-RoFormer lỗi ({roformer_err}), dùng audio gốc...")
                vocals_16k, _ = torchaudio.load(wav)
                self.prog.emit(25)

            # --- PYANNOTE DIARIZATION ---
            self.log.emit("🧠 [3/4] Tải pyannote speaker-diarization-3.1...")
            from pyannote.audio import Pipeline as PyannotePipeline
            import huggingface_hub
            hf_token = huggingface_hub.get_token()
            pipeline = PyannotePipeline.from_pretrained(
                "pyannote/speaker-diarization-3.1",
                token=hf_token
            )
            if torch.cuda.is_available():
                pipeline.to(torch.device("cuda"))
                self.log.emit("  ✅ Đang dùng GPU (CUDA)")
            else:
                self.log.emit("  ⚠️ Chạy trên CPU (sẽ chậm hơn)")
            self.prog.emit(35)

            self.log.emit("🔍 [4/4] Phân tích giọng nói (pyannote trên vocal sạch)...")
            audio_input = {"waveform": vocals_16k, "sample_rate": 16000}
            # Run diarization
            if self.expected_k > 0:
                self.log.emit(f"  🎯 Ép buộc {self.expected_k} người nói")
                diarization = pipeline(audio_input, num_speakers=self.expected_k)
            else:
                self.log.emit("  🎯 Tự động phát hiện số người nói (min=2, max=10)")
                diarization = pipeline(audio_input, min_speakers=2, max_speakers=10)
            self.prog.emit(80)

            # Collect unique speakers from pyannote output (v4.x API)
            speaker_labels = set()
            pyannote_segments = []
            annotation = diarization.speaker_diarization  # pyannote 4.x
            for turn, _, speaker in annotation.itertracks(yield_label=True):
                speaker_labels.add(speaker)
                pyannote_segments.append({
                    'start_ms': int(turn.start * 1000),
                    'end_ms': int(turn.end * 1000),
                    'speaker': speaker
                })

            # Map speaker label strings to integer indices
            label_to_idx = {label: idx for idx, label in enumerate(sorted(speaker_labels))}
            n_speakers = len(label_to_idx)
            self.log.emit(f"  📊 Pyannote phát hiện {n_speakers} người nói: {list(label_to_idx.keys())}")

            # Assign each SRT entry to the speaker with the most overlap
            for entry in self.entries:
                e_start = entry['start_ms']
                e_end = entry['end_ms']
                # Calculate overlap with each speaker
                speaker_overlap = {}
                for seg in pyannote_segments:
                    overlap_start = max(e_start, seg['start_ms'])
                    overlap_end = min(e_end, seg['end_ms'])
                    overlap = max(0, overlap_end - overlap_start)
                    if overlap > 0:
                        spk_idx = label_to_idx[seg['speaker']]
                        speaker_overlap[spk_idx] = speaker_overlap.get(spk_idx, 0) + overlap

                if speaker_overlap:
                    entry['speaker'] = max(speaker_overlap, key=speaker_overlap.get)
                else:
                    entry['speaker'] = 0  # Default if no overlap found

            self.log.emit(f"✅ Phân tích xong! Phát hiện {n_speakers} người nói.")
            self.prog.emit(100)
            self.done.emit(self.entries, n_speakers)
            self._wav_path = wav
        except Exception as ex:
            import traceback
            self.err.emit(str(ex)+"\n"+traceback.format_exc())

class NoScrollComboBox(QComboBox):
    def wheelEvent(self, e):
        e.ignore()

# ── Timeline Widget ─────────────────────────────────────────
class TimelineWidget(QWidget):
    clicked = pyqtSignal(int)

    def __init__(self):
        super().__init__()
        self.entries = []
        self.total_ms = 1
        self.hl = -1
        self.setMinimumHeight(55)
        self.setMaximumHeight(65)

    def set_data(self, entries, total_ms):
        self.entries = entries
        self.total_ms = max(total_ms, 1)
        self.update()

    def highlight(self, i):
        self.hl = i; self.update()

    def paintEvent(self, e):
        if not self.entries: return
        p = QPainter(self); p.setRenderHint(QPainter.Antialiasing)
        w, h, xo, yo = self.width()-20, self.height()-25, 10, 5
        p.fillRect(xo, yo, w, h, QColor('#f0f0f0')) # Light gray bg for light mode
        for i, en in enumerate(self.entries):
            x1 = xo+int(en['start_ms']/self.total_ms*w)
            x2 = xo+int(en['end_ms']/self.total_ms*w)
            sw = max(x2-x1, 2)
            spk = en.get('speaker', 0) or 0
            if i == self.hl:
                p.setPen(QPen(QColor('#000'), 2)) # Black border for highlight
            else:
                p.setPen(Qt.NoPen)
            p.setBrush(QBrush(COLORS[spk % len(COLORS)]))
            p.drawRect(x1, yo, sw, h)
        p.setPen(QColor('#aaa')); p.setFont(QFont('Consolas', 8))
        for pct in [0, .25, .5, .75, 1.0]:
            p.drawText(xo+int(pct*w)-15, self.height()-3, ms_fmt(int(pct*self.total_ms)))
        p.end()

    def mousePressEvent(self, ev):
        if not self.entries: return
        click_ms = int((ev.x()-10)/(self.width()-20)*self.total_ms)
        closest = min(range(len(self.entries)),
                      key=lambda i: abs((self.entries[i]['start_ms']+self.entries[i]['end_ms'])/2-click_ms))
        self.clicked.emit(closest)

# ── Voice Config ─────────────────────────────────────────────
def load_voices_config():
    if os.path.exists(VOICES_FILE):
        with open(VOICES_FILE, 'r', encoding='utf-8') as f:
            return json.load(f)
    return {}

# ── Dubbing Worker (OmniVoice TTS) ──────────────────────────
class DubbingWorker(QThread):
    log = pyqtSignal(str)
    prog = pyqtSignal(int)
    done = pyqtSignal(str)  # output folder
    err = pyqtSignal(str)

    def __init__(self, entries, voice_map, output_dir, tts_config):
        """
        entries: list of dicts with index, start_ms, end_ms, text, speaker (int)
        voice_map: dict {speaker_int: voice_name_str}
        output_dir: folder to save WAVs + manifest.json
        tts_config: dict with 'speed', 'steps', 'guidance'
        """
        super().__init__()
        self.entries = entries
        self.voice_map = voice_map
        self.output_dir = output_dir
        self.tts_config = tts_config

    def _clean_wav(self, wav, sr=24000):
        """Remove leading noise artifacts and trim silence from OmniVoice output.
        
        OmniVoice sometimes produces:
        1) A short noise burst/click at the very start ("xì", pop, etc.)
        2) Excess silence padding at the beginning and end
        
        Strategy:
        - Compute short-window RMS energy
        - Find where sustained speech begins (not just a single spike)
        - Trim silence from both ends
        """
        if len(wav) < sr * 0.1:  # Too short, skip
            return wav
        
        # Window size for energy analysis (10ms windows)
        win = int(sr * 0.01)
        hop = win // 2
        
        # Compute RMS energy per window
        n_frames = (len(wav) - win) // hop + 1
        if n_frames < 3:
            return wav
        rms = np.array([
            np.sqrt(np.mean(wav[i*hop : i*hop+win] ** 2))
            for i in range(n_frames)
        ])
        
        # Silence threshold: use the quietest 10% as noise floor, then 6x above that
        sorted_rms = np.sort(rms)
        noise_floor = np.mean(sorted_rms[:max(1, len(sorted_rms)//10)])
        threshold = max(noise_floor * 6, 0.005)  # Minimum absolute threshold
        
        # Find speech onset: first window where at least 3 consecutive windows exceed threshold
        # This skips isolated noise spikes at the start
        speech_start_frame = 0
        consecutive = 0
        min_consecutive = 3  # Need 3 windows (30ms) of sustained energy = real speech
        for i in range(n_frames):
            if rms[i] > threshold:
                consecutive += 1
                if consecutive >= min_consecutive:
                    speech_start_frame = i - min_consecutive + 1
                    break
            else:
                consecutive = 0
        
        # Find speech end: last window exceeding threshold
        speech_end_frame = n_frames - 1
        for i in range(n_frames - 1, -1, -1):
            if rms[i] > threshold:
                speech_end_frame = i
                break
        
        # Convert frame indices to sample indices
        # Start margin is small (5ms) to precisely cut noise spikes
        # End margin is larger (80ms) to preserve fading speech tails ("ạ", "à", "ơi"...)
        start_margin = int(sr * 0.005)
        end_margin = int(sr * 0.08)
        start_sample = max(0, speech_start_frame * hop - start_margin)
        end_sample = min(len(wav), (speech_end_frame + 1) * hop + end_margin)
        
        trimmed = wav[start_sample:end_sample]
        return trimmed if len(trimmed) > 0 else wav

    def run(self):
        try:
            os.makedirs(self.output_dir, exist_ok=True)
            self.log.emit("Loading OmniVoice model...")
            self.prog.emit(5)

            # Workaround: speechbrain lazy modules conflict with torch.distributed
            import sys as _sys
            _lazy = {k: v for k, v in _sys.modules.items()
                     if 'speechbrain' in k and hasattr(v, 'ensure_module')}
            for k in _lazy:
                del _sys.modules[k]

            from omnivoice import OmniVoice, OmniVoiceGenerationConfig
            model = OmniVoice.from_pretrained(TTS_DIR, device_map="cuda:0", dtype=torch.float16)
            self.prog.emit(15)

            voices_cfg = load_voices_config()
            prompts_cache = {}
            unique_voices = set(self.voice_map.values())
            for vname in unique_voices:
                vdata = voices_cfg.get(vname)
                if not vdata:
                    continue
                apath = os.path.join(TTS_DIR, "voice", vdata["audio"])
                if os.path.exists(apath):
                    self.log.emit(f"  Caching voice: {vname}")
                    prompts_cache[vname] = model.create_voice_clone_prompt(
                        ref_audio=apath, ref_text=vdata["text"])
            self.prog.emit(25)

            config = OmniVoiceGenerationConfig(
                num_step=self.tts_config.get('steps', 32),
                guidance_scale=self.tts_config.get('guidance', 5.0)
            )
            speed = self.tts_config.get('speed', 1.0)
            manifest = []
            total = len(self.entries)

            for i, e in enumerate(self.entries):
                spk_int = e.get('speaker', 0) or 0
                vname = self.voice_map.get(spk_int, "")
                prompt = prompts_cache.get(vname)
                if not prompt:
                    self.log.emit(f"  Skip #{e['index']}: no voice for speaker {spk_int}")
                    continue

                self.log.emit(f"  [{i+1}/{total}] #{e['index']}: {e['text'][:40]}...")
                sentences = [s.strip() for s in re.split(r'(?<=[.!?\n])\s+', e['text']) if s.strip()]
                if not sentences:
                    sentences = [e['text']]

                FILLER = "a lô một hai ba bốn năm"
                MIN_CHARS = 8
                chunks = []
                for sent in sentences:
                    if len(sent) >= MIN_CHARS:
                        # Normal inference for sentences with enough context
                        audio = model.generate(text=sent, language="vietnamese",
                                               voice_clone_prompt=prompt, generation_config=config, speed=speed)
                        chunks.append(audio[0])
                    else:
                        # Short sentence: pad with filler for context, then trim
                        self.log.emit(f"    ⚡ Câu ngắn '{sent}' → pad context")
                        padded_text = f"{sent}.. {FILLER}"
                        # 1) Infer the combined text (short sentence + filler)
                        audio_combined = model.generate(text=padded_text, language="vietnamese",
                                                        voice_clone_prompt=prompt, generation_config=config, speed=speed)
                        # 2) Infer the filler alone to measure its duration
                        audio_filler = model.generate(text=FILLER, language="vietnamese",
                                                      voice_clone_prompt=prompt, generation_config=config, speed=speed)
                        filler_len = len(audio_filler[0])
                        combined_len = len(audio_combined[0])
                        # 3) Trim: keep only the target portion (combined - filler)
                        # Add a small overlap margin (50ms) to avoid cutting into the target
                        margin = int(24000 * 0.05)
                        cut_point = max(0, combined_len - filler_len - margin)
                        if cut_point > 0:
                            chunks.append(audio_combined[0][:cut_point])
                        else:
                            chunks.append(audio_combined[0])
                    chunks.append(np.zeros(int(24000 * 0.05), dtype=np.float32))

                wav_data = np.concatenate(chunks)

                # ── POST-PROCESSING: Remove noise artifacts & trim silence ──
                wav_data = self._clean_wav(wav_data, sr=24000)

                fname = f"{e['index']:03d}_{ms_fmt(e['start_ms']).replace(':','m')}s.wav"
                fpath = os.path.join(self.output_dir, fname)
                sf.write(fpath, wav_data, 24000)

                manifest.append({
                    "index": e['index'], "start_ms": e['start_ms'], "end_ms": e['end_ms'],
                    "speaker": vname, "text": e['text'], "wav": fname
                })
                self.prog.emit(25 + int((i+1)/total*70))

            mpath = os.path.join(self.output_dir, "manifest.json")
            with open(mpath, 'w', encoding='utf-8') as f:
                json.dump(manifest, f, ensure_ascii=False, indent=2)

            self.log.emit(f"Done! {len(manifest)} WAVs saved to {self.output_dir}")
            self.prog.emit(100)
            self.done.emit(self.output_dir)

            del model
            __import__('torch').cuda.empty_cache()
        except Exception as ex:
            import traceback
            self.err.emit(str(ex) + "\n" + traceback.format_exc())

# ── CapCut Export Worker ─────────────────────────────────────
class CapCutWorker(QThread):
    log = pyqtSignal(str)
    prog = pyqtSignal(int)
    done = pyqtSignal(str)
    err = pyqtSignal(str)

    def __init__(self, video_path, dubbed_dir, project_name, video_speed=1.0, vocal_path=None, inst_path=None, stretch_threshold=1.0, min_gap_s=0.3):
        super().__init__()
        self.video_path = video_path.replace("\\", "/")
        self.dubbed_dir = dubbed_dir
        self.project_name = project_name
        self.video_speed = video_speed
        self.vocal_path = vocal_path.replace("\\", "/") if vocal_path else None
        self.inst_path = inst_path.replace("\\", "/") if inst_path else None
        self.stretch_threshold = stretch_threshold
        self.min_gap_us = int(min_gap_s * 1_000_000)
        self.local_appdata = os.getenv('LOCALAPPDATA')
        self.draft_dir = os.path.join(self.local_appdata, 'CapCut', 'User Data', 'Projects', 'com.lveditor.draft')

    def uid(self):
        return str(uuid.uuid4()).upper()

    def get_video_duration_us(self):
        import cv2
        cap = cv2.VideoCapture(self.video_path)
        if not cap.isOpened(): return 10_000_000
        fps = cap.get(cv2.CAP_PROP_FPS)
        fc = cap.get(cv2.CAP_PROP_FRAME_COUNT)
        cap.release()
        if fps <= 0: return 10_000_000
        return int((fc / fps) * 1_000_000)

    def make_speed(self, sid, speed=1.0):
        return {"curve_speed": None, "id": sid, "mode": 0, "speed": speed, "type": "speed"}

    def make_canvas(self, cid):
        return {"album_image": "", "blur": 0.0, "color": "", "id": cid, "image": "", "image_id": "", "image_name": "", "source_platform": 0, "team_id": "", "type": "canvas_color"}

    def make_vocal_sep(self, vid):
        return {"choice": 0, "enter_from": "", "final_algorithm": "", "id": vid, "production_path": "", "removed_sounds": [], "time_range": None, "type": "vocal_separation"}

    def make_video_mat(self, mid, path, dur_us, w=1920, h=1080):
        return {
            "aigc_type": "none", "audio_fade": None, "category_name": "local",
            "check_flag": 62978047, "crop": {"lower_left_x":0,"lower_left_y":1,"lower_right_x":1,"lower_right_y":1,"upper_left_x":0,"upper_left_y":0,"upper_right_x":1,"upper_right_y":0},
            "crop_ratio": "free", "crop_scale": 1.0, "duration": dur_us,
            "extra_type_option": 0, "has_audio": True, "height": h, "id": mid,
            "local_material_id": str(uuid.uuid4()), "material_name": os.path.basename(path),
            "path": path, "source": 0, "type": "video", "width": w,
            "reverse_path": "", "stable": {"matrix_path":"","stable_level":0,"time_range":{"duration":0,"start":0}},
            "video_algorithm": {"algorithms":[],"path":""},
        }

    def make_audio_mat(self, mid, path, dur_us):
        return {
            "app_id": 0, "category_name": "local", "check_flag": 1,
            "duration": dur_us, "id": mid, "local_material_id": self.uid(),
            "music_id": self.uid(), "name": os.path.basename(path), "path": path,
            "type": "extract_music",
        }

    def make_seg(self, sid, mid, start_us, dur_us, extras=None, track_idx=0):
        return {
            "caption_info": None, "cartoon": False,
            "clip": {"alpha":1.0,"flip":{"horizontal":False,"vertical":False},"rotation":0.0,"scale":{"x":1.0,"y":1.0},"transform":{"x":0.0,"y":0.0}},
            "common_keyframes": [], "enable_adjust": True, "enable_color_curves": True,
            "enable_color_wheels": True, "enable_lut": True, "enable_video_mask": True,
            "extra_material_refs": extras or [], "id": sid,
            "intensifies_audio": False, "is_placeholder": False,
            "last_nonzero_volume": 1.0, "material_id": mid,
            "render_index": 0, "reverse": False, "source": "segmentsourcenormal",
            "source_timerange": {"duration": dur_us, "start": 0}, "speed": 1.0,
            "target_timerange": {"duration": dur_us, "start": start_us},
            "track_attribute": 0, "track_render_index": track_idx,
            "uniform_scale": {"on": True, "value": 1.0}, "visible": True, "volume": 1.0
        }

    def make_audio_seg(self, sid, mid, start_us, dur_us, spd_id):
        return {
            "caption_info": None, "cartoon": False, "clip": None,
            "common_keyframes": [], "enable_adjust": False,
            "extra_material_refs": [spd_id], "id": sid,
            "is_placeholder": False, "last_nonzero_volume": 1.0,
            "material_id": mid, "render_index": 0,
            "source_timerange": {"duration": dur_us, "start": 0}, "speed": 1.0,
            "target_timerange": {"duration": dur_us, "start": start_us},
            "track_attribute": 0, "track_render_index": 0,
            "visible": True, "volume": 1.0
        }

    def make_text_mat(self, mid, text):
        content = json.dumps({
            "styles": [{"fill": {"alpha": 1.0, "content": {"render_type": "solid", "solid": {"alpha": 1.0, "color": [1.0, 1.0, 1.0]}}}, "font": {"id": "", "path": "C:/Windows/Fonts/arial.ttf"}, "range": [0, len(text)], "size": 5.0}],
            "text": text
        }, ensure_ascii=False)
        return {
            "id": mid, "type": "subtitle", "content": content,
            "text_color": "#FFFFFF", "font_size": 5.0, "alignment": 1,
            "layer_weight": 1, "text_alpha": 1.0, "font_path": "C:/Windows/Fonts/arial.ttf",
            "group_id": "import_" + str(int(time.time() * 1000)),
            "border_alpha": 1.0, "border_color": "", "border_width": 0.08, "border_mode": 0,
            "has_shadow": False, "shadow_alpha": 0.9,
            "shadow_angle": -45.0, "shadow_color": "", "shadow_distance": 5.0,
            "shadow_point": {"x": 0.636, "y": -0.636}, "shadow_smoothing": 0.45
        }

    def make_text_seg(self, sid, mid, start_us, dur_us, spd_id):
        return {
            "id": sid, "material_id": mid, "extra_material_refs": [spd_id],
            "target_timerange": {"duration": dur_us, "start": start_us},
            "source_timerange": None, "render_timerange": {"start": 0, "duration": 0},
            "speed": 1.0, "volume": 1.0, "last_nonzero_volume": 1.0,
            "clip": {"alpha": 1.0, "flip": {"horizontal": False, "vertical": False}, "rotation": 0.0, "scale": {"x": 1.0, "y": 1.0}, "transform": {"x": 0.0, "y": -0.8}},
            "render_index": 14000, "visible": True, "track_attribute": 0
        }

    def _run_demucs(self):
        """Run BS-RoFormer vocal separation and save vocal/instrumental WAV files."""
        import torch
        tmp = tempfile.mkdtemp()
        wav_path = os.path.join(tmp, 'audio.wav')
        subprocess.run([find_ffmpeg(), '-y', '-i', self.video_path, '-vn', wav_path],
                       capture_output=True)
        if not os.path.exists(wav_path):
            self.log.emit("  ⚠️ Không trích xuất được audio!"); return

        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        try:
            import yaml
            from bs_roformer import get_model_from_config, demix_track
            from ml_collections import ConfigDict

            roformer_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                                        "models", "roformer-model-bs-roformer-sw-by-jarredou")
            config_path = os.path.join(roformer_dir, "BS-Rofo-SW-Fixed.yaml")
            ckpt_path = os.path.join(roformer_dir, "BS-Rofo-SW-Fixed.ckpt")

            with open(config_path) as f:
                config = ConfigDict(yaml.load(f, Loader=yaml.SafeLoader))

            roformer_model = get_model_from_config("bs_roformer", config)
            roformer_model.load_state_dict(torch.load(ckpt_path, map_location="cpu"))
            roformer_model = roformer_model.to(device)
            roformer_model.eval()

            mix, sr_orig = sf.read(wav_path)
            if len(mix.shape) == 1:
                mix = np.stack([mix, mix], axis=-1)
            mixture = torch.tensor(mix.T, dtype=torch.float32)

            self.log.emit(f"  ⏳ BS-RoFormer: {mix.shape[0] / sr_orig:.0f}s audio...")
            res, _ = demix_track(config, roformer_model, mixture, device)

            vocals_np = res["vocals"].T
            instrumental_np = mix - vocals_np

            vid_base = os.path.splitext(self.video_path)[0]
            self.vocal_path = f"{vid_base}_vocal.wav".replace("\\", "/")
            self.inst_path = f"{vid_base}_instrumental.wav".replace("\\", "/")
            sf.write(self.vocal_path, vocals_np, sr_orig, subtype="FLOAT")
            sf.write(self.inst_path, instrumental_np, sr_orig, subtype="FLOAT")

            del roformer_model, mixture, res
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            self.log.emit("  ✅ Tách vocal/instrumental thành công! (BS-RoFormer)")
        except Exception as e:
            self.log.emit(f"  ⚠️ BS-RoFormer lỗi ({e}), dùng audio gốc.")
            self.vocal_path = None
            self.inst_path = None

    def run(self):
        try:
            manifest_path = os.path.join(self.dubbed_dir, "manifest.json")
            if not os.path.exists(manifest_path):
                self.err.emit("manifest.json not found!"); return

            with open(manifest_path, 'r', encoding='utf-8') as f:
                manifest = json.load(f)

            self.log.emit(f"Creating CapCut project: {self.project_name} (video speed: {self.video_speed}x)")
            vid_dur_us = self.get_video_duration_us()
            eff_dur_us = int(vid_dur_us / self.video_speed)

            # Auto-run Demucs if vocal/instrumental files don't exist
            if not self.vocal_path or not os.path.exists(self.vocal_path or ''):
                self.log.emit("🎤 Chạy Demucs tách vocal/instrumental...")
                self._run_demucs()
                self.prog.emit(30)

            project_folder = os.path.join(self.draft_dir, self.project_name)
            if os.path.exists(project_folder):
                shutil.rmtree(project_folder)
            os.makedirs(project_folder, exist_ok=True)
            
            for d in ['common_attachment', 'matting', 'smart_crop', 'adjust_mask', 'qr_upload', 'Resources', 'subdraft']:
                os.makedirs(os.path.join(project_folder, d), exist_ok=True)

            draft = {
                "canvas_config": {"background": None, "height": 1080, "ratio": "16:9", "width": 1920},
                "color_space": 0,
                "config": {
                    "adjust_max_index": 1, "attachment_info": [], "combination_max_index": 1,
                    "export_range": None, "extract_audio_last_index": 1,
                    "lyrics_recognition_id": "", "lyrics_sync": True, "lyrics_taskinfo": [],
                    "maintrack_adsorb": True, "material_save_mode": 0,
                    "multi_language_current": "none", "multi_language_list": [],
                    "multi_language_main": "none", "multi_language_mode": "none",
                    "original_sound_last_index": 1, "record_audio_last_index": 1,
                    "sticker_max_index": 1, "subtitle_keywords_config": None,
                    "subtitle_recognition_id": "", "subtitle_sync": True, "subtitle_taskinfo": [],
                    "system_font_list": [], "use_float_render": False,
                    "video_mute": False, "zoom_info_params": None
                },
                "cover": None,
                "create_time": 0,
                "duration": eff_dur_us,
                "extra_info": None,
                "fps": 30.0,
                "free_render_index_mode_on": False,
                "group_container": None,
                "id": self.uid(),
                "is_drop_frame_timecode": False,
                "keyframe_graph_list": [],
                "keyframes": {
                    "adjusts": [], "audios": [], "effects": [], "filters": [],
                    "handwrites": [], "stickers": [], "texts": [], "videos": []
                },
                "last_modified_platform": {
                    "app_id": 359289, "app_source": "cc", "app_version": "6.7.0",
                    "device_id": "", "hard_disk_id": "", "mac_address": "",
                    "os": "windows", "os_version": "10.0.19045"
                },
                "lyrics_effects": [],
                "materials": {
                    "ai_translates": [], "audio_balances": [], "audio_effects": [],
                    "audio_fades": [], "audio_track_indexes": [], "audios": [],
                    "beats": [], "canvases": [], "chromas": [], "color_curves": [],
                    "common_mask": [], "digital_humans": [], "drafts": [],
                    "effects": [], "flowers": [], "green_screens": [],
                    "handwrites": [], "hsl": [], "images": [],
                    "log_color_wheels": [], "loudnesses": [], "manual_beautys": [],
                    "manual_deformations": [], "material_animations": [],
                    "material_colors": [], "multi_language_refs": [],
                    "placeholder_infos": [], "placeholders": [], "plugin_effects": [],
                    "primary_color_wheels": [], "realtime_denoises": [],
                    "shapes": [], "smart_crops": [], "smart_relights": [],
                    "sound_channel_mappings": [], "speeds": [], "stickers": [],
                    "tail_leaders": [], "text_templates": [], "texts": [],
                    "time_marks": [], "transitions": [], "video_effects": [],
                    "video_trackings": [], "videos": [], "vocal_beautifys": [],
                    "vocal_separations": []
                },
                "mutable_config": None,
                "name": "",
                "new_version": "140.0.0",
                "path": "",
                "platform": {
                    "app_id": 359289, "app_source": "cc", "app_version": "6.7.0",
                    "device_id": "", "hard_disk_id": "", "mac_address": "",
                    "os": "windows", "os_version": "10.0.19045"
                },
                "relationships": [],
                "render_index_track_mode_on": True,
                "retouch_cover": None,
                "source": "default",
                "static_cover_image_path": "",
                "time_marks": None,
                "tracks": [],
                "uneven_animation_template_info": {
                    "composition": "", "content": "", "order": "",
                    "sub_template_info_list": []
                },
                "update_time": 0,
                "version": 360000
            }

            # ── BUILD TIMELINE WITH TIME-STRETCH LOGIC ──
            # Read dubbed WAV durations from manifest
            dub_items = []
            for item in manifest:
                wpath = os.path.join(self.dubbed_dir, item["wav"]).replace("\\", "/")
                if not os.path.exists(wpath):
                    self.log.emit(f"  SKIP missing: {item['wav']}")
                    continue
                data, sr_wav = sf.read(wpath)
                audio_dur_us = int(len(data) / sr_wav * 1_000_000)
                srt_dur_us = int((item["end_ms"] - item["start_ms"]) * 1000 / self.video_speed)
                # effective_dur includes the gap so stretch accounts for breathing room
                effective_dur_us = audio_dur_us + self.min_gap_us
                diff_s = (effective_dur_us - srt_dur_us) / 1_000_000
                dub_items.append({
                    **item,
                    "wpath": wpath,
                    "audio_dur_us": audio_dur_us,
                    "effective_dur_us": effective_dur_us,  # audio + gap
                    "srt_start_us": int(item["start_ms"] * 1000 / self.video_speed),
                    "srt_dur_us": srt_dur_us,
                    "diff_s": diff_s,
                    "needs_stretch": diff_s > self.stretch_threshold
                })
            self.prog.emit(45)

            # Build "cut points" where video needs to be split and stretched
            # new_dur_us = effective_dur (audio + gap) so video stretches enough for both
            stretches = []
            for d in dub_items:
                if d["needs_stretch"]:
                    stretches.append({
                        "src_start_us": int(d["start_ms"] * 1000),
                        "src_dur_us": int((d["end_ms"] - d["start_ms"]) * 1000),
                        "new_dur_us": d["effective_dur_us"],  # audio + gap
                        "item": d
                    })
            stretches.sort(key=lambda x: x["src_start_us"])

            self.log.emit(f"  📐 {len(stretches)} segments cần kéo dãn video (gap={self.min_gap_us/1e6:.1f}s)")

            # Build video/audio segments by splitting at stretch points
            # We walk through the video timeline, creating segments:
            # - Normal segments play at self.video_speed
            # - Stretched segments play slower to match audio duration
            def build_split_segments(material_id, total_src_dur_us, is_video=False):
                """Build split segments for video or audio track."""
                segments = []
                cursor_src = 0  # position in source (original video time, microseconds)
                cursor_target = 0  # position in target timeline

                for st in stretches:
                    st_start = st["src_start_us"]
                    st_src_dur = st["src_dur_us"]
                    st_new_dur = st["new_dur_us"]

                    # Segment BEFORE the stretch (normal speed)
                    if st_start > cursor_src:
                        gap_src = st_start - cursor_src
                        gap_target = int(gap_src / self.video_speed)
                        if gap_target > 0:
                            sp_id = self.uid()
                            draft["materials"]["speeds"].append(self.make_speed(sp_id, self.video_speed))
                            extras = [sp_id]
                            if is_video:
                                c_id, v_id = self.uid(), self.uid()
                                draft["materials"]["canvases"].append(self.make_canvas(c_id))
                                draft["materials"]["vocal_separations"].append(self.make_vocal_sep(v_id))
                                extras.extend([c_id, v_id])
                            seg = self.make_seg(self.uid(), material_id, cursor_target, gap_target, extras) if is_video else \
                                  self.make_audio_seg(self.uid(), material_id, cursor_target, gap_target, sp_id)
                            seg["source_timerange"] = {"duration": gap_src, "start": cursor_src}
                            seg["speed"] = self.video_speed
                            if is_video:
                                seg["volume"] = 0.0
                            segments.append(seg)
                            cursor_target += gap_target

                    # Stretched segment (slower speed to match audio)
                    stretch_speed = st_src_dur / st_new_dur * self.video_speed
                    sp_id = self.uid()
                    draft["materials"]["speeds"].append(self.make_speed(sp_id, stretch_speed))
                    extras = [sp_id]
                    if is_video:
                        c_id, v_id = self.uid(), self.uid()
                        draft["materials"]["canvases"].append(self.make_canvas(c_id))
                        draft["materials"]["vocal_separations"].append(self.make_vocal_sep(v_id))
                        extras.extend([c_id, v_id])
                    seg = self.make_seg(self.uid(), material_id, cursor_target, st_new_dur, extras) if is_video else \
                          self.make_audio_seg(self.uid(), material_id, cursor_target, st_new_dur, sp_id)
                    seg["source_timerange"] = {"duration": st_src_dur, "start": st_start}
                    seg["speed"] = stretch_speed
                    if is_video:
                        seg["volume"] = 0.0
                    segments.append(seg)
                    cursor_target += st_new_dur
                    cursor_src = st_start + st_src_dur

                # Final segment AFTER last stretch (normal speed)
                if cursor_src < total_src_dur_us:
                    remaining_src = total_src_dur_us - cursor_src
                    remaining_target = int(remaining_src / self.video_speed)
                    if remaining_target > 0:
                        sp_id = self.uid()
                        draft["materials"]["speeds"].append(self.make_speed(sp_id, self.video_speed))
                        extras = [sp_id]
                        if is_video:
                            c_id, v_id = self.uid(), self.uid()
                            draft["materials"]["canvases"].append(self.make_canvas(c_id))
                            draft["materials"]["vocal_separations"].append(self.make_vocal_sep(v_id))
                            extras.extend([c_id, v_id])
                        seg = self.make_seg(self.uid(), material_id, cursor_target, remaining_target, extras) if is_video else \
                              self.make_audio_seg(self.uid(), material_id, cursor_target, remaining_target, sp_id)
                        seg["source_timerange"] = {"duration": remaining_src, "start": cursor_src}
                        seg["speed"] = self.video_speed
                        if is_video:
                            seg["volume"] = 0.0
                        segments.append(seg)

                return segments

            # Track 0: Video (split at stretch points)
            vm = self.uid()
            draft["materials"]["videos"].append(self.make_video_mat(vm, self.video_path, vid_dur_us))
            vid_segs = build_split_segments(vm, vid_dur_us, is_video=True)
            draft["tracks"].append({
                "attribute": 0, "flag": 0, "id": self.uid(), "is_default_name": True,
                "name": "", "type": "video", "segments": vid_segs
            })
            self.prog.emit(55)

            # Track 1 & 2: Instrumental and Vocal (also split at same points)
            if self.inst_path and os.path.exists(self.inst_path) and self.vocal_path and os.path.exists(self.vocal_path):
                self.log.emit("Adding split Instrumental and Vocal tracks...")
                am1 = self.uid()
                draft["materials"]["audios"].append(self.make_audio_mat(am1, self.inst_path, vid_dur_us))
                inst_segs = build_split_segments(am1, vid_dur_us, is_video=False)
                draft["tracks"].append({
                    "attribute": 0, "flag": 0, "id": self.uid(), "is_default_name": True,
                    "name": "Instrumental", "type": "audio", "segments": inst_segs
                })
                am2 = self.uid()
                draft["materials"]["audios"].append(self.make_audio_mat(am2, self.vocal_path, vid_dur_us))
                vocal_segs = build_split_segments(am2, vid_dur_us, is_video=False)
                draft["tracks"].append({
                    "attribute": 0, "flag": 0, "id": self.uid(), "is_default_name": True,
                    "name": "Original Vocal", "type": "audio", "segments": vocal_segs
                })
            else:
                base = os.path.splitext(os.path.basename(self.video_path))[0]
                mp3 = os.path.join(os.path.dirname(self.video_path), f"{base}.mp3").replace("\\", "/")
                if not os.path.exists(mp3):
                    self.log.emit("Extracting original audio with ffmpeg...")
                    subprocess.run([find_ffmpeg(), "-y", "-i", self.video_path, "-q:a", "0", "-map", "a", mp3],
                                   capture_output=True)
                am = self.uid()
                draft["materials"]["audios"].append(self.make_audio_mat(am, mp3, vid_dur_us))
                orig_segs = build_split_segments(am, vid_dur_us, is_video=False)
                draft["tracks"].append({
                    "attribute": 0, "flag": 0, "id": self.uid(), "is_default_name": True,
                    "name": "Original Audio", "type": "audio", "segments": orig_segs
                })
            self.prog.emit(65)

            # Track 3+: Dubbed WAVs
            # For items that need_stretch: audio plays at 1.0x speed
            # For items that DON'T need_stretch but audio > srt: speed up audio to fit
            MAX_DUB_TRACKS = 3
            dub_tracks = [{"segs": [], "end_us": 0} for _ in range(MAX_DUB_TRACKS)]

            # We need to compute the actual target start for each dub item
            # because stretches shift the timeline
            def src_to_target(src_us):
                """Convert source video time to target timeline time, accounting for stretches."""
                target = 0
                prev_src = 0
                for st in stretches:
                    st_start = st["src_start_us"]
                    st_src_dur = st["src_dur_us"]
                    st_new_dur = st["new_dur_us"]
                    if src_us <= st_start:
                        # Before this stretch
                        target += int((src_us - prev_src) / self.video_speed)
                        return target
                    # Add normal part before stretch
                    target += int((st_start - prev_src) / self.video_speed)
                    if src_us <= st_start + st_src_dur:
                        # Inside this stretch
                        frac = (src_us - st_start) / st_src_dur
                        target += int(frac * st_new_dur)
                        return target
                    # After this stretch
                    target += st_new_dur
                    prev_src = st_start + st_src_dur
                # After all stretches
                target += int((src_us - prev_src) / self.video_speed)
                return target

            for i, d in enumerate(dub_items):
                target_start = src_to_target(int(d["start_ms"] * 1000))
                audio_dur = d["audio_dur_us"]

                if d["needs_stretch"]:
                    # Video was stretched to match audio+gap → play audio at 1.0x
                    audio_speed = 1.0
                    final_dur = audio_dur
                elif d["diff_s"] > 0:
                    # Audio+gap is longer but ≤threshold → speed up audio to fit SRT slot
                    target_end = src_to_target(int(d["end_ms"] * 1000))
                    slot_dur = target_end - target_start
                    if slot_dur > 0:
                        audio_speed = audio_dur / slot_dur
                        final_dur = slot_dur
                    else:
                        audio_speed = 1.0
                        final_dur = audio_dur
                else:
                    audio_speed = 1.0
                    final_dur = audio_dur

                # Find track with no overlap
                best_t = None
                for t_idx, t in enumerate(dub_tracks):
                    if target_start >= t["end_us"]:
                        best_t = t_idx
                        break
                if best_t is None:
                    best_t = min(range(MAX_DUB_TRACKS), key=lambda x: dub_tracks[x]["end_us"])

                mid, sid, spid = self.uid(), self.uid(), self.uid()
                draft["materials"]["audios"].append(self.make_audio_mat(mid, d["wpath"], audio_dur))
                draft["materials"]["speeds"].append(self.make_speed(spid, audio_speed))
                a_seg = self.make_audio_seg(sid, mid, target_start, final_dur, spid)
                a_seg["source_timerange"] = {"duration": audio_dur, "start": 0}
                a_seg["speed"] = audio_speed
                dub_tracks[best_t]["segs"].append(a_seg)
                dub_tracks[best_t]["end_us"] = target_start + final_dur
                self.prog.emit(65 + int((i+1)/len(dub_items)*20))

            for t in dub_tracks:
                if t["segs"]:
                    draft["tracks"].append({
                        "attribute": 0, "flag": 0, "id": self.uid(), "is_default_name": True,
                        "name": "", "type": "audio", "segments": t["segs"]
                    })

            # Subtitle track
            txt_segs = []
            for item in manifest:
                if "text" not in item: continue
                text = item["text"]
                target_start = src_to_target(int(item["start_ms"] * 1000))
                target_end = src_to_target(int(item["end_ms"] * 1000))
                dur_us = target_end - target_start
                if dur_us <= 0: continue
                mid, sid, spid = self.uid(), self.uid(), self.uid()
                draft["materials"]["texts"].append(self.make_text_mat(mid, text))
                draft["materials"]["speeds"].append(self.make_speed(spid))
                txt_segs.append(self.make_text_seg(sid, mid, target_start, dur_us, spid))

            if txt_segs:
                draft["tracks"].append({
                    "attribute": 0, "flag": 0, "id": self.uid(), "is_default_name": True,
                    "name": "", "type": "text", "segments": txt_segs
                })

            # Save
            draft_path = os.path.join(project_folder, "draft_content.json")
            with open(draft_path, 'w', encoding='utf-8') as f:
                json.dump(draft, f, ensure_ascii=False)

            ts = int(time.time() * 1_000_000)
            meta = {
                "draft_fold_path": project_folder.replace("\\","/"),
                "draft_id": draft["id"],
                "draft_name": self.project_name,
                "draft_root_path": self.draft_dir.replace("\\","/"),
                "tm_draft_create": ts,
                "tm_draft_modified": ts,
                "tm_duration": vid_dur_us
            }
            with open(os.path.join(project_folder, "draft_meta_info.json"), 'w', encoding='utf-8') as f:
                json.dump(meta, f, ensure_ascii=False)

            # Register
            root_meta_path = os.path.join(self.draft_dir, "root_meta_info.json")
            if os.path.exists(root_meta_path):
                with open(root_meta_path, 'r', encoding='utf-8') as f:
                    rm = json.load(f)
                store = [e for e in rm.get("all_draft_store", []) if e.get("draft_name") != self.project_name]
                store.insert(0, {**meta, "draft_cover": "", "draft_json_file": draft_path.replace("\\","/"),
                                 "draft_root_path": self.draft_dir.replace("\\","/")})
                rm["all_draft_store"] = store
                with open(root_meta_path, 'w', encoding='utf-8') as f:
                    json.dump(rm, f, ensure_ascii=False)

            self.log.emit(f"CapCut project created! Open CapCut to see: {self.project_name}")
            self.prog.emit(100)
            self.done.emit(project_folder)
        except Exception as ex:
            import traceback
            self.err.emit(str(ex) + "\n" + traceback.format_exc())

# ── Main App ────────────────────────────────────────────────
class App(QWidget):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("🎙️ Speaker Diarization - Ai Đang Nói?")
        self.resize(1200, 800)
        self.entries = []
        self.wav_path = None
        self.play_proc = None
        self._build_ui()

    def _build_ui(self):
        # ── Clean Light Theme ──
        self.setStyleSheet("""
            QWidget { font-size: 13px; font-family: 'Segoe UI', Arial, sans-serif; }
            QGroupBox { font-weight: bold; font-size: 14px; padding-top: 15px; margin-top: 10px; border: 1px solid #c0c0c0; border-radius: 6px; background: #fdfdfd; }
            QGroupBox::title { subcontrol-origin: margin; left: 10px; top: -5px; color: #2c3e50; }
            QPushButton { padding: 8px; border: 1px solid #b0b0b0; border-radius: 4px; background: #f0f0f0; }
            QPushButton:hover { background: #e0e0e0; border-color: #909090; }
            QLineEdit, QSpinBox, QDoubleSpinBox, QComboBox { padding: 5px; border: 1px solid #b0b0b0; border-radius: 3px; background: #fff; }
            QLineEdit:focus, QSpinBox:focus, QDoubleSpinBox:focus, QComboBox:focus { border: 1px solid #3498db; }
            QTableWidget { border: 1px solid #c0c0c0; gridline-color: #e0e0e0; alternate-background-color: #f9f9f9; }
            QHeaderView::section { background-color: #eaeaea; padding: 4px; border: 1px solid #c0c0c0; font-weight: bold; }
        """)

        main_layout = QVBoxLayout()
        self.prog = QProgressBar(); main_layout.addWidget(self.prog)

        h_panels = QHBoxLayout()

        # ── PANEL 1: INPUT & LOG (Left, 25%) ──
        panel_left = QVBoxLayout()
        
        g_input = QGroupBox("1. Đầu vào (Input)")
        v_in = QVBoxLayout()
        
        v_in.addWidget(QLabel("Video:"))
        h1 = QHBoxLayout()
        self.vid_in = QLineEdit(); self.vid_in.setPlaceholderText("File Video MP4...")
        b1 = QPushButton("📹 Chọn"); b1.clicked.connect(self._sel_vid)
        h1.addWidget(self.vid_in, 3); h1.addWidget(b1, 1)
        v_in.addLayout(h1)
        
        v_in.addWidget(QLabel("Phụ đề SRT:"))
        h2 = QHBoxLayout()
        self.srt_in = QLineEdit(); self.srt_in.setPlaceholderText("File SRT phụ đề...")
        b2 = QPushButton("📝 Chọn"); b2.clicked.connect(self._sel_srt)
        h2.addWidget(self.srt_in, 3); h2.addWidget(b2, 1)
        v_in.addLayout(h2)
        
        v_in.addSpacing(10)
        h_spk = QHBoxLayout()
        h_spk.addWidget(QLabel("Số người nói (0=Auto):"))
        self.spb_num_spk = QSpinBox(); self.spb_num_spk.setRange(0, 10); self.spb_num_spk.setValue(0)
        h_spk.addWidget(self.spb_num_spk)
        v_in.addLayout(h_spk)
        
        v_in.addSpacing(10)
        self.btn_go = QPushButton("🔍 PHÂN TÍCH GIỌNG NÓI")
        self.btn_go.setStyleSheet("background:#27ae60;color:#fff;padding:12px;font-size:15px;font-weight:bold;border:none;")
        self.btn_go.clicked.connect(self._analyze)
        v_in.addWidget(self.btn_go)
        
        g_input.setLayout(v_in)
        panel_left.addWidget(g_input)

        g_log = QGroupBox("Console Log")
        v_log = QVBoxLayout()
        self.console = QTextEdit(); self.console.setReadOnly(True)
        self.console.setStyleSheet("background:#f4f6f7;color:#2c3e50;font-family:Consolas;font-size:12px;border:1px solid #ccc;")
        v_log.addWidget(self.console)
        g_log.setLayout(v_log)
        panel_left.addWidget(g_log, 1)

        h_panels.addLayout(panel_left, 1)

        # ── PANEL 2: ANALYSIS (Middle, 50%) ──
        panel_mid = QVBoxLayout()
        g_analysis = QGroupBox("2. Phân Tích (Analysis)")
        v_ana = QVBoxLayout()
        
        self.timeline = TimelineWidget()
        self.timeline.setMinimumHeight(80) # Larger timeline
        self.timeline.clicked.connect(self._on_tl_click)
        v_ana.addWidget(self.timeline)

        self.legend = QLabel("")
        v_ana.addWidget(self.legend)

        self.table = QTableWidget()
        self.table.setColumnCount(5)
        self.table.setHorizontalHeaderLabels(['#', 'Thời gian', 'Nội dung', 'Speaker', '▶'])
        self.table.horizontalHeader().setSectionResizeMode(2, QHeaderView.Stretch)
        self.table.setSelectionBehavior(QAbstractItemView.SelectRows)
        self.table.setAlternatingRowColors(True)
        self.table.setColumnWidth(0, 40); self.table.setColumnWidth(1, 120)
        self.table.setColumnWidth(3, 120); self.table.setColumnWidth(4, 40)
        self.table.cellClicked.connect(self._on_cell)
        v_ana.addWidget(self.table)
        
        g_analysis.setLayout(v_ana)
        panel_mid.addWidget(g_analysis)
        h_panels.addLayout(panel_mid, 2)

        # ── PANEL 3: OUTPUT (Right, 25%) ──
        panel_right = QVBoxLayout()
        g_output = QGroupBox("3. Lồng Tiếng & Xuất (Output)")
        v_out = QVBoxLayout()

        v_out.addWidget(QLabel("🗣️ Gán Giọng Đọc (Voice Mapping):"))
        self.voice_map_layout = QVBoxLayout() # Vertical list of voices

        # Wrap in a scroll area so it doesn't break the layout with many speakers
        scroll_area = QScrollArea()
        scroll_area.setWidgetResizable(True)
        scroll_area.setFrameShape(QFrame.NoFrame)
        scroll_area.setStyleSheet("QScrollArea { background: transparent; border: none; }")
        
        scroll_content = QWidget()
        scroll_content.setStyleSheet("background: transparent;")
        scroll_content.setLayout(self.voice_map_layout)
        scroll_area.setWidget(scroll_content)
        
        v_out.addWidget(scroll_area, 1) # Give it stretch factor 1 so it expands

        v_out.addSpacing(15)
        v_out.addWidget(QLabel("⚙️ Cài đặt TTS:"))
        h_spd = QHBoxLayout(); h_spd.addWidget(QLabel("Speed:")); self.spb_speed = QDoubleSpinBox(); self.spb_speed.setRange(0.5, 2.0); self.spb_speed.setValue(1.0); self.spb_speed.setSingleStep(0.1); h_spd.addWidget(self.spb_speed); v_out.addLayout(h_spd)
        h_stp = QHBoxLayout(); h_stp.addWidget(QLabel("Steps:")); self.spb_steps = QSpinBox(); self.spb_steps.setRange(8, 100); self.spb_steps.setValue(32); h_stp.addWidget(self.spb_steps); v_out.addLayout(h_stp)
        h_gui = QHBoxLayout(); h_gui.addWidget(QLabel("Guidance:")); self.spb_guidance = QDoubleSpinBox(); self.spb_guidance.setRange(1.0, 10.0); self.spb_guidance.setValue(3.0); self.spb_guidance.setSingleStep(0.5); h_gui.addWidget(self.spb_guidance); v_out.addLayout(h_gui)

        v_out.addSpacing(15)
        self.btn_dub = QPushButton("🔊 TẠO WAV LỒNG TIẾNG"); self.btn_dub.setEnabled(False)
        self.btn_dub.setStyleSheet("background:#e67e22;color:#fff;padding:12px;font-size:14px;font-weight:bold;border:none;")
        self.btn_dub.clicked.connect(self._dubbing)
        v_out.addWidget(self.btn_dub)

        v_out.addSpacing(20)
        v_out.addWidget(QLabel("📦 CapCut Project Name:"))
        self.proj_name = QLineEdit("Dubbed_Project")
        v_out.addWidget(self.proj_name)
        self.btn_capcut = QPushButton("🎬 XUẤT VÀO CAPCUT"); self.btn_capcut.setEnabled(True)
        self.btn_capcut.setStyleSheet("background:#3498db;color:#fff;padding:12px;font-size:14px;font-weight:bold;border:none;")
        self.btn_capcut.clicked.connect(self._export_capcut)
        v_out.addWidget(self.btn_capcut)

        v_out.addStretch(1)
        self.btn_ex = QPushButton("💾 LƯU FILE SRT MỚI"); self.btn_ex.setEnabled(False)
        self.btn_ex.clicked.connect(self._export)
        v_out.addWidget(self.btn_ex)

        g_output.setLayout(v_out)
        panel_right.addWidget(g_output)
        h_panels.addLayout(panel_right, 1)

        main_layout.addLayout(h_panels)
        self.setLayout(main_layout)

    def _sel_vid(self):
        f, _ = QFileDialog.getOpenFileName(self, "Chọn Video", "", "Video (*.mp4 *.mkv *.avi)")
        if f: self.vid_in.setText(f)

    def _sel_srt(self):
        f, _ = QFileDialog.getOpenFileName(self, "Chọn SRT", "", "SRT (*.srt)")
        if f: self.srt_in.setText(f)

    def _log(self, t): self.console.append(t)

    def _analyze(self):
        vid = self.vid_in.text(); srt = self.srt_in.text()
        if not vid or not srt:
            self._log("❌ Chọn cả Video và SRT!"); return
        entries = parse_srt(srt)
        if not entries:
            self._log("❌ Không parse được SRT!"); return
        self._log(f"📋 Đã đọc {len(entries)} dòng phụ đề.")

        self.btn_ex.setEnabled(False)
        self.btn_go.setEnabled(False)
        self.btn_go.setText("⏳ ĐANG PHÂN TÍCH...")
        self.prog.setValue(0)
        self.console.clear()

        expected_k = self.spb_num_spk.value()
        self.worker = AnalysisWorker(vid, entries, expected_k)
        self.worker.log.connect(self._log)
        self.worker.prog.connect(self.prog.setValue)
        self.worker.done.connect(self._on_done)
        self.worker.err.connect(self._on_err)
        self.worker.start()

    def _on_done(self, entries, n_spk):
        self.entries = entries
        self.wav_path = getattr(self.worker, '_wav_path', None)
        self.vocal_path = getattr(self.worker, 'vocal_path', None)
        self.inst_path = getattr(self.worker, 'inst_path', None)
        self._populate_table()
        total_ms = max(e['end_ms'] for e in entries) if entries else 1
        self.timeline.set_data(entries, total_ms)
        # Legend
        parts = []
        for i in range(n_spk):
            c = COLORS[i % len(COLORS)]
            parts.append(f'<span style="background:{c.name()};color:#000;padding:3px 8px;border-radius:3px;'
                         f'margin:0 4px;font-weight:bold; border: 1px solid #aaa;">{NAMES[i]}</span>')
        self.legend.setText("  ".join(parts))
        self.legend.setTextFormat(Qt.RichText)
        self.btn_ex.setEnabled(True)
        self.btn_go.setEnabled(True)
        self.btn_go.setText("🔍 PHÂN TÍCH GIỌNG NÓI")
        self._build_voice_combos(n_spk)
        self.btn_dub.setEnabled(True)

    def _on_err(self, e):
        self._log(f"❌ {e}")
        self.btn_go.setEnabled(True)
        self.btn_go.setText("🔍 PHÂN TÍCH GIỌNG NÓI")

    def _populate_table(self):
        self.table.setRowCount(len(self.entries))
        for i, e in enumerate(self.entries):
            spk = e.get('speaker', 0) or 0
            bg = COLORS[spk % len(COLORS)]
            bg_light = QColor(bg.red(), bg.green(), bg.blue(), 30) # lighter for light mode

            # Index
            it = QTableWidgetItem(str(e['index']))
            it.setFlags(it.flags() & ~Qt.ItemIsEditable)
            it.setBackground(bg_light)
            self.table.setItem(i, 0, it)

            # Time
            it = QTableWidgetItem(f"{ms_fmt(e['start_ms'])} → {ms_fmt(e['end_ms'])}")
            it.setFlags(it.flags() & ~Qt.ItemIsEditable)
            it.setBackground(bg_light)
            self.table.setItem(i, 1, it)

            # Text
            it = QTableWidgetItem(e['text'])
            it.setFlags(it.flags() & ~Qt.ItemIsEditable)
            it.setBackground(bg_light)
            self.table.setItem(i, 2, it)

            # Speaker combo
            combo = NoScrollComboBox()
            combo.addItems(NAMES[:max(2, max(x.get('speaker',0) or 0 for x in self.entries)+2)])
            combo.setCurrentIndex(spk)
            combo.currentIndexChanged.connect(lambda val, row=i: self._change_speaker(row, val))
            self.table.setCellWidget(i, 3, combo)

            # Play button
            btn = QPushButton("▶")
            btn.setFixedWidth(35)
            btn.clicked.connect(lambda _, row=i: self._play(row))
            self.table.setCellWidget(i, 4, btn)

    def _change_speaker(self, row, val):
        self.entries[row]['speaker'] = val
        # Refresh row colors
        bg = COLORS[val % len(COLORS)]
        bg_light = QColor(bg.red(), bg.green(), bg.blue(), 30) # lighter for light mode
        for c in range(3):
            it = self.table.item(row, c)
            if it: it.setBackground(bg_light)
        total_ms = max(e['end_ms'] for e in self.entries)
        self.timeline.set_data(self.entries, total_ms)

    def _play(self, row):
        if not self.wav_path or not os.path.exists(self.wav_path): return
        e = self.entries[row]
        ss = e['start_ms'] / 1000
        dur = (e['end_ms'] - e['start_ms']) / 1000
        # Extract segment to temp wav and play with winsound
        import winsound
        winsound.PlaySound(None, winsound.SND_PURGE)  # Stop previous
        tmp_seg = os.path.join(tempfile.gettempdir(), '_spk_preview.wav')
        subprocess.run([find_ffmpeg(), '-y', '-ss', str(ss), '-t', str(dur),
                        '-i', self.wav_path, '-ar', '16000', '-ac', '1', tmp_seg],
                       capture_output=True)
        if os.path.exists(tmp_seg):
            winsound.PlaySound(tmp_seg, winsound.SND_FILENAME | winsound.SND_ASYNC)

    def _on_tl_click(self, idx):
        self.table.selectRow(idx)
        self.table.scrollToItem(self.table.item(idx, 0))
        self.timeline.highlight(idx)

    def _on_cell(self, row, col):
        self.timeline.highlight(row)
        if col == 4:
            self._play(row)

    def _export(self):
        if not self.entries: return
        f, _ = QFileDialog.getSaveFileName(self, "Lưu SRT mới", "", "SRT (*.srt)")
        if not f: return
        with open(f, 'w', encoding='utf-8') as out:
            for e in self.entries:
                spk = NAMES[e.get('speaker', 0) or 0]
                out.write(f"{e['index']}\n")
                out.write(f"{ms_srt(e['start_ms'])} --> {ms_srt(e['end_ms'])}\n")
                out.write(f"[{spk}] {e['text']}\n\n")
        self._log(f"💾 Đã xuất: {f}")
        QMessageBox.information(self, "Thành công", f"Đã lưu SRT mới tại:\n{f}")

    def _build_voice_combos(self, n_spk):
        # Clear old combos
        while self.voice_map_layout.count():
            item = self.voice_map_layout.takeAt(0)
            w = item.widget()
            if w: w.deleteLater()
        self.voice_combos = {}

        voices_cfg = load_voices_config()
        voice_names = list(voices_cfg.keys())
        if not voice_names:
            self._log("No voices found in voices.json!")
            return

        for i in range(n_spk):
            lbl = QLabel(f"{NAMES[i]}:")
            combo = NoScrollComboBox()
            combo.addItems(voice_names)
            if i < len(voice_names):
                combo.setCurrentIndex(i % len(voice_names))
            self.voice_combos[i] = combo
            btn_preview = QPushButton("▶")
            btn_preview.setFixedWidth(30)
            btn_preview.clicked.connect(lambda _, c=combo: self._preview_voice(c.currentText()))
            self.voice_map_layout.addWidget(lbl)
            self.voice_map_layout.addWidget(combo)
            self.voice_map_layout.addWidget(btn_preview)
        
        # Add stretch at the end to push all voice combos to the top
        self.voice_map_layout.addStretch(1)

    def _preview_voice(self, voice_name):
        voices_cfg = load_voices_config()
        vdata = voices_cfg.get(voice_name)
        if not vdata:
            return
        apath = os.path.join(TTS_DIR, "voice", vdata["audio"])
        if os.path.exists(apath):
            import winsound
            winsound.PlaySound(None, winsound.SND_PURGE)
            winsound.PlaySound(apath, winsound.SND_FILENAME | winsound.SND_ASYNC)
        else:
            self._log(f"File not found: {apath}")

    def _dubbing(self):
        if not self.entries:
            self._log("No entries!"); return
        vid = self.vid_in.text()
        out_dir = os.path.join(os.path.dirname(vid), "dubbed_wavs")
        voice_map = {k: v.currentText() for k, v in self.voice_combos.items()}
        
        tts_cfg = {
            'speed': self.spb_speed.value(),
            'steps': self.spb_steps.value(),
            'guidance': self.spb_guidance.value()
        }

        self.btn_dub.setEnabled(False)
        self.btn_dub.setText("DANG TAO WAV...")
        self.prog.setValue(0)

        self.dub_worker = DubbingWorker(self.entries, voice_map, out_dir, tts_cfg)
        self.dub_worker.log.connect(self._log)
        self.dub_worker.prog.connect(self.prog.setValue)
        self.dub_worker.done.connect(self._on_dub_done)
        self.dub_worker.err.connect(self._on_err)
        self.dub_worker.start()

    def _on_dub_done(self, out_dir):
        self._dubbed_dir = out_dir
        self.btn_dub.setEnabled(True)
        self.btn_dub.setText("TẠO WAV LỒNG TIẾNG")
        self.btn_capcut.setEnabled(True)
        self._log(f"WAVs ready at: {out_dir}")
        QMessageBox.information(self, "Done", f"WAVs saved to:\n{out_dir}")

    def _export_capcut(self):
        vid = self.vid_in.text()
        if not vid or not os.path.exists(vid):
            self._log("Chưa chọn Video!"); return

        # ── Popup Dialog ──
        dlg = QDialog(self)
        dlg.setWindowTitle("📦 Xuất vào CapCut")
        dlg.setMinimumWidth(500)
        lay = QVBoxLayout()

        # Folder picker
        lay.addWidget(QLabel("Thư mục chứa WAV lồng tiếng:"))
        h_folder = QHBoxLayout()
        default_dir = getattr(self, '_dubbed_dir', None) or os.path.join(os.path.dirname(vid), 'dubbed_wavs')
        folder_in = QLineEdit(default_dir)
        btn_browse = QPushButton("📁 Chọn...")
        btn_browse.clicked.connect(lambda: folder_in.setText(
            QFileDialog.getExistingDirectory(dlg, "Chọn thư mục WAV", os.path.dirname(vid)) or folder_in.text()))
        h_folder.addWidget(folder_in, 3); h_folder.addWidget(btn_browse, 1)
        lay.addLayout(h_folder)

        # Video speed
        h_spd = QHBoxLayout()
        h_spd.addWidget(QLabel("Tốc độ video gốc (speed):"))
        spd_box = QDoubleSpinBox()
        spd_box.setRange(0.1, 3.0)
        spd_box.setValue(1.0)
        spd_box.setSingleStep(0.1)
        h_spd.addWidget(spd_box)
        h_spd.addWidget(QLabel("(0.8 = chậm hơn, 1.0 = bình thường)"))
        lay.addLayout(h_spd)

        # Stretch threshold
        h_stretch = QHBoxLayout()
        h_stretch.addWidget(QLabel("Ngưỡng kéo dãn video (s):"))
        stretch_box = QDoubleSpinBox()
        stretch_box.setRange(0.0, 5.0)
        stretch_box.setValue(0.3)
        stretch_box.setSingleStep(0.1)
        stretch_box.setDecimals(1)
        h_stretch.addWidget(stretch_box)
        h_stretch.addWidget(QLabel("(Audio dài hơn SRT > ngưỡng này → kéo chậm video)"))
        lay.addLayout(h_stretch)

        # Min gap between audio clips
        h_gap = QHBoxLayout()
        h_gap.addWidget(QLabel("Khoảng nghỉ giữa câu (s):"))
        gap_box = QDoubleSpinBox()
        gap_box.setRange(0.0, 3.0)
        gap_box.setValue(0.3)
        gap_box.setSingleStep(0.1)
        gap_box.setDecimals(1)
        h_gap.addWidget(gap_box)
        h_gap.addWidget(QLabel("(Khoảng cách tối thiểu giữa 2 câu audio TTS)"))
        lay.addLayout(h_gap)

        # Project name
        h_name = QHBoxLayout()
        h_name.addWidget(QLabel("Tên Project CapCut:"))
        name_in = QLineEdit(self.proj_name.text().strip() or "Dubbed_Project")
        h_name.addWidget(name_in)
        lay.addLayout(h_name)

        # OK / Cancel
        h_btns = QHBoxLayout()
        btn_ok = QPushButton("✅ TẠO PROJECT"); btn_ok.setStyleSheet("background:#27ae60;color:#fff;padding:10px;font-weight:bold;")
        btn_ok.clicked.connect(dlg.accept)
        btn_cancel = QPushButton("Hủy")
        btn_cancel.clicked.connect(dlg.reject)
        h_btns.addWidget(btn_ok); h_btns.addWidget(btn_cancel)
        lay.addLayout(h_btns)

        dlg.setLayout(lay)
        if dlg.exec_() != QDialog.Accepted:
            return

        dubbed_dir = folder_in.text()
        if not os.path.exists(os.path.join(dubbed_dir, 'manifest.json')):
            self._log("Thư mục này không có manifest.json!"); return

        proj = name_in.text().strip() or "Dubbed_Project"
        v_speed = spd_box.value()
        v_stretch = stretch_box.value()
        v_gap = gap_box.value()

        self.btn_capcut.setEnabled(False)
        self.btn_capcut.setText("DANG XUẤT...")
        self.prog.setValue(0)

        # Auto-detect vocal/instrumental files if not in memory
        vocal_p = getattr(self, 'vocal_path', None)
        inst_p = getattr(self, 'inst_path', None)
        if not vocal_p or not os.path.exists(vocal_p or ''):
            vid_base = os.path.splitext(vid)[0]
            candidate_vocal = f"{vid_base}_vocal.wav"
            candidate_inst = f"{vid_base}_instrumental.wav"
            if os.path.exists(candidate_vocal) and os.path.exists(candidate_inst):
                vocal_p = candidate_vocal
                inst_p = candidate_inst
                self._log(f"  ✅ Tìm thấy file tách sẵn: {os.path.basename(candidate_vocal)}, {os.path.basename(candidate_inst)}")
            else:
                self._log("  ⚠️ Không tìm thấy file vocal/instrumental. Dùng audio gốc.")

        self.capcut_worker = CapCutWorker(vid, dubbed_dir, proj, v_speed, vocal_p, inst_p, v_stretch, v_gap)
        self.capcut_worker.log.connect(self._log)
        self.capcut_worker.prog.connect(self.prog.setValue)
        self.capcut_worker.done.connect(self._on_capcut_done)
        self.capcut_worker.err.connect(self._on_err)
        self.capcut_worker.start()

    def _on_capcut_done(self, folder):
        self.btn_capcut.setEnabled(True)
        self.btn_capcut.setText("XUẤT VÀO CAPCUT")
        QMessageBox.information(self, "Done", f"CapCut project created!\n{folder}")

if __name__ == '__main__':
    app = QApplication(sys.argv)
    w = App()
    w.show()
    sys.exit(app.exec_())