#!/usr/bin/python3
# Copyright (C) 2026 Danilo M. <danix@danix.xyz>
#
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License version 2 as published
# by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General
# Public License for more details.
"""fileops <op> <outdir> <file...>

One local operation on files, for the assistant canvas's Do row. Prints one
JSON line, {"ok": true, "out": [paths]} or {"ok": false, "error": "..."}.
Each result is written under a temporary name and renamed into place, to
the first free name; SIGTERM kills the running tool and removes what was
half written.
"""

import json
import os
import shutil
import signal
import subprocess
import sys
import tarfile
import zipfile

current = None          # the running tool, killed with its group on SIGTERM
temps = []              # partial results to remove if we do not finish

ARCHIVE_EXTS = (".tar.gz", ".tar.xz", ".tar.bz2", ".tar.zst", ".tgz", ".txz",
                ".zip", ".tar", ".7z", ".rar", ".gz", ".xz", ".bz2", ".zst")


class Fail(Exception):
    pass


def finish(result, code):
    print(json.dumps(result), flush=True)
    sys.exit(code)


def stop(*_):
    if current is not None:
        try:
            os.killpg(current.pid, signal.SIGKILL)
        except ProcessLookupError:
            pass
    for t in temps:
        shutil.rmtree(t, ignore_errors=True) if os.path.isdir(t) else _unlink(t)
    finish({"ok": False, "error": "stopped"}, 1)


def _unlink(path):
    try:
        os.unlink(path)
    except FileNotFoundError:
        pass


def run(argv):
    # Own process group, so a stop takes ffmpeg and its children too.
    global current
    if shutil.which(argv[0]) is None:
        hint = ": link it into ~/bin" if argv[0] == "trackcrop" else ""
        raise Fail(f"{argv[0]} not found{hint}")
    current = subprocess.Popen(argv, stdin=subprocess.DEVNULL, stdout=subprocess.PIPE,
                               stderr=subprocess.PIPE, text=True, start_new_session=True)
    out, err = current.communicate()
    rc, current = current.returncode, None
    if rc:
        last = err.strip().splitlines()[-1:] or [f"exit status {rc}"]
        raise Fail(f"{argv[0]}: {last[0]}"[:300])
    return out


def split(name):
    # ("report", ".pdf"), keeping a compound suffix whole: "a.tar.gz".
    low = name.lower()
    for ext in (".tar.gz", ".tar.xz", ".tar.bz2", ".tar.zst"):
        if low.endswith(ext):
            return name[:-len(ext)], name[-len(ext):]
    stem, ext = os.path.splitext(name)
    return stem, ext


def stem(path):
    name = os.path.basename(path)
    for ext in ARCHIVE_EXTS:
        if name.lower().endswith(ext) and len(name) > len(ext):
            return name[:-len(ext)]
    return os.path.splitext(name)[0]


def free(outdir, name):
    path = os.path.join(outdir, name)
    base, ext = split(name)
    n = 1
    while os.path.lexists(path):
        path = os.path.join(outdir, f"{base} ({n}){ext}")
        n += 1
    return path


def make(outdir, name, build, folder=False):
    # build(tmp) writes the result at tmp; it keeps the suffix, which is how
    # ffmpeg and magick choose the format.
    tmp = os.path.join(outdir, f".fileops-{os.getpid()}-{len(temps)}{split(name)[1]}")
    temps.append(tmp)
    if folder:
        os.mkdir(tmp)
    build(tmp)
    final = free(outdir, name)
    os.rename(tmp, final)
    temps.remove(tmp)
    return final


def audio_codec(path):
    return run(["ffprobe", "-v", "error", "-select_streams", "a:0", "-show_entries",
                "stream=codec_name", "-of", "csv=p=0", path]).strip()


# Copied as it is when the track already is one of these.
COPY = {"aac": ".m4a", "mp3": ".mp3", "opus": ".opus"}
GIF = "fps=12,scale=480:-1:flags=lanczos,split[a][b];[a]palettegen[p];[b][p]paletteuse"
# This ffmpeg has no libx264 and no AAC encoder: H.264 is the GPU's, through
# VAAPI, with MP3 audio, which MP4 carries.
VAAPI = ["-vaapi_device", "/dev/dri/renderD128", "-vf", "format=nv12,hwupload",
         "-c:v", "h264_vaapi", "-qp", "23"]


