#!/usr/bin/env python3
# uncompressed-resolve.py — upload your latest finished render to uncompressed.io
# straight from DaVinci Resolve's Workspace → Scripts menu.
#
# WHAT IT DOES
#   1. Finds the CURRENT project's most recent COMPLETED render job.
#   2. Uploads that file to uncompressed.io in 10 MB pieces, three at a time.
#   3. Prints the watch link (unlisted by default) and the manage link.
#
# SETUP (once)
#   1. Put this file in Resolve's Scripts folder:
#        macOS:   ~/Library/Application Support/Blackmagic Design/DaVinci Resolve/Fusion/Scripts/Utility
#        Windows: %APPDATA%\Blackmagic Design\DaVinci Resolve\Support\Fusion\Scripts\Utility
#   2. Create a file called .uncompressed in your home folder containing one line:
#        api_key=uc_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
#      Your key lives in the uncompressed.io dashboard, "API key" panel.
#      (The UNCOMPRESSED_API_KEY environment variable works too, and a
#      .uncompressed file sitting next to this script wins over both.)
#   3. Render with your usual master preset, then Workspace → Scripts →
#      Utility → uncompressed-resolve. Progress appears in the Console
#      (Workspace → Console).
#
# Pure Python 3 standard library — Resolve's bundled Python has no pip, and
# this script never needs it. Tinker away; every step is commented.

import json
import math
import os
import sys
import threading
import time
import urllib.error
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed

# ————— Tunables ——————————————————————————————————————————————————————————————
BASE_URL = os.environ.get("UNCOMPRESSED_BASE_URL", "https://uncompressed.io")
LANES = 3          # parallel upload lanes; 3 saturates most home fiber
RETRIES = 3        # attempts per part before giving up
TIMEOUT = 120      # seconds per HTTP call (a 10 MB part on slow DSL fits)


# ————— API key ———————————————————————————————————————————————————————————————
def _parse_config(path):
    """Read key=value lines from a config file; ignore everything else."""
    values = {}
    try:
        with open(path, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if "=" in line and not line.startswith("#"):
                    k, _, v = line.partition("=")
                    values[k.strip()] = v.strip()
    except OSError:
        pass
    return values


def load_api_key():
    """Config file next to the script wins, then ~/.uncompressed, then env."""
    here = os.path.dirname(os.path.abspath(__file__))
    for path in (os.path.join(here, ".uncompressed"),
                 os.path.join(os.path.expanduser("~"), ".uncompressed")):
        key = _parse_config(path).get("api_key", "")
        if key:
            return key
    return os.environ.get("UNCOMPRESSED_API_KEY", "")


# ————— HTTP helpers (urllib only) ————————————————————————————————————————————
class ApiError(Exception):
    def __init__(self, status, message):
        super().__init__(message)
        self.status = status


def api(path, payload, key):
    """POST JSON to the tool API, return the parsed JSON reply."""
    req = urllib.request.Request(
        BASE_URL + path,
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json",
                 "Authorization": "Bearer " + key},
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=TIMEOUT) as r:
            return json.loads(r.read().decode("utf-8"))
    except urllib.error.HTTPError as e:
        try:
            body = json.loads(e.read().decode("utf-8"))
            message = body.get("message") or body.get("error") or str(e)
        except Exception:
            message = str(e)
        raise ApiError(e.code, message)


def put_bytes(url, chunk):
    """PUT one part to its presigned URL; return the ETag header."""
    req = urllib.request.Request(url, data=chunk, method="PUT")
    with urllib.request.urlopen(req, timeout=TIMEOUT) as r:
        etag = r.headers.get("ETag")
        if not etag:
            raise RuntimeError("no ETag on part response")
        return etag


# ————— Finding the render inside Resolve —————————————————————————————————————
def get_resolve():
    """Scripts run from the menu get a `resolve` global; fall back to the
    scripting module for people running this from an external shell."""
    r = globals().get("resolve")
    if r is not None:
        return r
    try:
        import DaVinciResolveScript as dvr  # type: ignore
        return dvr.scriptapp("Resolve")
    except ImportError:
        return None


