#!/usr/bin/env python3
# uncompressed.io project archive for DaVinci Resolve.
#
# It gathers every clip used on the current timeline, exports the project as a
# .drp, picks up the finished master, and sends all of it, file by file,
# straight into your uncompressed.io cold archive as one project pack. Nothing
# is trimmed or re-encoded: the clips go up exactly as they are on disk. 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)
#
# Two ways to run it, with the project and timeline open in Resolve:
#   1. Terminal (recommended for big timelines, Resolve stays free and you see
#      the speed):   python3 "/path/to/uncompressed-resolve-archive.py"
#      Needs Resolve > Preferences > System > General > External scripting: Local.
#   2. Inside Resolve: put the file in Fusion/Scripts/Utility and pick it from
#      Workspace > Scripts. Output goes to Workspace > Console (Py3).
#      "uncompressed-resolve-archive-PREVIEW" next to it is the same script in
#      --dry-run mode.
#
# Options:
#   --master PATH     use this file as the finished master (default: the latest
#                     completed render in this project's render queue, if any)
#   --no-master       do not include a master
#   --dry-run         list what would be sent, upload nothing
#   --lanes N         parallel upload lanes per file (default 3)
#
# Re-running on the same timeline 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 re
import sys
import threading
import time
import urllib.error
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed

BASE_URL = os.environ.get("UNCOMPRESSED_BASE_URL", "https://uncompressed.io")
LANES = 3
SIGN_BATCH = 16          # parts signed per API call (each lane signs its own batch)
TTY = bool(getattr(sys.stdout, "isatty", lambda: False)())
UA = "uncompressed-resolve-archive/1.1"


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."""


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):
    """POST to the tool API. Rate limits (429) and transient errors are retried with backoff."""
    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 get_resolve():
    try:
        return resolve  # noqa  (predefined when run from Workspace -> Scripts)
    except NameError:
        pass
    dvr = None
    try:
        import DaVinciResolveScript as dvr  # noqa
    except ImportError:
        # Terminal run: find Resolve's scripting module where Blackmagic installs it.
        home = os.path.expanduser("~")
        candidates = [
            ("/Library/Application Support/Blackmagic Design/DaVinci Resolve/Developer/Scripting",
             "/Applications/DaVinci Resolve/DaVinci Resolve.app/Contents/Libraries/Fusion/fusionscript.so"),
            (os.path.join(home, "Library/Application Support/Blackmagic Design/DaVinci Resolve/Developer/Scripting"),
             "/Applications/DaVinci Resolve/DaVinci Resolve.app/Contents/Libraries/Fusion/fusionscript.so"),
            (os.path.join(os.environ.get("PROGRAMDATA", r"C:\ProgramData"), "Blackmagic Design", "DaVinci Resolve", "Support", "Developer", "Scripting"),
             r"C:\Program Files\Blackmagic Design\DaVinci Resolve\fusionscript.dll"),
            ("/opt/resolve/Developer/Scripting", "/opt/resolve/libs/Fusion/fusionscript.so"),
        ]
        for api_dir, lib in candidates:
            modules = os.path.join(api_dir, "Modules")
            if not os.path.isdir(modules):
                continue
            os.environ.setdefault("RESOLVE_SCRIPT_API", api_dir)
            os.environ.setdefault("RESOLVE_SCRIPT_LIB", lib)
            sys.path.insert(0, modules)
            try:
                import DaVinciResolveScript as dvr  # noqa
                break
            except ImportError:
                continue
    if dvr is None:
        return None
    try:
        return dvr.scriptapp("Resolve")
    except Exception:
        return None


def expand_sequence(path):
    """Resolve reports image sequences as one path with a [first-last] token: return the frames."""
    m = re.match(r"^(.*?)\[(\d+)-(\d+)\](.*)$", path)
    if not m:
        return [path]
    prefix, first, last, suffix = m.group(1), m.group(2), m.group(3), m.group(4)
    folder = os.path.dirname(prefix) or "."
    stem = os.path.basename(prefix)
    width = len(first)
    frames = []
    try:
        for name in sorted(os.listdir(folder)):
            if name.startswith(stem) and name.endswith(suffix):
                digits = name[len(stem):len(name) - len(suffix)]
                if digits.isdigit() and len(digits) == width and int(first) <= int(digits) <= int(last):
                    frames.append(os.path.join(folder, name))
    except OSError:
        pass
    return frames or [path]


def timeline_files(timeline):
    seen, files, skipped = set(), [], []
    for kind in ("video", "audio"):
        try:
            tracks = int(timeline.GetTrackCount(kind))
        except Exception:
            tracks = 0
        for t in range(1, tracks + 1):
            for item in timeline.GetItemListInTrack(kind, t) or []:
                try:
                    mpi = item.GetMediaPoolItem()
                except Exception:
                    mpi = None
                if not mpi:
                    skipped.append(getattr(item, "GetName", lambda: "timeline item")())
                    continue
                try:
                    path = mpi.GetClipProperty("File Path") or ""
                except Exception:
                    path = ""
                if not path:
                    try:
                        skipped.append(mpi.GetName())
                    except Exception:
                        skipped.append("clip without a file")
                    continue
                for p in expand_sequence(path):
                    if p in seen:
                        continue
                    seen.add(p)
                    if os.path.isfile(p):
                        files.append(p)
                    else:
                        skipped.append(p)
    return files, skipped


def latest_completed_render(project):
    best = None
    for job in project.GetRenderJobList() or []:
        try:
            status = project.GetRenderJobStatus(job.get("JobId"))
        except Exception:
            status = {}
        if (status or {}).get("JobStatus") != "Complete":
            continue
        target, name = job.get("TargetDir") or "", job.get("OutputFilename") or ""
        if not target or not name:
            continue
        path = os.path.join(target, name)
        if not os.path.isfile(path):
            # Resolve may add the extension/frame token itself; take the newest file with that stem.
            stem = os.path.splitext(name)[0]
            try:
                cands = [os.path.join(target, f) for f in os.listdir(target) if f.startswith(stem)]
            except OSError:
                cands = []
            cands = [c for c in cands if os.path.isfile(c)]
            if not cands:
                continue
            path = max(cands, key=os.path.getmtime)
        mtime = os.path.getmtime(path)
        if best is None or mtime > best[0]:
            best = (mtime, path)
    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(project_name, timeline_name):
    tag = hashlib.sha1(("%s|%s" % (project_name, timeline_name)).encode("utf-8")).hexdigest()[:12]
    return os.path.expanduser("~/.uncompressed-pack-%s.json" % tag)


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)  # the URL had expired: a fresh one
            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")


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

    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

    r = get_resolve()
    if r is None:
        print("Couldn't reach DaVinci Resolve. Open the project in Resolve first, and in")
        print("Preferences > System > General set 'External scripting using' to Local.")
        print("Or run this script from inside Resolve: Workspace > Scripts.")
        return 1
    pm = r.GetProjectManager()
    project = pm.GetCurrentProject()
    if project is None:
        print("No project open.")
        return 1
    timeline = project.GetCurrentTimeline()
    if timeline is None:
        print("No timeline open. Open the timeline you want to archive and run again.")
        return 1
    project_name, timeline_name = project.GetName(), timeline.GetName()

    print("Project:  %s" % project_name)
    print("Timeline: %s" % timeline_name)
    files, skipped = timeline_files(timeline)
    if not files:
        print("No clips with files on disk were found on this timeline.")
        return 1

    # The project file: Resolve writes it where we ask.
    export_dir = os.path.join(os.path.expanduser("~"), "Movies", "uncompressed-archives")
    os.makedirs(export_dir, exist_ok=True)
    drp_path = os.path.join(export_dir, re.sub(r"[^\w.-]+", "_", project_name) + ".drp")
    try:
        pm.SaveProject()
    except Exception:
        pass
    exported = False
    try:
        exported = bool(pm.ExportProject(project_name, drp_path))
    except Exception:
        exported = False
    if not exported or not os.path.isfile(drp_path):
        print("Could not export the project file (.drp); the pack will go up without it.")
        drp_path = None

    master = None if no_master else (master_arg or latest_completed_render(project))
    if master and not os.path.isfile(master):
        print("Master not found on disk: %s" % master)
        master = None

    root = common_root(files)
    plan = [(p, rel_path(p, root), "clip") for p in files]
    if drp_path:
        plan.append((drp_path, os.path.basename(drp_path), "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 (no file on disk, compound or generated clip): %d" % len(skipped))
        for s in skipped[:10]:
            print("  - %s" % s)
    if not master and not no_master:
        print("No completed render found; pass --master PATH to include the finished master.")
    if dry_run:
        print("\nDry run: nothing sent.")
        return 0

    # Resume support: one pack per project+timeline until it is finished.
    state = State(state_file(project_name, timeline_name))
    pack_id = state.data.get("packId")
    if not pack_id:
        try:
            begin = api("/api/tool/pack/begin", {
                "name": project_name, "nle": "resolve", "timeline": timeline_name, "totalBytes": total,
                "manifest": {"resolve": getattr(r, "GetVersionString", lambda: "")(), "fps": timeline.GetSetting("timelineFrameRate"),
                             "frames": [timeline.GetStartFrame(), timeline.GetEndFrame()], "root": root, "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:]))