def each(op, outdir, f):
    s, ext = stem(f), os.path.splitext(f)[1]
    ff = ["ffmpeg", "-v", "error", "-y", "-i", f]
    if op == "pdf-images":
        return make(outdir, f"{s} (pages)",
                    lambda t: run(["pdftoppm", "-r", "150", "-png", f, os.path.join(t, "page")]),
                    folder=True)
    if op == "audio":
        codec = audio_codec(f)
        if not codec:
            raise Fail(f"{os.path.basename(f)} has no audio track")
        if codec in COPY:
            return make(outdir, s + COPY[codec], lambda t: run(ff + ["-vn", "-c:a", "copy", t]))
        return make(outdir, s + ".mp3", lambda t: run(ff + ["-vn", "-c:a", "libmp3lame", "-q:a", "2", t]))
    if op == "mp4":
        return make(outdir, s + ".mp4",
                    lambda t: run(ff + VAAPI + ["-c:a", "libmp3lame", "-q:a", "2",
                                               "-movflags", "+faststart", t]))
    if op == "webm":
        return make(outdir, s + ".webm",
                    lambda t: run(ff + ["-c:v", "libvpx-vp9", "-b:v", "0", "-crf", "33",
                                        "-c:a", "libopus", t]))
    if op == "gif":
        return make(outdir, s + ".gif", lambda t: run(ff + ["-vf", GIF, t]))
    if op == "trackcrop":
        return make(outdir, f"{s} (tracked).mp4", lambda t: run(["trackcrop", f, t]))
    if op in ("png", "jpeg", "webp"):
        return make(outdir, f"{s}.{'jpg' if op == 'jpeg' else op}", lambda t: run(["magick", f, t]))
    if op == "half":
        return make(outdir, f"{s} (50%){ext}", lambda t: run(["magick", f, "-resize", "50%", t]))
    if op == "extract":
        return make(outdir, s, lambda t: run(["bsdtar", "-xf", f, "-C", t]), folder=True)
    raise Fail(f"unknown operation {op}")


def compress(op, outdir, files):
    ext = {"zip": ".zip", "tgz": ".tar.gz", "txz": ".tar.xz", "tbz2": ".tar.bz2"}[op]
    name = (stem(files[0]) if len(files) == 1 else "archive") + ext

    def build(t):
        if op == "zip":
            with zipfile.ZipFile(t, "w", zipfile.ZIP_DEFLATED) as z:
                for f in files:
                    if os.path.isdir(f):
                        for root, _, names in os.walk(f):
                            for n in names:
                                p = os.path.join(root, n)
                                z.write(p, os.path.relpath(p, os.path.dirname(f)))
                    else:
                        z.write(f, os.path.basename(f))
        else:
            with tarfile.open(t, "w:" + {"tgz": "gz", "txz": "xz", "tbz2": "bz2"}[op]) as tar:
                for f in files:
                    tar.add(f, arcname=os.path.basename(f))
    return make(outdir, name, build)


def main(argv):
    if len(argv) < 3:
        raise Fail("usage: fileops <op> <outdir> <file...>")
    op, outdir, files = argv[0], argv[1], argv[2:]
    if not os.path.isabs(outdir):
        raise Fail("the output folder must be an absolute path")
    for f in files:
        if not os.path.exists(f):
            raise Fail(f"cannot read {os.path.basename(f)}")
    os.makedirs(outdir, exist_ok=True)
    if op == "pdf-merge":
        if len(files) < 2:
            raise Fail("merging needs two or more PDFs")
        return [make(outdir, "merged.pdf",
                     lambda t: run(["qpdf", "--empty", "--pages", *files, "--", t]))]
    if op in ("zip", "tgz", "txz", "tbz2"):
        return [compress(op, outdir, files)]
    return [each(op, outdir, f) for f in files]


if __name__ == "__main__":
    signal.signal(signal.SIGTERM, stop)
    try:
        out = main(sys.argv[1:])
    except Fail as e:
        for t in temps:
            shutil.rmtree(t, ignore_errors=True) if os.path.isdir(t) else _unlink(t)
        finish({"ok": False, "error": str(e)}, 1)
    except Exception as e:      # a bug: say so rather than hang the card
        for t in temps:
            shutil.rmtree(t, ignore_errors=True) if os.path.isdir(t) else _unlink(t)
        finish({"ok": False, "error": f"{type(e).__name__}: {e}"}, 1)
    finish({"ok": True, "out": out}, 0)