#!/usr/bin/env python3
"""
WeBrake — HandBrake CLI web frontend for Alpine LXC containers.

Serves a single-page UI, accepts media uploads, queues HandBrakeCLI jobs
with full flag coverage (including a raw-arguments passthrough), reports
live progress, and serves finished files back for download.

Secrets live in token.json beside this file. That file is generated at
install time (or on first run if missing) and must NEVER be committed
to the repository.
"""

import json
import os
import re
import secrets
import shlex
import shutil
import signal
import subprocess
import threading
import time
import uuid
from functools import wraps
from pathlib import Path

from flask import (Flask, Response, abort, jsonify, request,
                   send_file, send_from_directory)

# --------------------------------------------------------------------------
# Paths & configuration
# --------------------------------------------------------------------------
APP_DIR = Path(__file__).resolve().parent
DATA_DIR = Path(os.environ.get("WEBRAKE_DATA", "/var/lib/webrake"))
UPLOAD_DIR = DATA_DIR / "uploads"
OUTPUT_DIR = DATA_DIR / "output"
LOG_DIR = DATA_DIR / "logs"
STATE_FILE = DATA_DIR / "jobs.json"
TOKEN_FILE = APP_DIR / "token.json"

HANDBRAKE = shutil.which("HandBrakeCLI") or "/usr/bin/HandBrakeCLI"
HOST = os.environ.get("WEBRAKE_HOST", "0.0.0.0")
PORT = int(os.environ.get("WEBRAKE_PORT", "8090"))

for d in (UPLOAD_DIR, OUTPUT_DIR, LOG_DIR):
    d.mkdir(parents=True, exist_ok=True)

# --------------------------------------------------------------------------
# token.json — generated at install, never stored in the repo
# --------------------------------------------------------------------------
def load_tokens() -> dict:
    if not TOKEN_FILE.exists():
        tokens = {
            "api_token": secrets.token_urlsafe(32),
            "secret_key": secrets.token_urlsafe(32),
            "note": "Generated locally by WeBrake. Do not commit this file.",
        }
        TOKEN_FILE.write_text(json.dumps(tokens, indent=2))
        os.chmod(TOKEN_FILE, 0o600)
        return tokens
    return json.loads(TOKEN_FILE.read_text())


TOKENS = load_tokens()

app = Flask(__name__)
app.secret_key = TOKENS["secret_key"]
app.config["MAX_CONTENT_LENGTH"] = 1024 * 1024 * 1024 * 64  # 64 GiB


def require_token(fn):
    @wraps(fn)
    def wrapper(*args, **kwargs):
        supplied = (request.headers.get("X-API-Token")
                    or request.args.get("token", ""))
        if not secrets.compare_digest(supplied, TOKENS["api_token"]):
            return jsonify({"error": "Invalid or missing API token."}), 401
        return fn(*args, **kwargs)
    return wrapper


# --------------------------------------------------------------------------
# Job model & persistence
# --------------------------------------------------------------------------
JOBS: dict = {}
JOB_LOCK = threading.Lock()
QUEUE_EVENT = threading.Event()
PROCS: dict = {}  # job_id -> Popen


def save_state():
    with JOB_LOCK:
        snapshot = {jid: {k: v for k, v in j.items()} for jid, j in JOBS.items()}
    STATE_FILE.write_text(json.dumps(snapshot, indent=2))


def load_state():
    if STATE_FILE.exists():
        try:
            data = json.loads(STATE_FILE.read_text())
            for jid, job in data.items():
                if job.get("status") in ("queued", "running", "scanning"):
                    job["status"] = "failed"
                    job["message"] = "Interrupted by server restart."
                JOBS[jid] = job
        except Exception:
            pass


# --------------------------------------------------------------------------
# GPU / capability detection
# --------------------------------------------------------------------------
def detect_gpu() -> dict:
    info = {"dri_devices": [], "nvidia": False, "vaapi": False, "notes": []}
    dri = Path("/dev/dri")
    if dri.exists():
        info["dri_devices"] = sorted(p.name for p in dri.iterdir())
        info["vaapi"] = any(n.startswith("renderD") for n in info["dri_devices"])
    if Path("/dev/nvidia0").exists() or Path("/dev/nvidiactl").exists():
        info["nvidia"] = True
    if not info["dri_devices"] and not info["nvidia"]:
        info["notes"].append(
            "No GPU devices visible in this container. Software encoders will be used. "
            "Pass /dev/dri (Intel/AMD) or /dev/nvidia* (NVIDIA) into the LXC to enable "
            "hardware encoding.")
    return info


