"""
SubtitleBurner — Auto-transcribe, translate to English, burn subtitles into video.
Uses faster-whisper for GPU/CPU transcription and ffmpeg for subtitle burning.
"""

import os
import sys
import time
import json
import shutil
import subprocess
import threading
import tempfile
import platform
from pathlib import Path

import tkinter as tk
from tkinter import ttk, filedialog, messagebox, scrolledtext

# ──────────────────────────────────────────────
# Resource detection
# ──────────────────────────────────────────────

def detect_resources():
    """Detect GPU, CUDA, RAM and pick the best whisper device/compute_type."""
    info = {
        "platform": platform.system(),
        "cpu_cores": os.cpu_count(),
        "ram_gb": None,
        "gpu": None,
        "cuda": False,
        "device": "cpu",
        "compute_type": "int8",
        "model_size": "medium",
    }

    # RAM
    try:
        import psutil
        info["ram_gb"] = round(psutil.virtual_memory().total / 1e9, 1)
    except ImportError:
        pass

    # CUDA / GPU
    try:
        import torch
        if torch.cuda.is_available():
            info["cuda"] = True
            info["gpu"] = torch.cuda.get_device_name(0)
            vram = torch.cuda.get_device_properties(0).total_memory / 1e9
            info["vram_gb"] = round(vram, 1)
            if vram >= 8:
                info["device"] = "cuda"
                info["compute_type"] = "float16"
                info["model_size"] = "large-v3"
            elif vram >= 4:
                info["device"] = "cuda"
                info["compute_type"] = "float16"
                info["model_size"] = "medium"
            else:
                info["device"] = "cuda"
                info["compute_type"] = "int8"
                info["model_size"] = "small"
        elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            info["gpu"] = "Apple Silicon MPS"
            info["device"] = "cpu"  # faster-whisper uses CPU on MPS
            info["compute_type"] = "int8"
            info["model_size"] = "medium"
    except ImportError:
        pass

    # CPU fallback model sizing
    if info["device"] == "cpu":
        ram = info.get("ram_gb") or 4
        if ram >= 16:
            info["model_size"] = "medium"
        elif ram >= 8:
            info["model_size"] = "small"
        else:
            info["model_size"] = "base"

    return info


def check_ffmpeg():
    """Return True if ffmpeg is available."""
    return shutil.which("ffmpeg") is not None


# ──────────────────────────────────────────────
# Core processing
# ──────────────────────────────────────────────

def transcribe_video(video_path, device, compute_type, model_size, translate, progress_cb, log_cb):
    """
    Run faster-whisper on the video, return list of (start, end, text) segments.
    If translate=True, translate to English.
    """
    from faster_whisper import WhisperModel

    log_cb(f"Loading Whisper model '{model_size}' on {device} ({compute_type})…")
    model = WhisperModel(model_size, device=device, compute_type=compute_type)

    task = "translate" if translate else "transcribe"
    log_cb(f"Starting {task}… (this may take a while for long videos)")

    segments_iter, info = model.transcribe(
        video_path,
        task=task,
        beam_size=5,
        vad_filter=True,
        word_timestamps=False,
    )

    log_cb(f"Detected language: {info.language} (confidence {info.language_probability:.0%})")

    segments = []
    for seg in segments_iter:
        segments.append((seg.start, seg.end, seg.text.strip()))
        progress_cb(min(95, int(seg.end / info.duration * 80) + 5) if info.duration else 50)
        log_cb(f"  [{seg.start:.1f}s → {seg.end:.1f}s] {seg.text.strip()}")

    return segments, info.language