def latest_completed_render(project):
    """Walk the render queue newest-first and return (path, job) for the last
    job whose status is Complete. Jobs stay in the queue after rendering,
    which is exactly what makes this lookup possible."""
    jobs = project.GetRenderJobList() or []
    for job in reversed(jobs):
        job_id = job.get("JobId")
        if not job_id:
            continue
        try:
            status = project.GetRenderJobStatus(job_id) or {}
        except Exception:
            status = {}
        if status.get("JobStatus") != "Complete":
            continue
        target_dir = job.get("TargetDir") or ""
        filename = job.get("OutputFilename") or ""
        if not target_dir or not filename:
            continue
        if "%" in filename:
            # An image sequence (name_%08d.exr) — not a single uploadable file.
            continue
        path = os.path.join(target_dir, filename)
        if os.path.isfile(path):
            return path, job
    return None, None


# ————— Render preset self-install ————————————————————————————————————————————
PRESET_NAME = "uncompressed.io master (HEVC)"


def ensure_render_preset(project):
    """First run only: save an 'uncompressed.io master (HEVC)' preset into the
    Deliver page, so every future MANUAL export can use it too. H.265 with
    EncodingProfile Main10 (that's the 10-bit switch) in an MP4, restricted
    to 150,000 kb/s — the 4K master target from uncompressed.io/nle. For
    1080p work, duplicate it in the Deliver page and drop to ~40,000.

    Everything sits in try/except: older Resolve builds miss some of these
    calls, and a preset is a convenience, never a reason to stop the upload.
    """
    try:
        presets = project.GetRenderPresetList() or []
        if PRESET_NAME in presets:
            return
        ok = bool(project.SetCurrentRenderFormatAndCodec("mp4", "H265"))
        # VideoQuality is the max bitrate in kb/s (0 would mean "automatic").
        # The API has no AAC-bitrate knob; set audio to 320 kb/s by hand once.
        ok = bool(project.SetRenderSettings({
            "SelectAllFrames": True,
            "ExportVideo": True,
            "ExportAudio": True,
            "VideoQuality": 150000,
            "EncodingProfile": "Main10",
            "AudioCodec": "aac",
            "AudioSampleRate": 48000,
        })) and ok
        if ok and project.SaveAsNewRenderPreset(PRESET_NAME):
            print("Installed render preset '%s' — it now lives in the" % PRESET_NAME)
            print("Deliver page's preset list for all your future exports.")
        else:
            raise RuntimeError("the Resolve API declined one of the calls")
    except Exception as e:
        print("Couldn't auto-install the '%s' preset (%s)." % (PRESET_NAME, e))
        print("No harm done — the settings to recreate it are at %s/nle" % BASE_URL)


# ————— Upload ————————————————————————————————————————————————————————————————
class Progress:
    """Thread-safe counter that prints a line whenever the percent moves."""

    def __init__(self, total_parts):
        self.lock = threading.Lock()
        self.done = 0
        self.total = total_parts
        self.last_pct = -1

    def bump(self):
        with self.lock:
            self.done += 1
            pct = int(self.done * 100 / self.total)
            if pct != self.last_pct:
                self.last_pct = pct
                print("  %d%%  (%d/%d parts)" % (pct, self.done, self.total))
                sys.stdout.flush()


def upload_part(path, upload_key, upload_id, part_number, part_size, api_key, progress):
    """Sign, read, PUT one part — with retries and a fresh URL each attempt.
    Each part opens its own file handle, so lanes never fight over a seek."""
    offset = (part_number - 1) * part_size
    with open(path, "rb") as f:
        f.seek(offset)
        chunk = f.read(part_size)

    last_err = None
    for attempt in range(1, RETRIES + 1):
        try:
            signed = api("/api/tool/sign",
                         {"key": upload_key, "uploadId": upload_id,
                          "partNumbers": [part_number]}, api_key)
            url = signed["urls"][str(part_number)]
            etag = put_bytes(url, chunk)
            progress.bump()
            return {"partNumber": part_number, "etag": etag}
        except ApiError:
            raise  # auth/quota problems won't heal on retry
        except Exception as e:  # network hiccup, expired URL, R2 5xx…
            last_err = e
            if attempt < RETRIES:
                time.sleep(2 ** attempt)  # 2 s, 4 s
    raise RuntimeError("part %d failed after %d tries: %s"
                       % (part_number, RETRIES, last_err))