def hb_help_text() -> str:
    try:
        out = subprocess.run([HANDBRAKE, "--help"], capture_output=True,
                             text=True, timeout=30)
        return out.stdout + out.stderr
    except Exception:
        return ""


def detect_encoders() -> list:
    """Parse the encoder list out of HandBrakeCLI --help."""
    text = hb_help_text()
    encoders = []
    m = re.search(r"--encoder\b.*?Select video encoder:(.*?)(?:--|\Z)",
                  text, re.S)
    block = m.group(1) if m else text
    for token in re.findall(r"^\s{6,}([a-z0-9_]+)\s*$", block, re.M):
        encoders.append(token)
    if not encoders:
        # Sensible fallback list; UI marks unavailable ones after a scan fails.
        encoders = ["svt_av1", "x264", "x264_10bit", "x265", "x265_10bit",
                    "x265_12bit", "mpeg4", "mpeg2", "VP8", "VP9", "theora",
                    "nvenc_h264", "nvenc_h265", "nvenc_av1",
                    "qsv_h264", "qsv_h265", "qsv_av1",
                    "vce_h264", "vce_h265"]
    return encoders


def detect_presets() -> list:
    try:
        out = subprocess.run([HANDBRAKE, "--preset-list"], capture_output=True,
                             text=True, timeout=60)
        text = out.stdout + out.stderr
        presets, category = [], ""
        for line in text.splitlines():
            cat = re.match(r"^([A-Za-z].*)/$", line.strip())
            if cat:
                category = cat.group(1)
                continue
            item = re.match(r"^\s{4}(\S.*)$", line)
            if item and category and not line.strip().startswith("+"):
                name = item.group(1).strip()
                if name and not name.startswith("-"):
                    presets.append({"category": category, "name": name})
        return presets
    except Exception:
        return []


CAPS_CACHE: dict = {}


def capabilities() -> dict:
    if not CAPS_CACHE:
        version = ""
        try:
            v = subprocess.run([HANDBRAKE, "--version"], capture_output=True,
                               text=True, timeout=30)
            version = (v.stdout + v.stderr).strip().splitlines()[0] if (v.stdout or v.stderr) else ""
        except Exception:
            version = "HandBrakeCLI not found — install it inside the container."
        CAPS_CACHE.update({
            "handbrake": version,
            "binary": HANDBRAKE,
            "encoders": detect_encoders(),
            "presets": detect_presets(),
            "gpu": detect_gpu(),
        })
    return CAPS_CACHE


# --------------------------------------------------------------------------
# HandBrake command builder — maps UI options onto CLI flags
# --------------------------------------------------------------------------
SAFE_NAME = re.compile(r"[^A-Za-z0-9._ ()\[\]-]")


def sanitize_name(name: str) -> str:
    return SAFE_NAME.sub("_", Path(name).name)[:200] or "media"


def add(cmd: list, flag: str, value=None):
    if value is None:
        cmd.append(flag)
    else:
        cmd.extend([flag, str(value)])


