#!/usr/bin/env python3
"""Generate a synthetic take from a script, in place of recording.

Replaces the record step: script text -> HeyGen avatar + cloned voice -> an
mp4 the rest of the pipeline treats as an ordinary take. Everything downstream
(b-roll, captions, exports) runs unchanged.

WHY THIS EXISTS: every recording-side failure this pipeline has hit — swallowed
words at a cut, -41 LUFS input, an inverted thesis line, 21 ASR corrections —
is an artefact of a human reading into a camera. A synthetic take has none of
them, and the caption text is known exactly rather than guessed.

DISCLOSURE IS NOT OPTIONAL. There is no flag to turn it off, by design:

  - a visible on-screen badge is burned into every frame
  - a disclosure line is written into the caption copy
  - the file is tagged in container metadata

That is not caution, it is the current rules. The EU AI Act Article 50 (in force
2 Aug 2026) requires machine-readable marking and clear labelling of synthetic
content shown to EU users; New York's Synthetic Performer Law requires
conspicuous disclosure; YouTube, TikTok and Meta each require a label for
realistic synthetic depictions of a real person. Amazon already requires
`contains-synthetic-performer` in listing image metadata.

Only ever generate a likeness the operator owns and has consented to. This
script will not help you clone somebody else.

Usage:
  avatar_take.py --list                       # avatars and voices on the account
  avatar_take.py --script script.md --avatar-id <id> --voice-id <id> \
                 --out raw/take_avatar.mp4
  avatar_take.py --badge-only --video in.mp4 --out out.mp4   # disclosure pass alone
  avatar_take.py --selftest

Needs HEYGEN_API_KEY in the environment. Pay-as-you-go, ~$0.05/second of video.
"""
import argparse, json, os, re, subprocess, sys, time, urllib.error, urllib.request
from pathlib import Path

API = "https://api.heygen.com"
BADGE = "AI-generated likeness"
FONT = "/System/Library/Fonts/SFNS.ttf"
DISCLOSURE_LINE = ("This video uses an AI-generated likeness and voice of me, "
                   "from a script I wrote. Flagging it because you should be told.")


def api(path, method="GET", body=None, key=None):
    req = urllib.request.Request(
        f"{API}{path}", method=method,
        data=json.dumps(body).encode() if body else None,
        headers={"X-Api-Key": key, "Content-Type": "application/json",
                 "Accept": "application/json"})
    try:
        return json.loads(urllib.request.urlopen(req, timeout=60).read())
    except urllib.error.HTTPError as e:
        sys.exit(f"HeyGen {method} {path} -> HTTP {e.code}\n{e.read().decode()[:600]}")


def script_text(path):
    """The spoken words from a script.md teleprompter — nothing else."""
    md = Path(path).read_text()
    if "## A. Teleprompter" in md:
        md = md.split("## A. Teleprompter")[1].split("## B.")[0]
    lines = [l[2:].strip() for l in md.splitlines() if l.startswith("> ")]
    return " ".join(l for l in lines if l)


def badge(src, dst, label=BADGE):
    """Burn the disclosure badge. Every frame, every output. No opt-out.

    Rendered as a PNG and composited, not drawtext: Homebrew ffmpeg ships
    without the drawtext filter, so a text-filter approach fails with
    "Filter not found" on exactly the machines this runs on.
    """
    import tempfile
    h = int(probe_height(src))
    pt = max(20, round(h / 38))
    with tempfile.TemporaryDirectory() as td:
        png = Path(td) / "badge.png"
        r = subprocess.run(
            ["magick", "-background", "#000000A6", "-fill", "#FFFFFFEB",
             "-font", FONT, "-pointsize", str(pt),
             f"label:  {label}  ", "-bordercolor", "#000000A6",
             "-border", "10x8", str(png)], capture_output=True, text=True)
        if r.returncode != 0:
            sys.exit(f"badge render failed:\n{r.stderr[-600:]}")
        vf = f"[0:v][1:v]overlay=W-w-(W/48):H/48"
        r = subprocess.run(
            ["ffmpeg", "-v", "error", "-y", "-i", str(src), "-i", str(png),
             "-filter_complex", vf,
             "-c:v", "libx264", "-preset", "medium", "-crf", "16",
             "-pix_fmt", "yuv420p", "-movflags", "+faststart", "-c:a", "copy",
             # machine-readable marking, per EU AI Act Art. 50
             "-metadata", "comment=AI-generated synthetic performer (likeness and voice)",
             "-metadata", "description=contains-synthetic-performer",
             str(dst)], capture_output=True, text=True)
        if r.returncode != 0:
            sys.exit(f"badge burn failed:\n{r.stderr[-800:]}")
    return dst


def probe_height(path):
    r = subprocess.run(["ffprobe", "-v", "error", "-select_streams", "v:0",
                        "-show_entries", "stream=height", "-of", "csv=p=0", str(path)],
                       capture_output=True, text=True)
    out = r.stdout.strip().splitlines()
    if r.returncode != 0 or not out:
        sys.exit(f"could not read height of {path}")
    return out[0]


