#!/usr/bin/env python3
"""Rename subtitle files to match their corresponding video files.

Designed as a GNOME/Nautilus right-click script, but also runnable from a
terminal. Subtitles are paired to videos by episode number (S01E02, 1x02, ...)
with a fallback to 1-to-1 matching when a folder holds a single video + single
subtitle. Trailing language codes (e.g. `.en`, `.eng`) are preserved and
existing files are never overwritten. A Zenity dialog lists the planned renames
for confirmation before anything happens.

Invocation:
  * From Nautilus: reads NAUTILUS_SCRIPT_SELECTED_FILE_PATHS.
  * From a terminal: `./rename-subtitles [PATH ...]` (files and/or folders).
    With no arguments it operates on the current directory.
"""

import os
import re
import shutil
import subprocess
import sys

VIDEO_EXTS = {
    ".mkv", ".mp4", ".avi", ".m4v", ".mov", ".wmv",
    ".flv", ".webm", ".mpg", ".mpeg", ".ts", ".m2ts",
}
SUBTITLE_EXTS = {
    ".srt", ".sub", ".ssa", ".ass", ".vtt", ".idx", ".sup",
}

# Language codes we recognise as a trailing token and preserve.
LANG_CODES = {
    # ISO 639-1 (2-letter) — common subset
    "en", "fr", "es", "de", "it", "nl", "pt", "sv", "da", "no", "fi",
    "pl", "ru", "cs", "hu", "ro", "tr", "el", "he", "ar", "ja", "ko",
    "zh", "hi", "th", "vi", "id", "uk", "bg", "hr", "sr", "sk", "sl",
    "et", "lv", "lt", "fa",
    # ISO 639-2 (3-letter) and common English names
    "eng", "english", "fre", "fra", "french", "spa", "esp", "spanish",
    "ger", "deu", "german", "ita", "italian", "dut", "nld", "dutch",
    "por", "portuguese", "swe", "swedish", "dan", "danish", "nor",
    "norwegian", "fin", "finnish", "pol", "polish", "rus", "russian",
    "cze", "ces", "czech", "hun", "hungarian", "rum", "ron", "romanian",
    "tur", "turkish", "gre", "ell", "greek", "heb", "hebrew", "ara",
    "arabic", "jpn", "japanese", "kor", "korean", "chi", "zho", "chinese",
    "hin", "hindi", "tha", "thai", "vie", "vietnamese", "ind", "indonesian",
    "ukr", "ukrainian",
}

# Episode-key regexes, tried in order. Each must yield (season, episode).
_EP_PATTERNS = [
    re.compile(r"[Ss](\d{1,2})[\._ -]?[Ee](\d{1,3})"),      # S01E02, s1e2, S01.E02
    re.compile(r"(?<!\d)(\d{1,2})[xX](\d{2,3})(?!\d)"),      # 1x02, 12x05
    re.compile(r"[Ss]eason[\._ -]?(\d{1,2})[\._ -]?"
               r"[Ee]pisode[\._ -]?(\d{1,3})"),               # Season 1 Episode 2
]
# Lighter fallback: a lone episode marker, treated as season 1.
_EP_ONLY = re.compile(r"(?:^|[\._ -])[Ee](\d{1,3})(?:[\._ -]|$)")


def episode_key(name):
    """Return (season, episode) parsed from *name*, or None."""
    stem = os.path.splitext(name)[0]
    for pat in _EP_PATTERNS:
        m = pat.search(stem)
        if m:
            return (int(m.group(1)), int(m.group(2)))
    m = _EP_ONLY.search(stem)
    if m:
        return (1, int(m.group(1)))
    return None


def split_lang(subname):
    """Split a subtitle filename into (stem_without_lang, lang_or_None, ext).

    The token immediately before the extension is treated as a language code
    only if it appears in LANG_CODES.
    """
    root, ext = os.path.splitext(subname)
    ext = ext.lower()
    stem, sep, last = root.rpartition(".")
    if sep and last.lower() in LANG_CODES:
        return stem, last, ext
    return root, None, ext


def gather_files(paths):
    """Resolve *paths* (files/dirs) into a flat, de-duplicated file list."""
    files = []
    seen = set()

    def add(p):
        ap = os.path.abspath(p)
        if ap not in seen and os.path.isfile(ap):
            seen.add(ap)
            files.append(ap)

    for p in paths:
        if os.path.isdir(p):
            for entry in sorted(os.listdir(p)):
                add(os.path.join(p, entry))
        elif os.path.isfile(p):
            add(p)
    return files


def classify(files):
    videos, subs = [], []
    for f in files:
        ext = os.path.splitext(f)[1].lower()
        if ext in VIDEO_EXTS:
            videos.append(f)
        elif ext in SUBTITLE_EXTS:
            subs.append(f)
    return videos, subs


