#!/usr/bin/env python3
# uncompressed.io project archive for Final Cut Pro.
#
# Final Cut has no scripting door into a timeline, so the archive starts from
# the XML you export: in Final Cut, select the project, File > Export XML...
# (this makes a .fcpxmld package, or a single .fcpxml on older versions). This
# script reads it, finds every clip used in the project, and sends the original
# media files, the XML itself and the finished master, file by file, straight
# into your uncompressed.io cold archive as one project pack. Nothing is
# trimmed or re-encoded; proxies and optimized media are ignored, only the
# originals go up. Each file then lives in your dashboard under Project
# archives, where it can be downloaded, re-hydrated, shared or sent back to
# cold on its own.
#
# Setup, once: ~/.uncompressed containing   api_key=uc_...   (dashboard > Settings > API key)
#
# Run it from Terminal:
#   python3 uncompressed-fcp-archive.py "/path/to/My Film.fcpxmld"
#   python3 uncompressed-fcp-archive.py            (no path: a file chooser opens)
#
# Options:
#   --master PATH     use this file as the finished master (default: the newest
#                     video file named like the project, next to the XML or in
#                     Movies, Desktop, Downloads)
#   --pick-master     choose the master in a file dialog
#   --no-master       do not include a master
#   --project NAME    which project to archive when the XML holds several
#   --dry-run         list what would be sent, upload nothing
#   --lanes N         parallel upload lanes per file (default 3)
#
# Re-running on the same project resumes: finished files are skipped, and a
# file that was interrupted continues from the last part that made it.

import hashlib
import json
import math
import os
import subprocess
import sys
import threading
import time
import urllib.error
import urllib.parse
import urllib.request
import xml.etree.ElementTree as ET
from concurrent.futures import ThreadPoolExecutor, as_completed

BASE_URL = os.environ.get("UNCOMPRESSED_BASE_URL", "https://uncompressed.io")
LANES = 3
SIGN_BATCH = 16
TTY = bool(getattr(sys.stdout, "isatty", lambda: False)())
UA = "uncompressed-fcp-archive/1.0"
VIDEO_EXT = (".mov", ".mp4", ".m4v", ".mxf", ".mkv", ".avi", ".prores", ".webm")


class ApiError(Exception):
    def __init__(self, status, payload):
        Exception.__init__(self, "%s: %s" % (status, payload.get("message") or payload.get("error") or payload))
        self.status = status
        self.payload = payload


class PartRefused(Exception):
    """A presigned part URL was refused (expired or rejected): re-sign and retry."""


# ---------------------------------------------------------------- API + upload

def load_api_key():
    key = os.environ.get("UNCOMPRESSED_API_KEY", "").strip()
    if key:
        return key
    try:
        with open(os.path.expanduser("~/.uncompressed")) as f:
            for line in f:
                line = line.strip()
                if line.startswith("api_key="):
                    return line.split("=", 1)[1].strip()
    except OSError:
        pass
    return ""


def api(path, body, api_key, tries=6):
    data = json.dumps(body).encode("utf-8")
    for attempt in range(1, tries + 1):
        req = urllib.request.Request(BASE_URL + path, data=data, method="POST", headers={
            "Content-Type": "application/json", "Authorization": "Bearer " + api_key, "User-Agent": UA})
        try:
            with urllib.request.urlopen(req, timeout=120) as r:
                return json.loads(r.read().decode("utf-8") or "{}")
        except urllib.error.HTTPError as e:
            try:
                payload = json.loads(e.read().decode("utf-8") or "{}")
            except ValueError:
                payload = {}
            final = payload.get("error") in ("packs_unavailable", "cold_unavailable")
            if e.code in (429, 500, 502, 503, 504) and attempt < tries and not final:
                try:
                    wait = float(e.headers.get("Retry-After") or 0)
                except ValueError:
                    wait = 0
                time.sleep(wait or min(60, 2 ** attempt))
                continue
            raise ApiError(e.code, payload)
        except (urllib.error.URLError, OSError):
            if attempt == tries:
                raise
            time.sleep(min(60, 2 ** attempt))