def build_command(job: dict) -> list:
    o = job["options"]
    src = UPLOAD_DIR / job["source"]
    dst = OUTPUT_DIR / job["output"]
    cmd = [HANDBRAKE, "-i", str(src), "-o", str(dst)]

    # ---- General / source -------------------------------------------------
    if o.get("preset"):
        add(cmd, "--preset", o["preset"])
    if o.get("format"):
        add(cmd, "--format", o["format"])
    if o.get("title"):
        add(cmd, "--title", o["title"])
    if o.get("chapters"):
        add(cmd, "--chapters", o["chapters"])
    if o.get("start_at"):
        add(cmd, "--start-at", o["start_at"])
    if o.get("stop_at"):
        add(cmd, "--stop-at", o["stop_at"])
    if o.get("angle"):
        add(cmd, "--angle", o["angle"])
    if o.get("markers"):
        add(cmd, "--markers")
    if o.get("optimize"):
        add(cmd, "--optimize")
    if o.get("ipod_atom"):
        add(cmd, "--ipod-atom")
    if o.get("align_av"):
        add(cmd, "--align-av")
    if o.get("inline_parameter_sets"):
        add(cmd, "--inline-parameter-sets")

    # ---- Video ------------------------------------------------------------
    if o.get("encoder"):
        add(cmd, "--encoder", o["encoder"])
    rc = o.get("rate_control", "quality")
    if rc == "quality" and o.get("quality") not in (None, ""):
        add(cmd, "--quality", o["quality"])
    elif rc == "bitrate" and o.get("vb"):
        add(cmd, "--vb", o["vb"])
        if o.get("two_pass"):
            add(cmd, "--multi-pass")
            if o.get("turbo"):
                add(cmd, "--turbo")
    if o.get("encoder_preset"):
        add(cmd, "--encoder-preset", o["encoder_preset"])
    if o.get("encoder_tune"):
        add(cmd, "--encoder-tune", o["encoder_tune"])
    if o.get("encoder_profile"):
        add(cmd, "--encoder-profile", o["encoder_profile"])
    if o.get("encoder_level"):
        add(cmd, "--encoder-level", o["encoder_level"])
    if o.get("encopts"):
        add(cmd, "--encopts", o["encopts"])
    if o.get("framerate"):
        add(cmd, "--rate", o["framerate"])
    fr_mode = o.get("framerate_mode")
    if fr_mode == "cfr":
        add(cmd, "--cfr")
    elif fr_mode == "vfr":
        add(cmd, "--vfr")
    elif fr_mode == "pfr":
        add(cmd, "--pfr")

    # ---- Dimensions ---------------------------------------------------------
    if o.get("width"):
        add(cmd, "--width", o["width"])
    if o.get("height"):
        add(cmd, "--height", o["height"])
    if o.get("max_width"):
        add(cmd, "--maxWidth", o["max_width"])
    if o.get("max_height"):
        add(cmd, "--maxHeight", o["max_height"])
    if o.get("crop"):
        add(cmd, "--crop", o["crop"])
    if o.get("crop_mode"):
        add(cmd, "--crop-mode", o["crop_mode"])
    anam = o.get("anamorphic")
    if anam in ("auto", "loose", "custom", "non"):
        add(cmd, f"--{'non-' if anam == 'non' else ''}anamorphic"
            if anam == "non" else f"--{anam}-anamorphic")
    if o.get("display_width"):
        add(cmd, "--display-width", o["display_width"])
    if o.get("pixel_aspect"):
        add(cmd, "--pixel-aspect", o["pixel_aspect"])
    if o.get("modulus"):
        add(cmd, "--modulus", o["modulus"])
    if o.get("color_matrix"):
        add(cmd, "--color-matrix", o["color_matrix"])

    # ---- Filters ------------------------------------------------------------
    def filt(key, flag):
        v = o.get(key)
        if v is True or v == "default":
            add(cmd, flag)
        elif v:
            add(cmd, flag, v)

    filt("comb_detect", "--comb-detect")
    filt("deinterlace", "--deinterlace")
    filt("decomb", "--decomb")
    filt("detelecine", "--detelecine")
    if o.get("denoise_filter") == "hqdn3d":
        filt("denoise", "--hqdn3d")
    elif o.get("denoise_filter") == "nlmeans":
        filt("denoise", "--nlmeans")
        if o.get("nlmeans_tune"):
            add(cmd, "--nlmeans-tune", o["nlmeans_tune"])
    filt("chroma_smooth", "--chroma-smooth")
    if o.get("chroma_smooth_tune"):
        add(cmd, "--chroma-smooth-tune", o["chroma_smooth_tune"])
    if o.get("sharpen_filter") == "unsharp":
        filt("sharpen", "--unsharp")
        if o.get("sharpen_tune"):
            add(cmd, "--unsharp-tune", o["sharpen_tune"])
    elif o.get("sharpen_filter") == "lapsharp":
        filt("sharpen", "--lapsharp")
        if o.get("sharpen_tune"):
            add(cmd, "--lapsharp-tune", o["sharpen_tune"])
    filt("deblock", "--deblock")
    if o.get("deblock_tune"):
        add(cmd, "--deblock-tune", o["deblock_tune"])
    if o.get("rotate"):
        add(cmd, "--rotate", o["rotate"])
    if o.get("pad"):
        add(cmd, "--pad", o["pad"])
    if o.get("colorspace"):
        add(cmd, "--colorspace", o["colorspace"])
    if o.get("grayscale"):
        add(cmd, "--grayscale")
    if o.get("no_dvdnav"):
        add(cmd, "--no-dvdnav")

    # ---- Audio --------------------------------------------------------------
    if o.get("all_audio"):
        add(cmd, "--all-audio")
    elif o.get("audio_tracks"):
        add(cmd, "--audio", o["audio_tracks"])
    if o.get("audio_encoder"):
        add(cmd, "--aencoder", o["audio_encoder"])
    if o.get("audio_bitrate"):
        add(cmd, "--ab", o["audio_bitrate"])
    if o.get("audio_quality"):
        add(cmd, "--aq", o["audio_quality"])
    if o.get("mixdown"):
        add(cmd, "--mixdown", o["mixdown"])
    if o.get("samplerate"):
        add(cmd, "--arate", o["samplerate"])
    if o.get("drc"):
        add(cmd, "--drc", o["drc"])
    if o.get("gain"):
        add(cmd, "--gain", o["gain"])
    if o.get("audio_names"):
        add(cmd, "--aname", o["audio_names"])
    if o.get("audio_copy_mask"):
        add(cmd, "--audio-copy-mask", o["audio_copy_mask"])
    if o.get("audio_fallback"):
        add(cmd, "--audio-fallback", o["audio_fallback"])
    if o.get("normalize_mix"):
        add(cmd, "--normalize-mix", o["normalize_mix"])

    # ---- Subtitles ------------------------------------------------------------
    if o.get("all_subtitles"):
        add(cmd, "--all-subtitles")
    elif o.get("subtitle_tracks"):
        add(cmd, "--subtitle", o["subtitle_tracks"])
    if o.get("subtitle_burned"):
        add(cmd, "--subtitle-burned", o["subtitle_burned"])
    if o.get("subtitle_default"):
        add(cmd, "--subtitle-default", o["subtitle_default"])
    if o.get("subtitle_forced"):
        add(cmd, "--subtitle-forced", o["subtitle_forced"])
    if o.get("native_language"):
        add(cmd, "--native-language", o["native_language"])
    if o.get("srt_file"):
        add(cmd, "--srt-file", str(UPLOAD_DIR / sanitize_name(o["srt_file"])))
        if o.get("srt_codeset"):
            add(cmd, "--srt-codeset", o["srt_codeset"])
        if o.get("srt_lang"):
            add(cmd, "--srt-lang", o["srt_lang"])
        if o.get("srt_burn"):
            add(cmd, "--srt-burn")

    # ---- Raw passthrough: guarantees EVERY HandBrake flag is reachable -------
    if o.get("raw_args"):
        cmd.extend(shlex.split(o["raw_args"]))

    return cmd