def plan_renames(files):
    """Return (renames, skips).

    renames: list of (src_abs, dst_abs, old_name, new_name).
    skips:   list of (name, reason).
    """
    videos, subs = classify(files)
    renames, skips = [], []

    # Map episode key -> video (skip ambiguous duplicate keys).
    video_by_key = {}
    dup_keys = set()
    for v in videos:
        k = episode_key(os.path.basename(v))
        if k is None:
            continue
        if k in video_by_key:
            dup_keys.add(k)
        else:
            video_by_key[k] = v
    for k in dup_keys:
        video_by_key.pop(k, None)

    single_fallback = (len(videos) == 1 and len(subs) == 1
                       and episode_key(os.path.basename(subs[0])) is None
                       and episode_key(os.path.basename(videos[0])) is None)

    planned_dsts = set()
    for sub in subs:
        subname = os.path.basename(sub)
        subkey = episode_key(subname)

        if subkey is not None and subkey in video_by_key:
            video = video_by_key[subkey]
        elif subkey is not None and subkey in dup_keys:
            skips.append((subname, "multiple videos share this episode number"))
            continue
        elif single_fallback:
            video = videos[0]
        else:
            skips.append((subname, "no matching video found"))
            continue

        _, lang, sub_ext = split_lang(subname)
        video_stem = os.path.splitext(os.path.basename(video))[0]
        new_name = video_stem + (("." + lang) if lang else "") + sub_ext
        dst = os.path.join(os.path.dirname(sub), new_name)

        if os.path.abspath(dst) == os.path.abspath(sub):
            skips.append((subname, "already named correctly"))
            continue
        if os.path.exists(dst) or dst in planned_dsts:
            skips.append((subname, "target '%s' already exists" % new_name))
            continue

        planned_dsts.add(dst)
        renames.append((sub, dst, subname, new_name))

    return renames, skips


def build_message(renames, skips):
    lines = []
    if renames:
        lines.append("The following subtitles will be renamed:\n")
        for _, _, old, new in renames:
            lines.append("  %s\n    →  %s\n" % (old, new))
    if skips:
        lines.append("\nSkipped:")
        for name, reason in skips:
            lines.append("  %s  (%s)" % (name, reason))
    return "\n".join(lines)


def confirm(message, use_zenity):
    """Return True if the user approves the plan."""
    if use_zenity:
        rc = subprocess.run(
            ["zenity", "--question", "--no-wrap",
             "--title", "Rename subtitles to match videos",
             "--ok-label", "Rename", "--cancel-label", "Cancel",
             "--text", message],
        ).returncode
        return rc == 0
    print(message)
    try:
        return input("\nProceed with rename? [y/N] ").strip().lower() in ("y", "yes")
    except EOFError:
        return False


def notify(message, use_zenity, kind="info"):
    if use_zenity:
        subprocess.run(
            ["zenity", "--" + kind, "--no-wrap",
             "--title", "Rename subtitles to match videos",
             "--text", message],
        )
    else:
        print(message)


def main():
    env = os.environ.get("NAUTILUS_SCRIPT_SELECTED_FILE_PATHS", "")
    if env.strip():
        paths = [p for p in env.split("\n") if p]
    elif len(sys.argv) > 1:
        paths = sys.argv[1:]
    else:
        paths = [os.getcwd()]

    use_zenity = shutil.which("zenity") is not None and bool(os.environ.get("DISPLAY") or os.environ.get("WAYLAND_DISPLAY"))

    files = gather_files(paths)
    if not files:
        notify("No files found to process.", use_zenity, "warning")
        return 1

    renames, skips = plan_renames(files)

    if not renames:
        msg = "No subtitles needed renaming."
        if skips:
            msg += "\n\n" + build_message([], skips)
        notify(msg, use_zenity, "info")
        return 0

    if not confirm(build_message(renames, skips), use_zenity):
        return 0

    done, failed = 0, []
    for src, dst, old, new in renames:
        try:
            # Final guard against overwriting a file created in the meantime.
            if os.path.exists(dst):
                failed.append("%s (target appeared)" % old)
                continue
            os.rename(src, dst)
            done += 1
        except OSError as exc:
            failed.append("%s (%s)" % (old, exc))

    summary = "Renamed %d subtitle%s." % (done, "" if done == 1 else "s")
    if skips:
        summary += "\nSkipped %d." % len(skips)
    if failed:
        summary += "\nFailed %d:\n  %s" % (len(failed), "\n  ".join(failed))
    notify(summary, use_zenity, "error" if failed else "info")
    return 1 if failed else 0


if __name__ == "__main__":
    sys.exit(main())