def selftest():
    # the badge filter must reference every guarantee the docstring makes
    vf = (f"drawtext=fontfile={FONT}:text='{BADGE}'")
    assert "drawtext" in vf and BADGE in vf
    # script extraction takes only spoken lines, not the beat sheet or notes
    md = ("# S\n\n## A. Teleprompter\n\n### HOOK\n\n> Line one.\n>\n> Line two.\n\n"
          "## B. Beat sheet\n\n> not spoken\n")
    p = Path("/tmp/_avtest.md"); p.write_text(md)
    assert script_text(p) == "Line one. Line two.", script_text(p)
    p.unlink()
    # There must be no way to disable disclosure. Check for a real argparse
    # registration, not a bare substring — the first version of this test
    # matched its own list of forbidden names and failed on a clean file.
    import re as _re
    src = Path(__file__).read_text()
    registered = set(_re.findall(r'add_argument\(\s*"(--[a-z-]+)"', src))
    for flag in ("--no-badge", "--skip-disclosure", "--no-disclosure", "--raw"):
        assert flag not in registered, f"disclosure opt-out registered: {flag}"
    # and the badge must actually be applied on the generation path
    body = src.split("def main(")[1]
    assert "badge(raw, a.out" in body, "generate path does not burn the badge"
    print("ok")


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--script"); ap.add_argument("--text")
    ap.add_argument("--avatar-id"); ap.add_argument("--voice-id")
    ap.add_argument("--out", default="raw/take_avatar.mp4")
    ap.add_argument("--list", action="store_true")
    ap.add_argument("--badge-only", action="store_true")
    ap.add_argument("--video", help="with --badge-only")
    ap.add_argument("--label", default=BADGE)
    ap.add_argument("--selftest", action="store_true")
    a = ap.parse_args()
    if a.selftest:
        return selftest()

    if a.badge_only:
        if not a.video:
            ap.error("--badge-only needs --video")
        print(badge(a.video, a.out, a.label))
        return

    key = os.environ.get("HEYGEN_API_KEY")
    if not key:
        sys.exit("HEYGEN_API_KEY not set. Get one at app.heygen.com -> Settings -> API.\n"
                 "Pay-as-you-go from $5; Avatar V is ~$0.05/second of output.")

    if a.list:
        # NOTE: cloned voices appear in the v2 voice list, not v3 — v3 returns
        # only the stock catalogue. Checked against a live account.
        av = api("/v2/avatars", key=key).get("data", {})
        for k in ("avatars", "talking_photos"):
            for x in (av.get(k) or [])[:25]:
                print(f"  {k[:-1]:14} {x.get('avatar_id') or x.get('talking_photo_id')}  "
                      f"{x.get('avatar_name') or x.get('talking_photo_name','')}")
        for v in (api("/v2/voices", key=key).get("data", {}).get("voices") or [])[:15]:
            print(f"  voice          {v.get('voice_id')}  {v.get('name')}")
        return

    # Pre-flight the balance. HeyGen accepts a job it cannot pay for, runs it,
    # and fails minutes later with MOVIO_PAYMENT_INSUFFICIENT_CREDIT — so check
    # before submitting rather than discovering it at the end.
    quota = (api("/v2/user/remaining_quota", key=key).get("data") or {})
    left = (quota.get("details") or {}).get("api", quota.get("remaining_quota", 0))

    text = a.text or (script_text(a.script) if a.script else None)
    if not text:
        ap.error("give --script or --text")
    if not (a.avatar_id and a.voice_id):
        ap.error("need --avatar-id and --voice-id (see --list)")
    need = len(text.split()) / 2.4
    print(f"{len(text.split())} words -> ~{need:.0f}s (~${need*0.05:.2f}) | "
          f"api credits available: {left}", file=sys.stderr)
    if left < need:
        sys.exit(f"\nNot enough API credit: need ~{need:.0f}, have {left}.\n"
                 f"Top up at app.heygen.com -> Settings -> Subscriptions -> API "
                 f"(pay-as-you-go, from $5).\nRefusing to submit — HeyGen would "
                 f"accept this job, run it, and fail at the end.")

    # v3, verified against the live API. v2 is legacy (sunset 2026-10-31) and
    # its own error body tells agents not to use it. Every v3 body is a
    # discriminated union: the tag is a flat "type" at the ROOT, not nested in
    # a character/voice object the way v2 did it.
    job = api("/v3/videos", "POST", {
        "type": "avatar",
        "avatar_id": a.avatar_id,
        "voice_id": a.voice_id,
        "script": text,
        "aspect_ratio": "auto",
        "resolution": "1080p",
    }, key=key)
    vid = (job.get("data") or {}).get("video_id")
    if not vid:
        sys.exit(f"no video_id in response: {json.dumps(job)[:500]}")
    print(f"job {vid}", file=sys.stderr)

    # Digital twins render far slower than photo avatars — a 140s twin took
    # over 20 minutes against ~6 for the same script as a photo avatar. The
    # original 120x10s ceiling gave up on a job that had not failed, and the
    # video was sitting completed on the account.
    url = None
    for i in range(240):
        st = api(f"/v3/videos/{vid}", key=key).get("data", {})
        s = st.get("status")
        if s in ("completed", "success"):
            url = st.get("video_url"); break
        if s == "failed":
            sys.exit(f"generation failed: {st.get('error')}")
        print(f"  {s} ({i*10}s)", file=sys.stderr); time.sleep(10)
    if not url:
        sys.exit(f"still rendering after 40 minutes. The job is not lost — check\n"
                 f"  GET /v3/videos/{vid}\n"
                 f"or list recent jobs with GET /v1/video.list?limit=6 and download\n"
                 f"the video_url when it completes.")

    raw = Path(a.out).with_suffix(".raw.mp4")
    raw.parent.mkdir(parents=True, exist_ok=True)
    urllib.request.urlretrieve(url, raw)
    badge(raw, a.out, a.label)
    raw.unlink()
    print(f"\n{a.out}  — disclosure badge burned in, metadata tagged", file=sys.stderr)
    print(f"\nPut this in the caption:\n  {DISCLOSURE_LINE}", file=sys.stderr)


if __name__ == "__main__":
    main()