# --------------------------------------------------------------------------
# Worker: runs queued jobs one at a time, parses live progress
# --------------------------------------------------------------------------
PROGRESS_RE = re.compile(
    r"Encoding:.*?(\d+\.\d+)\s?%"
    r"(?:.*?(\d+\.\d+)\s?fps"
    r".*?avg\s+(\d+\.\d+)\s?fps"
    r".*?ETA\s+(\d+h\d+m\d+s))?", re.S)
TASK_RE = re.compile(r"task (\d+) of (\d+)")


def run_job(job_id: str):
    with JOB_LOCK:
        job = JOBS[job_id]
        job["status"] = "running"
        job["started"] = time.time()
    save_state()

    cmd = job["command"]
    log_path = LOG_DIR / f"{job_id}.log"
    try:
        with open(log_path, "w") as log:
            log.write("$ " + " ".join(shlex.quote(c) for c in cmd) + "\n\n")
            log.flush()
            proc = subprocess.Popen(cmd, stdout=subprocess.PIPE,
                                    stderr=subprocess.STDOUT, text=True,
                                    bufsize=1, errors="replace",
                                    preexec_fn=os.setsid)
            PROCS[job_id] = proc
            buf = ""
            while True:
                chunk = proc.stdout.read(256)
                if not chunk:
                    break
                log.write(chunk)
                log.flush()
                buf = (buf + chunk)[-2000:]
                m = None
                for m in PROGRESS_RE.finditer(buf):
                    pass
                if m:
                    with JOB_LOCK:
                        job["progress"] = float(m.group(1))
                        if m.group(2):
                            job["fps"] = float(m.group(2))
                            job["avg_fps"] = float(m.group(3))
                            job["eta"] = m.group(4)
                    t = None
                    for t in TASK_RE.finditer(buf):
                        pass
                    if t:
                        with JOB_LOCK:
                            job["task"] = f"{t.group(1)}/{t.group(2)}"
            proc.wait()
        PROCS.pop(job_id, None)
        out_file = OUTPUT_DIR / job["output"]
        with JOB_LOCK:
            if job.get("status") == "cancelled":
                pass
            elif proc.returncode == 0 and out_file.exists():
                job["status"] = "done"
                job["progress"] = 100.0
                job["size"] = out_file.stat().st_size
                job["message"] = "Ready to download."
            else:
                job["status"] = "failed"
                job["message"] = f"HandBrakeCLI exited with code {proc.returncode}. See log."
            job["finished"] = time.time()
    except Exception as exc:
        PROCS.pop(job_id, None)
        with JOB_LOCK:
            job["status"] = "failed"
            job["message"] = f"{type(exc).__name__}: {exc}"
            job["finished"] = time.time()
    save_state()