def put_part(url, chunk):
    for attempt in range(1, 4):
        req = urllib.request.Request(url, data=chunk, method="PUT",
                                     headers={"Content-Type": "application/octet-stream", "User-Agent": UA})
        try:
            with urllib.request.urlopen(req, timeout=600) as r:
                etag = r.headers.get("ETag") or ""
                if etag:
                    return etag.strip('"')
                raise RuntimeError("no ETag after upload")
        except urllib.error.HTTPError as e:
            if e.code in (400, 403, 404):
                raise PartRefused(str(e.code))
            if attempt == 3:
                raise
        except (urllib.error.URLError, OSError):
            if attempt == 3:
                raise
        time.sleep(2 * attempt)


def fmt_bytes(n):
    n = float(n)
    for unit in ("B", "KB", "MB", "GB", "TB"):
        if n < 1000 or unit == "TB":
            return ("%.0f %s" if unit in ("B", "KB") else "%.2f %s") % (n, unit)
        n /= 1000.0


def fmt_eta(seconds):
    if seconds is None or seconds != seconds or seconds == float("inf"):
        return "--"
    seconds = int(seconds)
    if seconds < 90:
        return "%ds" % seconds
    if seconds < 5400:
        return "%dm" % (seconds // 60)
    return "%dh%02dm" % (seconds // 3600, (seconds % 3600) // 60)


class State(object):
    """~/.uncompressed-pack-<id>.json: the pack id, finished files, and the parts of the file in flight."""

    def __init__(self, path):
        self.path = path
        self.lock = threading.Lock()
        self.data = {}
        try:
            with open(path) as f:
                self.data = json.load(f)
        except (OSError, ValueError):
            self.data = {}
        self.data.setdefault("done", {})
        self.data.setdefault("partial", {})

    def save(self):
        tmp = self.path + ".tmp"
        with open(tmp, "w") as f:
            json.dump(self.data, f)
        os.replace(tmp, self.path)

    def remove(self):
        try:
            os.remove(self.path)
        except OSError:
            pass


def upload_file(path, rel, role, pack_id, api_key, lanes, state):
    size = os.path.getsize(path)
    filename = os.path.basename(path)
    partial = state.data["partial"].get(rel)
    if partial and partial.get("size") == size and partial.get("key") and partial.get("uploadId"):
        key, upload_id, part_size = partial["key"], partial["uploadId"], int(partial["partSize"])
        parts_done = dict((int(k), v) for k, v in (partial.get("parts") or {}).items())
        if parts_done:
            print("    resuming, %d parts already sent" % len(parts_done))
    else:
        begin = api("/api/tool/pack/item/begin", {"packId": pack_id, "filename": filename, "sizeBytes": size, "role": role}, api_key)
        key, upload_id, part_size = begin["key"], begin["uploadId"], int(begin["partSize"])
        parts_done = {}
        partial = {"key": key, "uploadId": upload_id, "size": size, "partSize": part_size, "parts": {}}
        with state.lock:
            state.data["partial"][rel] = partial
            state.save()

    total_parts = max(1, int(math.ceil(size / float(part_size)))) if size else 1
    pending = [n for n in range(1, total_parts + 1) if n not in parts_done]
    sent_before = sum(min(part_size, size - (n - 1) * part_size) for n in parts_done)
    progress = {"bytes": sent_before, "session": 0, "t0": time.time(), "printed": 0.0}

    def show(force=False):
        # A terminal gets one line refreshed every 2 s; Resolve's Console (no tty) gets a new line every 30 s.
        now = time.time()
        if not force and now - progress["printed"] < (2 if TTY else 30):
            return
        progress["printed"] = now
        elapsed = max(0.001, now - progress["t0"])
        rate = progress["session"] / elapsed
        remaining = size - progress["bytes"]
        eta = remaining / rate if rate > 0 else None
        line = "    %-44s %3d%%  %s / %s  %s/s  ETA %s   " % (
            filename[:44], 100 * progress["bytes"] // max(1, size), fmt_bytes(progress["bytes"]), fmt_bytes(size),
            fmt_bytes(rate), fmt_eta(eta))
        sys.stdout.write(("\r" + line) if TTY else (line + "\n"))
        sys.stdout.flush()

    def sign(numbers):
        urls = api("/api/tool/pack/item/sign", {"key": key, "uploadId": upload_id, "partNumbers": numbers}, api_key)["urls"]
        return dict((int(k), v) for k, v in urls.items())

    def lane(batch):
        urls = sign(batch)
        for n in batch:
            with open(path, "rb") as f:
                f.seek((n - 1) * part_size)
                chunk = f.read(part_size)
            try:
                etag = put_part(urls[n], chunk)
            except PartRefused:
                etag = put_part(sign([n])[n], chunk)
            with state.lock:
                parts_done[n] = etag
                partial["parts"][str(n)] = etag
                progress["bytes"] += len(chunk)
                progress["session"] += len(chunk)
                state.save()
                show()
        return len(batch)

    if pending:
        show(force=True)
        batches = [pending[i:i + SIGN_BATCH] for i in range(0, len(pending), SIGN_BATCH)]
        with ThreadPoolExecutor(max_workers=lanes) as pool:
            for fut in as_completed([pool.submit(lane, b) for b in batches]):
                fut.result()
        show(force=True)
        if TTY:
            sys.stdout.write("\n")
    parts = [{"partNumber": n, "etag": parts_done[n]} for n in sorted(parts_done)]
    fin = api("/api/tool/pack/item/finish", {"packId": pack_id, "key": key, "uploadId": upload_id, "parts": parts,
                                              "sizeBytes": size, "filename": filename, "path": rel, "role": role}, api_key)
    with state.lock:
        state.data["partial"].pop(rel, None)
        state.data["done"][rel] = fin.get("videoId")
        state.save()
    return fin.get("videoId")


# ---------------------------------------------------------------- FCPXML

def choose_file(prompt):
    """macOS file dialog; returns a POSIX path or None (Cancel)."""
    if sys.platform != "darwin":
        return None
    try:
        out = subprocess.run(["osascript", "-e", 'POSIX path of (choose file with prompt "%s")' % prompt.replace('"', "'")],
                             capture_output=True, text=True, timeout=600)
    except (OSError, subprocess.SubprocessError):
        return None
    path = out.stdout.strip()
    return path.rstrip("/") if out.returncode == 0 and path else None


def locate_xml(arg):
    """Accept a .fcpxml file or a .fcpxmld package; return (xml_path, bundle_dir_or_None)."""
    arg = os.path.abspath(os.path.expanduser(arg.rstrip("/")))
    if os.path.isdir(arg):
        inner = os.path.join(arg, "Info.fcpxml")
        if os.path.isfile(inner):
            return inner, arg
        for name in os.listdir(arg):
            if name.lower().endswith(".fcpxml"):
                return os.path.join(arg, name), arg
        return None, None
    if os.path.isfile(arg):
        return arg, None
    return None, None


def src_to_path(src, base_dir):
    if not src:
        return ""
    if src.startswith("file:"):
        u = urllib.parse.urlparse(src)
        return urllib.parse.unquote(u.path)
    return os.path.normpath(os.path.join(base_dir, urllib.parse.unquote(src)))


def collect_refs(node, resources, seen_media, refs):
    """Walk a timeline (sequence / spine / clips) and collect every resource id it references."""
    for el in node.iter():
        ref = el.get("ref")
        if not ref:
            continue
        refs.add(ref)
        res = resources.get(ref)
        # A media resource (compound clip, multicam) is a timeline of its own: dig in once.
        if res is not None and res.tag == "media" and ref not in seen_media:
            seen_media.add(ref)
            collect_refs(res, resources, seen_media, refs)


def project_files(root, project, base_dir):
    resources = {}
    for res in root.find("resources") or []:
        rid = res.get("id")
        if rid:
            resources[rid] = res
    refs, seen_media = set(), set()
    for seq in project.findall("sequence"):
        collect_refs(seq, resources, seen_media, refs)

    files, skipped, seen = [], [], set()
    for rid in sorted(refs):
        res = resources.get(rid)
        if res is None or res.tag != "asset":
            continue
        src = ""
        for rep in res.findall("media-rep"):
            if (rep.get("kind") or "original-media") == "original-media":
                src = rep.get("src") or ""
                break
        if not src:
            src = res.get("src") or ""  # FCPXML 1.8 and older keep the path on the asset
        path = src_to_path(src, base_dir)
        name = res.get("name") or rid
        if not path:
            skipped.append("%s (no original media, proxy only?)" % name)
            continue
        if path in seen:
            continue
        seen.add(path)
        if os.path.isfile(path):
            files.append(path)
        else:
            skipped.append("%s: %s (not on disk)" % (name, path))
    return files, skipped


def find_master(project_name, near_dir):
    stem = project_name.lower()
    home = os.path.expanduser("~")
    best = None
    for folder in [near_dir, os.path.join(home, "Movies"), os.path.join(home, "Desktop"), os.path.join(home, "Downloads")]:
        if not folder or not os.path.isdir(folder):
            continue
        try:
            names = os.listdir(folder)
        except OSError:
            continue
        for n in names:
            if n.lower().startswith(stem) and n.lower().endswith(VIDEO_EXT):
                p = os.path.join(folder, n)
                if os.path.isfile(p):
                    m = os.path.getmtime(p)
                    if best is None or m > best[0]:
                        best = (m, p)
    return best[1] if best else None


def common_root(paths):
    try:
        root = os.path.commonpath(paths)
    except ValueError:
        return ""
    return root if os.path.isdir(root) else os.path.dirname(root)


def rel_path(path, root):
    if root and path.startswith(root):
        return path[len(root):].lstrip("/")
    return os.path.basename(path)


def state_file(library, project_name):
    tag = hashlib.sha1(("%s|%s" % (library, project_name)).encode("utf-8")).hexdigest()[:12]
    return os.path.expanduser("~/.uncompressed-pack-%s.json" % tag)


def main(argv):
    xml_arg, master_arg, pick_master, no_master, project_arg, dry_run, lanes = None, None, False, False, None, False, LANES
    args = list(argv)
    while args:
        a = args.pop(0)
        if a == "--master" and args:
            master_arg = args.pop(0)
        elif a == "--pick-master":
            pick_master = True
        elif a == "--no-master":
            no_master = True
        elif a == "--project" and args:
            project_arg = args.pop(0)
        elif a == "--dry-run":
            dry_run = True
        elif a == "--lanes" and args:
            lanes = max(1, min(8, int(args.pop(0))))
        elif not a.startswith("--"):
            xml_arg = a

    api_key = load_api_key()
    if not api_key and not dry_run:
        print("No API key found. Create ~/.uncompressed containing:  api_key=YOUR_KEY")
        print("Your key is in the uncompressed.io dashboard, Settings > API key.")
        return 1

    if not xml_arg:
        xml_arg = choose_file("Choose the XML exported from Final Cut Pro (File > Export XML)")
        if not xml_arg:
            print("Usage: python3 uncompressed-fcp-archive.py \"/path/to/My Film.fcpxmld\"")
            return 1
    xml_path, bundle = locate_xml(xml_arg)
    if not xml_path:
        print("No FCPXML found at %s. In Final Cut Pro, select the project and use File > Export XML..." % xml_arg)
        return 1
    base_dir = bundle or os.path.dirname(xml_path)

    try:
        root = ET.parse(xml_path).getroot()
    except ET.ParseError as e:
        print("That XML could not be read: %s" % e)
        return 1
    if root.tag != "fcpxml":
        print("That file is not an FCPXML export.")
        return 1

    projects = []
    for ev in root.iter("event"):
        for pr in ev.findall("project"):
            projects.append((ev.get("name") or "", pr))
    for pr in root.findall("project") + [p for lib in root.findall("library") for p in lib.findall("project")]:
        projects.append(("", pr))
    if not projects:
        print("No project inside that XML. Export the project itself (select it in the browser, then File > Export XML).")
        return 1
    chosen = None
    if project_arg:
        for ev, pr in projects:
            if (pr.get("name") or "").lower() == project_arg.lower():
                chosen = (ev, pr)
    elif len(projects) == 1:
        chosen = projects[0]
    else:
        print("This XML holds %d projects; pick one with --project NAME:" % len(projects))
        for ev, pr in projects:
            print("  - %s%s" % (pr.get("name") or "(untitled)", ("  [event: %s]" % ev) if ev else ""))
        return 1
    if not chosen:
        print("No project named %s in that XML." % project_arg)
        return 1
    event_name, project = chosen
    project_name = project.get("name") or os.path.splitext(os.path.basename(bundle or xml_path))[0]
    lib = root.find("library")
    library = (lib.get("location") if lib is not None else "") or ""

    print("Project: %s" % project_name)
    if event_name:
        print("Event:   %s" % event_name)
    files, skipped = project_files(root, project, base_dir)
    if not files:
        print("No original media files were found for this project's clips.")
        for s in skipped[:10]:
            print("  - %s" % s)
        return 1

    # The project file: the XML export itself (every file of a .fcpxmld package).
    project_paths = []
    if bundle:
        for dp, _, names in os.walk(bundle):
            for n in names:
                if not n.startswith("."):
                    project_paths.append(os.path.join(dp, n))
    else:
        project_paths.append(xml_path)
    xml_root_name = os.path.basename(bundle or xml_path)

    master = None
    if not no_master:
        master = master_arg or (choose_file("Choose the finished master (Cancel to skip)") if pick_master else None) \
            or find_master(project_name, os.path.dirname(bundle or xml_path))
    if master and not os.path.isfile(master):
        print("Master not found on disk: %s" % master)
        master = None

    root_dir = common_root(files)
    plan = [(p, rel_path(p, root_dir), "clip") for p in files]
    for p in project_paths:
        inner = os.path.relpath(p, os.path.dirname(bundle)) if bundle else os.path.basename(p)
        plan.append((p, "project/" + inner, "project"))
    if master:
        plan.append((master, "master/" + os.path.basename(master), "master"))
    total = sum(os.path.getsize(p) for p, _, _ in plan)

    print("")
    print("%d files, %s" % (len(plan), fmt_bytes(total)))
    for p, rel, role in plan:
        print("  [%-7s] %-60s %10s" % (role, rel[:60], fmt_bytes(os.path.getsize(p))))
    if skipped:
        print("Skipped (%d): generators, titles, or media that is not on disk" % len(skipped))
        for s in skipped[:10]:
            print("  - %s" % s)
    if not master and not no_master:
        print("No master found next to the XML or in Movies/Desktop/Downloads; pass --master PATH or --pick-master to include it.")
    if dry_run:
        print("\nDry run: nothing sent.")
        return 0

    state = State(state_file(library or base_dir, project_name))
    pack_id = state.data.get("packId")
    if not pack_id:
        try:
            begin = api("/api/tool/pack/begin", {
                "name": project_name, "nle": "fcp", "timeline": event_name or None, "totalBytes": total,
                "manifest": {"fcpxml": root.get("version") or "", "library": library, "event": event_name,
                             "export": xml_root_name, "root": root_dir, "files": len(plan)},
            }, api_key)
        except ApiError as e:
            if e.status == 401:
                print("Your API key was refused. Check ~/.uncompressed against the key in your dashboard.")
            else:
                print("Couldn't open the pack: %s" % e)
            return 1
        pack_id = begin["packId"]
        state.data["packId"] = pack_id
        state.save()

    print("\nSending to your cold archive (%d lanes per file). Interrupt any time; running again resumes." % lanes)
    t_start = time.time()
    for p, rel, role in plan:
        if rel in state.data["done"]:
            print("  already sent: %s" % rel)
            continue
        print("  %s" % rel)
        try:
            upload_file(p, rel, role, pack_id, api_key, lanes, state)
        except ApiError as e:
            print("\n  %s: %s" % (rel, e))
            if e.status == 413:
                print("  Add cold storage on your plan page, then run the script again; finished files are kept.")
            return 1
        except KeyboardInterrupt:
            print("\n  Stopped. Run the script again to resume from where it left off.")
            return 1
        except Exception as e:  # network, disk
            print("\n  %s failed: %s. Run the script again to resume." % (rel, e))
            return 1

    fin = api("/api/tool/pack/finish", {"packId": pack_id}, api_key)
    state.remove()
    print("\nDone: %d files, %s in the cold archive, %s." % (
        fin.get("itemCount", len(plan)), fmt_bytes(fin.get("totalBytes", total)), fmt_eta(time.time() - t_start)))
    print("Open it: %s" % fin.get("url", BASE_URL + "/dashboard/packs/" + pack_id))
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))