#!/usr/bin/env python3
"""Convert one 2D video to side-by-side 3D with the Owl3D API.

    export OWL3D_API_KEY=o3d_live_...
    python3 convert.py my-video.mp4

Needs Python 3.9+, `pip install requests`, and ffprobe (part of FFmpeg) on PATH.
"""
import hashlib, json, os, subprocess, sys, time
import requests

API = os.environ.get("OWL3D_API_URL", "https://api.owl3d.com").rstrip("/")
KEY = os.environ["OWL3D_API_KEY"]
EDGE = 4 * 1024 * 1024  # the fingerprint hashes the first and the last 4 MiB


def call(method, path, body=None):
    """One API call. Waits and tries again on 429 (Retry-After) and on 5xx."""
    delay = 1
    while True:
        r = requests.request(method, API + path, json=body, timeout=60,
                             headers={"Authorization": f"Bearer {KEY}"})
        if r.status_code == 429:
            time.sleep(int(r.headers.get("Retry-After", "15")))
            continue
        if r.status_code >= 500:
            time.sleep(delay)
            delay = min(delay * 2, 60)
            continue
        if r.status_code >= 400:
            sys.exit(f"{method} {path}: {r.status_code} {r.text}")
        return r.json()


def fingerprint(path):
    size = os.path.getsize(path)
    edge = min(EDGE, size)
    with open(path, "rb") as f:
        head = hashlib.sha256(f.read(edge)).hexdigest()
        f.seek(size - edge)
        tail = hashlib.sha256(f.read(edge)).hexdigest()
    return {"size": size, "headSha256": head, "tailSha256": tail, "edgeBytes": edge}


def probe(path):
    out = subprocess.run(["ffprobe", "-v", "error", "-print_format", "json",
                          "-show_streams", "-show_format", path],
                         check=True, capture_output=True, text=True).stdout
    p = json.loads(out)
    v = next(s for s in p["streams"] if s["codec_type"] == "video")
    n, d = (v.get("avg_frame_rate") or v["r_frame_rate"]).split("/")
    codec = {"h264": "H.264", "hevc": "HEVC", "av1": "AV1", "vp9": "VP9"}.get(v["codec_name"], v["codec_name"])
    return {
        "durationSeconds": float(p["format"]["duration"]),
        "width": v["width"], "height": v["height"],
        "fps": int(n) / int(d) if int(d) else int(n),
        "codec": codec,
        "bitDepth": 10 if "10" in v.get("pix_fmt", "") else 8,
        "chroma": "4:2:0",
    }


def main(path):
    fp = fingerprint(path)
    facts = {"name": os.path.basename(path), "size": fp["size"], **probe(path), "fingerprint": fp}

    # 1. Estimate: the price in credits. Nothing is charged.
    est = call("POST", "/v1/estimates", facts)
    if not est.get("ok"):
        sys.exit(f"estimate refused: {est}")
    print(f"price: {est['credits']['total']} credits")

    # 2. Create the job: reserves the credits, returns the upload URLs.
    job = call("POST", "/v1/jobs", {"estimateId": est["id"]})
    job_id, up = job["id"], job["upload"]

    # 3. Upload each part straight to storage, keep the ETags, then complete.
    etags = []
    with open(path, "rb") as f:
        for url in up["parts"]:
            r = requests.put(url, data=f.read(up["partSize"]), timeout=600)
            r.raise_for_status()
            etags.append(r.headers["ETag"])
    call("POST", f"/v1/jobs/{job_id}/upload-complete", {"etags": etags})
    print(f"job {job_id} queued")

    # 4. Poll every 15 s until the job ends.
    while True:
        st = call("GET", f"/v1/jobs/{job_id}")
        print(f"  {st['status']:<10} {round(st.get('progress', 0) * 100)}%")
        if st["status"] in ("done", "failed", "cancelled", "insufficient_credits"):
            break
        time.sleep(15)
    if st["status"] == "insufficient_credits":
        sys.exit(f"add {st['shortBy']} credits, then POST {st['resumeUrl']} before {st['deleteAt']}")
    if st["status"] != "done":
        sys.exit(f"job ended: {st['status']} {st.get('error')}")

    # 5. Download: the link works for 15 minutes.
    res = call("GET", f"/v1/jobs/{job_id}/result")
    with requests.get(res["url"], stream=True, timeout=600) as r:
        r.raise_for_status()
        with open(res["filename"], "wb") as out:
            for chunk in r.iter_content(1 << 20):
                out.write(chunk)
    print(f"saved {res['filename']}")


if __name__ == "__main__":
    if len(sys.argv) != 2:
        sys.exit("usage: python3 convert.py <video>")
    main(sys.argv[1])