def worker_loop():
    while True:
        QUEUE_EVENT.wait(timeout=2)
        QUEUE_EVENT.clear()
        while True:
            with JOB_LOCK:
                pending = [j for j in JOBS.values() if j["status"] == "queued"]
                pending.sort(key=lambda j: j["created"])
                nxt = pending[0]["id"] if pending else None
            if not nxt:
                break
            run_job(nxt)


threading.Thread(target=worker_loop, daemon=True).start()


# --------------------------------------------------------------------------
# Routes
# --------------------------------------------------------------------------
@app.route("/")
def index():
    return send_from_directory(APP_DIR, "index.html")


@app.route("/api/ping")
def ping():
    supplied = request.headers.get("X-API-Token", "")
    ok = secrets.compare_digest(supplied, TOKENS["api_token"])
    return jsonify({"ok": ok})


@app.route("/api/capabilities")
@require_token
def api_capabilities():
    return jsonify(capabilities())


@app.route("/api/upload", methods=["POST"])
@require_token
def api_upload():
    f = request.files.get("file")
    if not f or not f.filename:
        return jsonify({"error": "No file supplied."}), 400
    name = sanitize_name(f.filename)
    dest = UPLOAD_DIR / name
    stem, suffix, n = dest.stem, dest.suffix, 1
    while dest.exists():
        dest = UPLOAD_DIR / f"{stem}({n}){suffix}"
        n += 1
    f.save(dest)
    return jsonify({"filename": dest.name, "size": dest.stat().st_size})


@app.route("/api/sources")
@require_token
def api_sources():
    files = []
    for p in sorted(UPLOAD_DIR.iterdir()):
        if p.is_file():
            files.append({"name": p.name, "size": p.stat().st_size,
                          "mtime": p.stat().st_mtime})
    return jsonify(files)


@app.route("/api/sources/<path:name>", methods=["DELETE"])
@require_token
def api_delete_source(name):
    p = UPLOAD_DIR / sanitize_name(name)
    if p.exists():
        p.unlink()
    return jsonify({"ok": True})


