"""
SubtitleBurner — Auto-transcribe, translate to English, burn subtitles into video.
Automatically scales subtitles to any video resolution.
"""

import os
import re
import shutil
import subprocess
import threading
import platform
from pathlib import Path
from datetime import timedelta

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

# ----------------------------------------------------------------------
# Helper functions
# ----------------------------------------------------------------------

def get_video_dimensions(video_path):
    try:
        cmd = ["ffprobe", "-v", "error", "-select_streams", "v:0",
               "-show_entries", "stream=width,height", "-of", "csv=p=0", str(video_path)]
        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
        w, h = map(int, result.stdout.strip().split(','))
        return w, h
    except:
        return 1920, 1080

def scale_value(value, video_height, ref_height=1080):
    return max(1, int(value * video_height / ref_height))

def hex_to_ass_color(hex_color):
    hex_color = hex_color.lstrip('#')
    r, g, b = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16)
    return f"&H00{b:02x}{g:02x}{r:02x}".upper()

def seconds_to_ass_time(seconds):
    hours = int(seconds // 3600)
    minutes = int((seconds % 3600) // 60)
    secs = seconds % 60
    cs = int((secs - int(secs)) * 100)
    return f"{hours}:{minutes:02d}:{int(secs):02d}.{cs:02d}"

# ----------------------------------------------------------------------
# Strict line wrapping (no line longer than max_chars, may break words)
# ----------------------------------------------------------------------

def wrap_text_by_chars(text, max_chars, max_lines):
    """
    Enforce that no line exceeds max_chars. Words may be split.
    Preserves spaces as they are.
    """
    text = text.strip()
    if not text:
        return text

    # Normalise multiple spaces to single, but keep single spaces
    text = ' '.join(text.split())

    if len(text) <= max_chars:
        return text

    lines = []
    remaining = text

    for line_idx in range(max_lines):
        if not remaining:
            break
        if len(remaining) <= max_chars:
            lines.append(remaining)
            break

        if line_idx == max_lines - 1:
            lines.append(remaining[:max_chars])
            break

        # Take first max_chars characters
        chunk = remaining[:max_chars]
        # Try to break at last space within chunk
        last_space = chunk.rfind(' ')
        if last_space > 0:
            lines.append(chunk[:last_space])
            remaining = remaining[last_space+1:]  # skip the space
        else:
            # No space: break exactly at max_chars
            lines.append(chunk)
            remaining = remaining[max_chars:]

    return '\n'.join(lines)

# ----------------------------------------------------------------------
# Build ASS subtitles with dynamic scaling
# ----------------------------------------------------------------------

def build_ass_subtitle(segments, video_width, video_height,
                       angle=135, max_chars=35, max_lines=2, y_offset_percent=0):
    """
    y_offset_percent: negative = move subtitles UP (higher on screen)
                      positive = move DOWN (lower on screen)
    """
    # Scale font size relative to height (42px at 1080p)
    base_font = scale_value(42, video_height, 1080)
    font_size = max(16, min(base_font, 48))
    shadow = scale_value(2, video_height, 1080)
    outline = 0  # no black outline
    margin_h = scale_value(10, video_height, 1080)

    # Default bottom margin (3% of video height)
    default_bottom = max(5, int(video_height * 0.03))
    offset_pixels = int(video_height * abs(y_offset_percent) / 100)
    if y_offset_percent < 0:
        margin_v = max(5, default_bottom - offset_pixels)  # move up
    else:
        margin_v = default_bottom + offset_pixels           # move down

    gold = hex_to_ass_color("#FEBD00")

    # Dynamically adjust max_chars based on video width (narrow = fewer chars)
    effective_max_chars = max(20, min(45, int(max_chars * video_width / 1920)))

    header = f"""[Script Info]
Title: SubtitleBurner
ScriptType: v4.00+
Collisions: Normal
PlayDepth: 0
Timer: 100.0000
VideoWidth: {video_width}
VideoHeight: {video_height}
WrapStyle: 0

[V4+ Styles]
Format: Name, Fontname, Fontsize, PrimaryColour, SecondaryColour, OutlineColour, BackColour, Bold, Italic, Underline, StrikeOut, ScaleX, ScaleY, Spacing, Angle, BorderStyle, Outline, Shadow, Alignment, MarginL, MarginR, MarginV, Encoding
Style: Default,Arial,{font_size},{gold},&H00FFFFFF,&H00000000,&H00000000,-1,0,0,0,100,100,0,0,1,{outline},{shadow},2,{margin_h},{margin_h},{margin_v},1

[Events]
Format: Layer, Start, End, Style, Name, MarginL, MarginR, MarginV, Effect, Text
"""

    lines = [header]
    for start, end, text in segments:
        wrapped = wrap_text_by_chars(text, effective_max_chars, max_lines)
        ass_text = wrapped.replace('\n', '\\N')
        styled = f"{{\\an2\\ang{angle}}}{ass_text}"
        lines.append(f"Dialogue: 0,{seconds_to_ass_time(start)},{seconds_to_ass_time(end)},Default,,0,0,0,,{styled}")

    return '\n'.join(lines)

# ----------------------------------------------------------------------
# Transcription (faster-whisper)
# ----------------------------------------------------------------------

def transcribe_video(video_path, device, compute_type, model_size, translate, progress_cb, log_cb):
    from faster_whisper import WhisperModel
    log_cb(f"Loading {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"Transcribing...")
    segments_iter, info = model.transcribe(video_path, task=task, beam_size=5, vad_filter=True)
    log_cb(f"Language: {info.language} ({info.language_probability:.0%})")
    segments = []
    total_dur = info.duration if hasattr(info, 'duration') else 0
    for seg in segments_iter:
        segments.append((seg.start, seg.end, seg.text.strip()))
        pct = min(95, int(seg.end / total_dur * 80) + 5) if total_dur else 50
        progress_cb(pct)
        log_cb(f"  [{seg.start:.1f}s → {seg.end:.1f}s] {seg.text.strip()}")
    return segments, info.language

# ----------------------------------------------------------------------
# Burn subtitles with ffmpeg
# ----------------------------------------------------------------------

def burn_ass_subtitles(video_path, ass_path, output_path, progress_cb, log_cb):
    log_cb("Burning subtitles...")
    escaped = str(ass_path).replace('\\', '/').replace(':', '\\:')
    cmd = ["ffmpeg", "-y", "-i", str(video_path), "-vf", f"ass='{escaped}'", "-c:a", "copy", "-preset", "fast", str(output_path)]
    process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, universal_newlines=True)
    dur = None
    for line in process.stdout:
        line = line.strip()
        if "Duration:" in line and dur is None:
            try:
                dstr = line.split("Duration:")[1].split(",")[0].strip()
                h,m,s = dstr.split(":")
                dur = int(h)*3600 + int(m)*60 + float(s)
            except: pass
        if "time=" in line and dur:
            try:
                tstr = line.split("time=")[1].split(" ")[0]
                h,m,s = tstr.split(":")
                cur = int(h)*3600 + int(m)*60 + float(s)
                pct = 95 + int(4 * cur / dur)
                progress_cb(min(99, pct))
            except: pass
    process.wait()
    if process.returncode != 0:
        raise RuntimeError("ffmpeg failed")
    log_cb("✓ Burning done")

# ----------------------------------------------------------------------
# Resource detection
# ----------------------------------------------------------------------

def detect_resources():
    import psutil, torch
    info = {
        "cpu_cores": os.cpu_count(),
        "ram_gb": round(psutil.virtual_memory().total / 1e9, 1),
        "gpu": None,
        "device": "cpu",
        "compute_type": "int8",
        "model_size": "medium"
    }
    if torch.cuda.is_available():
        info["gpu"] = torch.cuda.get_device_name(0)
        vram = torch.cuda.get_device_properties(0).total_memory / 1e9
        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 MPS"
    if info["device"] == "cpu":
        if info["ram_gb"] >= 16:
            info["model_size"] = "medium"
        elif info["ram_gb"] >= 8:
            info["model_size"] = "small"
        else:
            info["model_size"] = "base"
    return info

def check_ffmpeg():
    return shutil.which("ffmpeg") is not None

def check_ffprobe():
    return shutil.which("ffprobe") is not None

# ----------------------------------------------------------------------
# GUI
# ----------------------------------------------------------------------

class SubtitleBurnerApp(tk.Tk):
    def __init__(self):
        super().__init__()
        self.title("SubtitleBurner")
        self.geometry("880x720")
        self.minsize(800, 650)
        self.resources = None
        self.processing = False
        self._build_ui()
        self._detect_resources_async()

    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("TProgressbar", thickness=8)

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

        # Header
        ttk.Label(main, text="SubtitleBurner", font=("Segoe UI", 20, "bold")).pack(anchor="w")
        ttk.Label(main, text="Auto-transcribe · Translate · Burn with automatic scaling", foreground="#666").pack(anchor="w", pady=(0,12))

        # Resources
        res_frame = ttk.LabelFrame(main, text="System Resources", padding=8)
        res_frame.pack(fill="x", pady=(0,10))
        self.res_label = ttk.Label(res_frame, text="Detecting...")
        self.res_label.pack(anchor="w")

        # Input
        inp_frame = ttk.LabelFrame(main, 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).pack(side="left", fill="x", expand=True)
        ttk.Button(row, text="Browse...", command=self._browse_input).pack(side="left", padx=(6,0))

        # Output
        out_frame = ttk.LabelFrame(main, 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).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(main, text="Styling Options (auto-scaled to video)", padding=8)
        opt_frame.pack(fill="x", pady=(0,10))

        # Row 1
        r1 = ttk.Frame(opt_frame)
        r1.pack(fill="x", pady=2)
        self.translate_var = tk.BooleanVar(value=True)
        ttk.Checkbutton(r1, text="Translate to English", variable=self.translate_var).pack(side="left")
        ttk.Label(r1, text="   Format:").pack(side="left", padx=(10,0))
        self.format_var = tk.StringVar(value="ass")
        ttk.Combobox(r1, textvariable=self.format_var, values=["ass","srt"], width=6, state="readonly").pack(side="left", padx=4)
        ttk.Label(r1, text="   Angle (°):").pack(side="left", padx=(10,0))
        self.angle_var = tk.StringVar(value="135")
        ttk.Entry(r1, textvariable=self.angle_var, width=6).pack(side="left", padx=4)
        ttk.Label(r1, text="   Font adjust (px):").pack(side="left", padx=(10,0))
        self.font_adj_var = tk.StringVar(value="0")
        ttk.Combobox(r1, textvariable=self.font_adj_var, values=["-8","-4","0","+4","+8"], width=5, state="readonly").pack(side="left", padx=4)
        ttk.Label(r1, text="   Model:").pack(side="left", padx=(10,0))
        self.model_var = tk.StringVar(value="auto")
        ttk.Combobox(r1, textvariable=self.model_var, values=["auto","tiny","base","small","medium","large-v3"], width=10, state="readonly").pack(side="left", padx=4)

        # Row 2 - Line wrapping
        r2 = ttk.Frame(opt_frame)
        r2.pack(fill="x", pady=2)
        ttk.Label(r2, text="Line wrapping:", font=("Segoe UI",10,"bold")).pack(side="left")
        ttk.Label(r2, text="   Max chars/line:").pack(side="left", padx=(10,0))
        self.chars_var = tk.IntVar(value=35)
        ttk.Spinbox(r2, from_=15, to=60, width=5, textvariable=self.chars_var).pack(side="left", padx=4)
        ttk.Label(r2, text="   Max lines:").pack(side="left", padx=(10,0))
        self.lines_var = tk.IntVar(value=3)   # allow 3 lines by default to avoid truncation
        ttk.Combobox(r2, textvariable=self.lines_var, values=[1,2,3,4,5], width=5, state="readonly").pack(side="left", padx=4)
        ttk.Button(r2, text="Preview wrapping", command=self._preview_wrapping, width=14).pack(side="left", padx=(15,0))
        self.save_sub_var = tk.BooleanVar(value=True)
        ttk.Checkbutton(r2, text="Keep subtitle file", variable=self.save_sub_var).pack(side="left", padx=(15,0))

        # Row 3 - Vertical position (percentage)
        r3 = ttk.Frame(opt_frame)
        r3.pack(fill="x", pady=2)
        ttk.Label(r3, text="Vertical position:", font=("Segoe UI",10,"bold")).pack(side="left")
        ttk.Label(r3, text="   Offset (% of height):").pack(side="left", padx=(10,0))
        self.yoffset_var = tk.IntVar(value=0)
        y_slider = ttk.Scale(r3, from_=-30, to=20, orient="horizontal", variable=self.yoffset_var, length=200)
        y_slider.pack(side="left", padx=6)
        self.yoffset_label = ttk.Label(r3, text="0%", width=6)
        self.yoffset_label.pack(side="left")
        ttk.Button(r3, text="Reset", command=lambda: self.yoffset_var.set(0), width=5).pack(side="left", padx=6)
        ttk.Label(r3, text="   ( – = higher on screen,  + = lower )", foreground="#666").pack(side="left", padx=(10,0))
        self.yoffset_var.trace_add("write", lambda *a: self.yoffset_label.config(text=f"{self.yoffset_var.get()}%"))

        # Progress
        prog_frame = ttk.LabelFrame(main, text="Progress", padding=8)
        prog_frame.pack(fill="x", pady=(0,10))
        self.progress_var = tk.IntVar()
        ttk.Progressbar(prog_frame, variable=self.progress_var, maximum=100).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(main, text="Log", padding=8)
        log_frame.pack(fill="both", expand=True)
        self.log_box = scrolledtext.ScrolledText(log_frame, height=8, font=("Courier New",9), bg="#F4F3EF", wrap="word")
        self.log_box.pack(fill="both", expand=True)

        # Buttons
        btn_frame = ttk.Frame(main)
        btn_frame.pack(fill="x", pady=(10,0))
        self.run_btn = ttk.Button(btn_frame, text="▶ Start Processing", command=self._start, style="Accent.TButton")
        self.run_btn.pack(side="left")
        ttk.Button(btn_frame, text="Clear log", command=self._clear_log).pack(side="left", padx=8)
        ttk.Button(btn_frame, text="Video info", command=self._show_video_info).pack(side="left", padx=8)

        if not check_ffmpeg():
            self._log("⚠ ffmpeg not found – install it and add to PATH")
        if not check_ffprobe():
            self._log("⚠ ffprobe not found – video info will be limited")

    # ------------------------------------------------------------------
    # GUI callbacks
    # ------------------------------------------------------------------
    def _log(self, msg):
        self.after(0, lambda: self.log_box.insert("end", msg + "\n") or self.log_box.see("end"))

    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, val):
        self.after(0, lambda: self.progress_var.set(val))

    def _detect_resources_async(self):
        def run():
            r = detect_resources()
            self.resources = r
            text = f"CPU: {r['cpu_cores']} cores, RAM: {r['ram_gb']} GB"
            if r.get("gpu"):
                text += f", GPU: {r['gpu']}"
            text += f" → using {r['device'].upper()} / {r['model_size']}"
            self.after(0, lambda: self.res_label.config(text=text))
        threading.Thread(target=run, daemon=True).start()

    def _browse_input(self):
        path = filedialog.askopenfilename(filetypes=[("Video", "*.mp4 *.mkv *.mov *.avi *.webm")])
        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.mp4")))
            self._log(f"Selected: {path}")
            try:
                w, h = get_video_dimensions(path)
                self._log(f"Video size: {w}x{h}")
                suggested_chars = max(20, min(45, int(35 * w / 1920)))
                self.chars_var.set(suggested_chars)
                self._log(f"Auto-set max chars/line to {suggested_chars} based on width.")
            except:
                pass

    def _browse_output(self):
        path = filedialog.asksaveasfilename(defaultextension=".mp4", filetypes=[("MP4", "*.mp4")])
        if path:
            self.output_var.set(path)

    def _preview_wrapping(self):
        sample = "This is a long sample subtitle to demonstrate how the automatic line wrapping works with your current settings. It may break words if necessary."
        maxc = self.chars_var.get()
        maxl = self.lines_var.get()
        wrapped = wrap_text_by_chars(sample, maxc, maxl)
        messagebox.showinfo("Preview", f"Max chars/line: {maxc}\nMax lines: {maxl}\n\nResult:\n{wrapped}")

    def _show_video_info(self):
        path = self.input_var.get().strip()
        if not path or not Path(path).exists():
            messagebox.showwarning("No video", "Select a video file first.")
            return
        w, h = get_video_dimensions(path)
        font_size = max(16, min(48, int(42 * h / 1080)))
        rec_chars = max(20, min(45, int(35 * w / 1920)))
        msg = f"""Video: {w} x {h}
Aspect ratio: {w/h:.2f}
Recommended settings:
- Font size: {font_size} px
- Max chars/line: {rec_chars}
- Max lines: 3
- Vertical offset: 0% (negative moves up, positive down)"""
        messagebox.showinfo("Video Information", msg)

    # ------------------------------------------------------------------
    # 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 or not Path(video_in).exists():
            messagebox.showerror("Error", "Valid input video required")
            return
        if not video_out:
            messagebox.showerror("Error", "Output path required")
            return
        if not check_ffmpeg():
            messagebox.showerror("Error", "ffmpeg not found")
            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 = self.model_var.get()
            if model == "auto":
                model = r["model_size"]
            translate = self.translate_var.get()
            fmt = self.format_var.get()
            angle = int(self.angle_var.get())
            font_adj = int(self.font_adj_var.get())
            max_chars = self.chars_var.get()
            max_lines = self.lines_var.get()
            yoffset_pct = self.yoffset_var.get()

            w, h = get_video_dimensions(video_in)
            base_font = max(16, min(48, int(42 * h / 1080)))
            final_font = base_font + font_adj

            self._log(f"\n{'='*60}")
            self._log(f"Video: {w}x{h}")
            self._log(f"Model: {model}, Device: {r['device']}, Translate: {translate}")
            self._log(f"Format: {fmt.upper()}, Angle: {angle}°, Font: {final_font}px (auto-scaled)")
            self._log(f"Line wrap: {max_chars} chars/line, max {max_lines} lines")
            self._log(f"Vertical offset: {yoffset_pct}% (negative = up, positive = down)")
            self._log(f"{'='*60}")

            self._set_status("Transcribing audio...")
            self._set_progress(2)
            segments, lang = transcribe_video(video_in, r["device"], r["compute_type"], model, translate,
                                              self._set_progress, self._log)

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

            self._log(f"Got {len(segments)} segments")
            self._set_status("Generating subtitles...")
            self._set_progress(85)

            if fmt == "ass":
                ass_content = build_ass_subtitle(segments, w, h, angle, max_chars, max_lines, yoffset_pct)
                sub_path = Path(video_out).with_suffix(".ass")
                sub_path.write_text(ass_content, encoding="utf-8")
                self._log(f"ASS file: {sub_path}")
                self._set_status("Burning subtitles...")
                self._set_progress(88)
                burn_ass_subtitles(video_in, sub_path, video_out, self._set_progress, self._log)
            else:
                # SRT fallback
                srt_lines = []
                for i, (start, end, text) in enumerate(segments, 1):
                    wrapped = wrap_text_by_chars(text, max_chars, max_lines)
                    srt_lines.append(str(i))
                    start_str = str(timedelta(seconds=start)).replace('.', ',')[:11]
                    end_str = str(timedelta(seconds=end)).replace('.', ',')[:11]
                    srt_lines.append(f"{start_str} --> {end_str}")
                    srt_lines.append(wrapped)
                    srt_lines.append("")
                sub_path = Path(video_out).with_suffix(".srt")
                sub_path.write_text("\n".join(srt_lines), encoding="utf-8")
                self._log(f"SRT file: {sub_path}")
                # Burn SRT
                escaped = str(sub_path).replace('\\', '/')
                if platform.system() == "Windows":
                    escaped = escaped.replace(':', '\\:')
                style = f"FontName=Arial,FontSize={final_font},PrimaryColour=&H0000BDFE,OutlineColour=&H00000000,Bold=-1,Shadow=1,Alignment=2"
                cmd = ["ffmpeg", "-y", "-i", video_in, "-vf", f"subtitles='{escaped}':force_style='{style}'", "-c:a", "copy", "-preset", "fast", video_out]
                subprocess.run(cmd, capture_output=True, check=True)

            if not self.save_sub_var.get():
                sub_path.unlink(missing_ok=True)

            self._set_progress(100)
            self._set_status(f"✅ Done: {video_out}")
            self._log(f"\n✅ Finished → {video_out}")
            self.after(0, lambda: messagebox.showinfo("Complete", 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("Error", str(e)))
        finally:
            self.processing = False
            self.after(0, lambda: self.run_btn.config(state="normal", text="▶ Start Processing"))

# ----------------------------------------------------------------------
# Run
# ----------------------------------------------------------------------
if __name__ == "__main__":
    app = SubtitleBurnerApp()
    app.mainloop()