def segments_to_srt(segments):
    """Convert segments list to SRT format string."""
    def fmt_time(t):
        h = int(t // 3600)
        m = int((t % 3600) // 60)
        s = int(t % 60)
        ms = int((t - int(t)) * 1000)
        return f"{h:02d}:{m:02d}:{s:02d},{ms:03d}"

    lines = []
    for i, (start, end, text) in enumerate(segments, 1):
        lines.append(str(i))
        lines.append(f"{fmt_time(start)} --> {fmt_time(end)}")
        # Wrap long lines
        words = text.split()
        wrapped = []
        line = ""
        for w in words:
            if len(line) + len(w) + 1 > 42:
                wrapped.append(line.strip())
                line = w + " "
            else:
                line += w + " "
        if line.strip():
            wrapped.append(line.strip())
        lines.append("\n".join(wrapped))
        lines.append("")
    return "\n".join(lines)


def burn_subtitles(video_path, srt_path, output_path, progress_cb, log_cb):
    """Use ffmpeg to burn subtitles into video."""
    log_cb("Burning subtitles into video with ffmpeg…")

    # Escape path for ffmpeg subtitles filter (Windows needs special handling)
    srt_escaped = str(srt_path).replace("\\", "/").replace(":", "\\:")

    cmd = [
        "ffmpeg", "-y",
        "-i", str(video_path),
        "-vf", f"subtitles='{srt_escaped}':force_style='FontSize=22,PrimaryColour=&HFFFFFF,OutlineColour=&H000000,Outline=2,Shadow=1,Alignment=2'",
        "-c:a", "copy",
        "-preset", "fast",
        str(output_path),
    ]

    log_cb(f"Running: {' '.join(cmd[:6])} …")

    process = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        universal_newlines=True,
    )

    duration_sec = None
    for line in process.stdout:
        line = line.strip()
        if "Duration:" in line and duration_sec is None:
            try:
                dur_str = line.split("Duration:")[1].split(",")[0].strip()
                h, m, s = dur_str.split(":")
                duration_sec = int(h) * 3600 + int(m) * 60 + float(s)
            except Exception:
                pass
        if "time=" in line and duration_sec:
            try:
                time_str = line.split("time=")[1].split(" ")[0]
                h, m, s = time_str.split(":")
                current = int(h) * 3600 + int(m) * 60 + float(s)
                pct = int(95 + (current / duration_sec) * 4)
                progress_cb(min(99, pct))
            except Exception:
                pass

    process.wait()
    if process.returncode != 0:
        raise RuntimeError("ffmpeg failed — check the log for details.")
    log_cb("✓ Subtitle burning complete.")


# ──────────────────────────────────────────────
# GUI
# ──────────────────────────────────────────────

class App(tk.Tk):
    def __init__(self):
        super().__init__()
        self.title("SubtitleBurner")
        self.resizable(True, True)
        self.minsize(680, 520)

        self.resources = None
        self.processing = False

        self._build_ui()
        self._detect_resources_async()

    # ── UI construction ──────────────────────

    def _build_ui(self):
        self.configure(bg="#FAFAF8")
        style = ttk.Style(self)
        style.theme_use("clam")
        style.configure("TFrame", background="#FAFAF8")
        style.configure("TLabel", background="#FAFAF8", font=("Segoe UI", 10))
        style.configure("TButton", font=("Segoe UI", 10), padding=6)
        style.configure("Accent.TButton", font=("Segoe UI", 10, "bold"), padding=6)
        style.configure("TCheckbutton", background="#FAFAF8", font=("Segoe UI", 10))
        style.configure("TLabelframe", background="#FAFAF8", font=("Segoe UI", 10, "bold"))
        style.configure("TLabelframe.Label", background="#FAFAF8", font=("Segoe UI", 10, "bold"))
        style.configure("TProgressbar", thickness=8)

        root = ttk.Frame(self, padding="16 16 16 16")
        root.pack(fill="both", expand=True)

        # ── Header
        hdr = ttk.Frame(root)
        hdr.pack(fill="x", pady=(0, 12))
        ttk.Label(hdr, text="SubtitleBurner", font=("Segoe UI", 18, "bold")).pack(side="left")
        ttk.Label(hdr, text="Auto-transcribe · Translate · Burn", foreground="#888", font=("Segoe UI", 10)).pack(side="left", padx=12, pady=4)

        # ── Resources panel
        res_frame = ttk.LabelFrame(root, text="Detected resources", padding=8)
        res_frame.pack(fill="x", pady=(0, 10))
        self.res_label = ttk.Label(res_frame, text="Detecting…", foreground="#666")
        self.res_label.pack(anchor="w")

        # ── Input file
        inp_frame = ttk.LabelFrame(root, text="Input video", padding=8)
        inp_frame.pack(fill="x", pady=(0, 10))

        row = ttk.Frame(inp_frame)
        row.pack(fill="x")
        self.input_var = tk.StringVar()
        ttk.Entry(row, textvariable=self.input_var, font=("Segoe UI", 10)).pack(side="left", fill="x", expand=True)
        ttk.Button(row, text="Browse…", command=self._browse_input).pack(side="left", padx=(6, 0))

        # ── Output file
        out_frame = ttk.LabelFrame(root, text="Output video", padding=8)
        out_frame.pack(fill="x", pady=(0, 10))

        row2 = ttk.Frame(out_frame)
        row2.pack(fill="x")
        self.output_var = tk.StringVar()
        ttk.Entry(row2, textvariable=self.output_var, font=("Segoe UI", 10)).pack(side="left", fill="x", expand=True)
        ttk.Button(row2, text="Browse…", command=self._browse_output).pack(side="left", padx=(6, 0))

        # ── Options
        opt_frame = ttk.LabelFrame(root, text="Options", padding=8)
        opt_frame.pack(fill="x", pady=(0, 10))

        opts_row = ttk.Frame(opt_frame)
        opts_row.pack(fill="x")

        self.translate_var = tk.BooleanVar(value=True)
        ttk.Checkbutton(opts_row, text="Translate to English", variable=self.translate_var).pack(side="left")

        ttk.Label(opts_row, text="   Model override:").pack(side="left")
        self.model_var = tk.StringVar(value="auto")
        model_combo = ttk.Combobox(opts_row, textvariable=self.model_var, width=12, state="readonly",
                                   values=["auto", "tiny", "base", "small", "medium", "large-v3"])
        model_combo.pack(side="left", padx=6)

        ttk.Label(opts_row, text="  Save .srt:").pack(side="left")
        self.save_srt_var = tk.BooleanVar(value=True)
        ttk.Checkbutton(opts_row, variable=self.save_srt_var).pack(side="left")

        # ── Progress
        prog_frame = ttk.LabelFrame(root, text="Progress", padding=8)
        prog_frame.pack(fill="x", pady=(0, 10))

        self.progress_var = tk.IntVar(value=0)
        self.progress_bar = ttk.Progressbar(prog_frame, variable=self.progress_var, maximum=100, length=400)
        self.progress_bar.pack(fill="x")
        self.status_label = ttk.Label(prog_frame, text="Ready.", foreground="#555")
        self.status_label.pack(anchor="w", pady=(4, 0))

        # ── Log
        log_frame = ttk.LabelFrame(root, text="Log", padding=8)
        log_frame.pack(fill="both", expand=True, pady=(0, 10))
        self.log_box = scrolledtext.ScrolledText(log_frame, height=8, font=("Courier New", 9),
                                                 bg="#F4F3EF", relief="flat", wrap="word")
        self.log_box.pack(fill="both", expand=True)

        # ── Buttons
        btn_row = ttk.Frame(root)
        btn_row.pack(fill="x")
        self.run_btn = ttk.Button(btn_row, text="▶  Start Processing", command=self._start,
                                  style="Accent.TButton")
        self.run_btn.pack(side="left")
        ttk.Button(btn_row, text="Clear log", command=self._clear_log).pack(side="left", padx=8)

        # FFmpeg warning
        if not check_ffmpeg():
            self._log("⚠  ffmpeg not found in PATH — please install it (see README).")

    # ── Resource detection ───────────────────

    def _detect_resources_async(self):
        def _run():
            try:
                r = detect_resources()
                self.resources = r
                parts = [f"CPU: {r['cpu_cores']} cores"]
                if r.get("ram_gb"):
                    parts.append(f"RAM: {r['ram_gb']} GB")
                if r.get("gpu"):
                    parts.append(f"GPU: {r['gpu']}")
                    if r.get("vram_gb"):
                        parts.append(f"VRAM: {r['vram_gb']} GB")
                parts.append(f"→ Will use: {r['device'].upper()} / {r['compute_type']} / model={r['model_size']}")
                self.after(0, lambda: self.res_label.config(text="   ".join(parts), foreground="#333"))
                self.after(0, lambda: self._log("Resource detection complete: " + " | ".join(parts)))
            except Exception as e:
                self.after(0, lambda: self.res_label.config(text=f"Detection error: {e}", foreground="red"))
        threading.Thread(target=_run, daemon=True).start()

    # ── File browsing ────────────────────────

    def _browse_input(self):
        path = filedialog.askopenfilename(
            title="Select video file",
            filetypes=[("Video files", "*.mp4 *.mkv *.mov *.avi *.webm *.m4v *.flv"), ("All files", "*.*")]
        )
        if path:
            self.input_var.set(path)
            if not self.output_var.get():
                p = Path(path)
                self.output_var.set(str(p.parent / (p.stem + "_subtitled" + p.suffix)))

    def _browse_output(self):
        path = filedialog.asksaveasfilename(
            title="Save output video",
            defaultextension=".mp4",
            filetypes=[("MP4 video", "*.mp4"), ("MKV video", "*.mkv"), ("All files", "*.*")]
        )
        if path:
            self.output_var.set(path)

    # ── Logging ──────────────────────────────

    def _log(self, msg):
        def _do():
            self.log_box.insert("end", msg + "\n")
            self.log_box.see("end")
        self.after(0, _do)

    def _clear_log(self):
        self.log_box.delete("1.0", "end")

    def _set_status(self, msg):
        self.after(0, lambda: self.status_label.config(text=msg))

    def _set_progress(self, pct):
        self.after(0, lambda: self.progress_var.set(pct))

    # ── Main processing ──────────────────────

    def _start(self):
        if self.processing:
            return

        video_in = self.input_var.get().strip()
        video_out = self.output_var.get().strip()

        if not video_in:
            messagebox.showerror("Error", "Please select an input video.")
            return
        if not Path(video_in).exists():
            messagebox.showerror("Error", f"Input file not found:\n{video_in}")
            return
        if not video_out:
            messagebox.showerror("Error", "Please specify an output path.")
            return
        if not check_ffmpeg():
            messagebox.showerror("Error", "ffmpeg not found in PATH.\nPlease install ffmpeg and ensure it is on your PATH.")
            return

        self.processing = True
        self.run_btn.config(state="disabled", text="Processing…")
        self._set_progress(0)
        threading.Thread(target=self._process_thread, args=(video_in, video_out), daemon=True).start()

    def _process_thread(self, video_in, video_out):
        try:
            r = self.resources or detect_resources()
            model_size = self.model_var.get()
            if model_size == "auto":
                model_size = r["model_size"]
            device = r["device"]
            compute_type = r["compute_type"]
            translate = self.translate_var.get()

            self._log(f"\n{'='*60}")
            self._log(f"Input:  {video_in}")
            self._log(f"Output: {video_out}")
            self._log(f"Model: {model_size}  Device: {device}  Compute: {compute_type}  Translate: {translate}")
            self._log(f"{'='*60}")
            self._set_status("Transcribing audio…")
            self._set_progress(2)

            # Check faster-whisper import
            try:
                import faster_whisper  # noqa
            except ImportError:
                self._log("❌ faster-whisper not installed. Run: pip install faster-whisper")
                raise RuntimeError("faster-whisper missing — see log.")

            segments, lang = transcribe_video(
                video_in, device, compute_type, model_size, translate,
                progress_cb=self._set_progress,
                log_cb=self._log,
            )

            if not segments:
                raise RuntimeError("No speech detected in the video.")

            self._log(f"\nDetected {len(segments)} subtitle segments.")
            self._set_status("Generating SRT…")
            self._set_progress(85)

            srt_content = segments_to_srt(segments)

            # Write SRT alongside output (or temp)
            srt_path = Path(video_out).with_suffix(".srt")
            srt_path.write_text(srt_content, encoding="utf-8")
            self._log(f"SRT saved: {srt_path}")

            self._set_status("Burning subtitles (ffmpeg)…")
            self._set_progress(88)

            burn_subtitles(
                video_in, srt_path, video_out,
                progress_cb=self._set_progress,
                log_cb=self._log,
            )

            # Optionally delete SRT
            if not self.save_srt_var.get():
                srt_path.unlink(missing_ok=True)

            self._set_progress(100)
            self._set_status(f"✅ Done! Saved to {video_out}")
            self._log(f"\n✅ All done → {video_out}")
            self.after(0, lambda: messagebox.showinfo("Done", f"Video saved:\n{video_out}"))

        except Exception as e:
            self._log(f"\n❌ Error: {e}")
            self._set_status(f"Error: {e}")
            self.after(0, lambda: messagebox.showerror("Processing failed", str(e)))

        finally:
            self.processing = False
            self.after(0, lambda: self.run_btn.config(state="normal", text="▶  Start Processing"))


# ──────────────────────────────────────────────
# Entry point
# ──────────────────────────────────────────────

if __name__ == "__main__":
    app = App()
    app.mainloop()