@app.route("/api/scan", methods=["POST"])
@require_token
def api_scan():
    """Scan a source with HandBrakeCLI to enumerate titles/tracks."""
    name = sanitize_name(request.json.get("filename", ""))
    src = UPLOAD_DIR / name
    if not src.exists():
        return jsonify({"error": "Source not found."}), 404
    try:
        out = subprocess.run(
            [HANDBRAKE, "-i", str(src), "--scan", "--title", "0", "--json"],
            capture_output=True, text=True, timeout=600)
        text = out.stdout
        m = re.search(r"JSON Title Set:\s*(\{.*)", text, re.S)
        titles = []
        if m:
            data = json.loads(m.group(1)[:m.group(1).rfind("}") + 1])
            for t in data.get("TitleList", []):
                titles.append({
                    "index": t.get("Index"),
                    "duration": t.get("Duration"),
                    "geometry": t.get("Geometry"),
                    "framerate": t.get("FrameRate"),
                    "audio": [{"track": i + 1,
                               "description": a.get("Description", ""),
                               "language": a.get("Language", "")}
                              for i, a in enumerate(t.get("AudioList", []))],
                    "subtitles": [{"track": i + 1,
                                   "name": s.get("Name") or s.get("Language", ""),
                                   "format": s.get("SourceName", "")}
                                  for i, s in enumerate(t.get("SubtitleList", []))],
                })
        return jsonify({"titles": titles})
    except subprocess.TimeoutExpired:
        return jsonify({"error": "Scan timed out."}), 500
    except Exception as exc:
        return jsonify({"error": str(exc)}), 500


@app.route("/api/jobs", methods=["GET"])
@require_token
def api_jobs():
    with JOB_LOCK:
        jobs = sorted(JOBS.values(), key=lambda j: j["created"], reverse=True)
    return jsonify(jobs)


@app.route("/api/jobs", methods=["POST"])
@require_token
def api_create_job():
    body = request.json or {}
    source = sanitize_name(body.get("source", ""))
    if not (UPLOAD_DIR / source).exists():
        return jsonify({"error": "Source file not found — upload it first."}), 400
    options = body.get("options", {})
    fmt = options.get("format", "av_mkv")
    ext = {"av_mp4": ".mp4", "av_mkv": ".mkv", "av_webm": ".webm"}.get(fmt, ".mkv")
    out_name = body.get("output") or (Path(source).stem + ".webrake" + ext)
    out_name = sanitize_name(out_name)

    job_id = uuid.uuid4().hex[:12]
    job = {
        "id": job_id,
        "source": source,
        "output": out_name,
        "options": options,
        "status": "queued",
        "progress": 0.0,
        "created": time.time(),
        "message": "Waiting in queue.",
    }
    job["command"] = build_command(job)
    job["command_preview"] = " ".join(shlex.quote(c) for c in job["command"])
    with JOB_LOCK:
        JOBS[job_id] = job
    save_state()
    QUEUE_EVENT.set()
    return jsonify(job), 201


@app.route("/api/jobs/<job_id>/cancel", methods=["POST"])
@require_token
def api_cancel(job_id):
    with JOB_LOCK:
        job = JOBS.get(job_id)
        if not job:
            abort(404)
        job["status"] = "cancelled"
        job["message"] = "Cancelled by user."
    proc = PROCS.get(job_id)
    if proc:
        try:
            os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
        except Exception:
            pass
    save_state()
    return jsonify({"ok": True})


@app.route("/api/jobs/<job_id>", methods=["DELETE"])
@require_token
def api_delete_job(job_id):
    with JOB_LOCK:
        job = JOBS.pop(job_id, None)
    if job:
        for p in (OUTPUT_DIR / job["output"], LOG_DIR / f"{job_id}.log"):
            if p.exists():
                p.unlink()
    save_state()
    return jsonify({"ok": True})


@app.route("/api/jobs/<job_id>/log")
@require_token
def api_log(job_id):
    p = LOG_DIR / f"{job_id}.log"
    if not p.exists():
        abort(404)
    tail = p.read_text(errors="replace")[-20000:]
    return Response(tail, mimetype="text/plain")


@app.route("/api/download/<job_id>")
@require_token
def api_download(job_id):
    with JOB_LOCK:
        job = JOBS.get(job_id)
    if not job or job["status"] != "done":
        abort(404)
    return send_file(OUTPUT_DIR / job["output"], as_attachment=True,
                     download_name=job["output"])


if __name__ == "__main__":
    load_state()
    print(f"WeBrake listening on http://{HOST}:{PORT}")
    print(f"API token: {TOKENS['api_token']}")
    app.run(host=HOST, port=PORT, threaded=True)