def main():
    api_key = load_api_key()
    if not api_key:
        print("No API key found.")
        print("Create ~/.uncompressed containing:  api_key=YOUR_KEY")
        print("Your key is in the uncompressed.io dashboard, 'API key' panel.")
        return

    r = get_resolve()
    if r is None:
        print("Couldn't reach the Resolve scripting API — run this from")
        print("Workspace → Scripts inside DaVinci Resolve.")
        return
    project = r.GetProjectManager().GetCurrentProject()
    if project is None:
        print("No project open.")
        return

    # One-time nicety: put the master preset into the Deliver page.
    ensure_render_preset(project)

    path, job = latest_completed_render(project)
    if not path:
        print("No completed render found in this project's render queue.")
        print("Render your timeline first (your master preset), then run this again.")
        return

    size = os.path.getsize(path)
    filename = os.path.basename(path)
    title = os.path.splitext(filename)[0] or project.GetName()
    print("Uploading: %s  (%.2f GB)" % (path, size / 1024 ** 3))

    # Step 1 — open the upload. The server picks the part size (10 MB until
    # files get enormous) and enforces your plan's storage quota.
    try:
        begin = api("/api/tool/begin",
                    {"filename": filename, "sizeBytes": size, "title": title},
                    api_key)
    except ApiError as e:
        if e.status == 401:
            print("Your API key was refused — check ~/.uncompressed (or")
            print("UNCOMPRESSED_API_KEY) against the key in your dashboard.")
        else:
            print("Couldn't start the upload: %s" % e)
        return

    upload_key, upload_id = begin["key"], begin["uploadId"]
    part_size = int(begin["partSize"])
    total_parts = max(1, math.ceil(size / part_size))
    progress = Progress(total_parts)
    print("  %d parts of %d MB, %d lanes" % (total_parts, part_size // 1024 ** 2, LANES))

    # Step 2 — the parts, three lanes wide. Memory stays tiny: one part per
    # lane in flight, so ~30 MB even for a 50 GB master.
    parts = []
    try:
        with ThreadPoolExecutor(max_workers=LANES) as pool:
            futures = [pool.submit(upload_part, path, upload_key, upload_id,
                                   n, part_size, api_key, progress)
                       for n in range(1, total_parts + 1)]
            for fut in as_completed(futures):
                parts.append(fut.result())  # re-raises the first failure
    except ApiError as e:
        print("Upload stopped: %s" % e)
        return
    except Exception as e:
        print("Upload failed: %s" % e)
        print("Nothing was published. Re-run the script to try again.")
        return

    # Step 3 — close the upload. Width/height/frame rate ride along when the
    # render job knows them; the video lands unlisted either way.
    meta = {"key": upload_key, "uploadId": upload_id, "parts": parts, "title": title}
    fps = None
    try:
        fps = float(job.get("FrameRate") or 0) or None
    except (TypeError, ValueError):
        pass
    if job.get("FormatWidth"):
        meta["width"] = int(job["FormatWidth"])
    if job.get("FormatHeight"):
        meta["height"] = int(job["FormatHeight"])
    if fps:
        meta["fps"] = fps
        mark_in, mark_out = job.get("MarkIn"), job.get("MarkOut")
        if isinstance(mark_in, int) and isinstance(mark_out, int) and mark_out > mark_in:
            meta["duration"] = (mark_out - mark_in + 1) / fps

    try:
        done = api("/api/tool/finish", meta, api_key)
    except ApiError as e:
        print("The parts uploaded but finishing failed: %s" % e)
        return

    print("")
    print("Done — your master is live (unlisted).")
    print("  Watch:  %s" % done["watchUrl"])
    print("  Manage: %s  (title, poster frame, visibility)" % done["manageUrl"])


if __name__ == "__main__":
    main()