#!/usr/bin/env python3
"""
Submit each video segment to Replicate's heygen/video-translate model in parallel,
poll until each prediction completes, and download the translated output.

Usage:
    python3 translate.py
"""

from __future__ import annotations

import json
import os
import pathlib
import sys
import time
import urllib.request
import urllib.error

ROOT = pathlib.Path(__file__).resolve().parent


def _path_env(var: str, default: pathlib.Path) -> pathlib.Path:
    raw = os.environ.get(var)
    if not raw:
        return default
    p = pathlib.Path(raw)
    return p if p.is_absolute() else ROOT / p


SEG_DIR = _path_env("SEG_DIR", ROOT / "segments")
OUT_DIR = _path_env("OUT_DIR", ROOT / "translated")
STATE_FILE = _path_env("STATE_FILE", ROOT / "predictions.json")

# replicate_upload: POST files to Replicate (works when ngrok fetch times out).
# url: serve segments yourself and set NGROK_URL (or any HTTPS base).
INPUT_DELIVERY = os.environ.get("INPUT_DELIVERY", "replicate_upload")
NGROK_URL = os.environ.get("NGROK_URL", "").rstrip("/")
TOKEN = os.environ["REPLICATE_API_TOKEN"]
LANGUAGE = os.environ.get("OUTPUT_LANGUAGE", "English")
MODE = os.environ.get("MODE", "precision")

API = "https://api.replicate.com/v1"
HEADERS = {
    "Authorization": f"Bearer {TOKEN}",
    "Content-Type": "application/json",
    "User-Agent": "translate-script/1.0",
}


def upload_replicate_file(path: pathlib.Path) -> str:
    """Upload a local file to Replicate; return the authenticated file URL for model input."""
    import subprocess

    abs_path = path.resolve()
    proc = subprocess.run(
        [
            "curl",
            "-s",
            "-S",
            "-X",
            "POST",
            f"{API}/files",
            "-H",
            f"Authorization: Token {TOKEN}",
            "-F",
            f"content=@{abs_path};type=video/mp4;filename={path.name}",
            "-F",
            "metadata={};type=application/json",
        ],
        capture_output=True,
        text=True,
        timeout=1200,
        check=False,
    )
    if proc.returncode != 0:
        raise RuntimeError(f"upload curl rc={proc.returncode} stderr={proc.stderr!r}")
    try:
        data = json.loads(proc.stdout)
    except json.JSONDecodeError as exc:
        raise RuntimeError(f"upload bad JSON: {proc.stdout[:800]}") from exc
    if proc.stderr:
        print(proc.stderr, file=sys.stderr)
    url = (data.get("urls") or {}).get("get")
    if not url:
        raise RuntimeError(f"upload missing urls.get: {proc.stdout[:800]}")
    return url


def resolve_video_url(seg: pathlib.Path, existing: dict | None) -> str:
    if INPUT_DELIVERY == "url":
        if not NGROK_URL:
            print("NGROK_URL is required when INPUT_DELIVERY=url", file=sys.stderr)
            sys.exit(1)
        return f"{NGROK_URL}/{seg.name}"
    cached = (existing or {}).get("input_video_url")
    if cached:
        return cached
    print(f"[{seg.name}] uploading to Replicate ({seg.stat().st_size // (1024 * 1024)} MiB)...")
    return upload_replicate_file(seg)


def http(method: str, url: str, body: dict | None = None) -> dict:
    data = json.dumps(body).encode() if body is not None else None
    req = urllib.request.Request(url, data=data, method=method, headers=HEADERS)
    try:
        with urllib.request.urlopen(req, timeout=60) as resp:
            return json.loads(resp.read())
    except urllib.error.HTTPError as e:
        raise RuntimeError(f"HTTP {e.code} on {method} {url}: {e.read().decode()}") from None


def submit(seg: pathlib.Path, video_url: str) -> dict:
    payload = {
        "input": {
            "video": video_url,
            "output_language": LANGUAGE,
            "mode": MODE,
        }
    }
    return http(
        "POST",
        f"{API}/models/heygen/video-translate/predictions",
        body=payload,
    )


def poll(prediction_id: str) -> dict:
    return http("GET", f"{API}/predictions/{prediction_id}")


def download(url: str, dest: pathlib.Path) -> None:
    req = urllib.request.Request(url, headers={"User-Agent": HEADERS["User-Agent"]})
    with urllib.request.urlopen(req, timeout=900) as resp, dest.open("wb") as f:
        while True:
            chunk = resp.read(1 << 20)
            if not chunk:
                break
            f.write(chunk)


def load_state() -> dict:
    if STATE_FILE.exists():
        return json.loads(STATE_FILE.read_text())
    return {}


def save_state(state: dict) -> None:
    STATE_FILE.write_text(json.dumps(state, indent=2))


def main() -> int:
    OUT_DIR.mkdir(exist_ok=True)
    segments = sorted(SEG_DIR.glob("seg_*.mp4"))
    if not segments:
        print("No segments found.", file=sys.stderr)
        return 1

    state = load_state()

    print(f"Language: {LANGUAGE}   Mode: {MODE}   Delivery: {INPUT_DELIVERY}   Segments: {len(segments)}")
    if INPUT_DELIVERY == "url":
        print(f"Public base: {NGROK_URL}\n")
    else:
        print("Using Replicate file uploads (no public URL needed).\n")

    for seg in segments:
        if seg.name in state and state[seg.name].get("status") in ("starting", "processing", "succeeded"):
            print(f"[{seg.name}] reusing existing prediction {state[seg.name]['id']}")
            continue
        video_url = resolve_video_url(seg, state.get(seg.name))
        pred = submit(seg, video_url)
        state[seg.name] = {
            "id": pred["id"],
            "status": pred["status"],
            "submitted_at": pred.get("created_at"),
            "input_video_url": video_url,
        }
        save_state(state)
        print(f"[{seg.name}] submitted -> {pred['id']} ({pred['status']})")

    print("\nPolling all predictions until done...\n")
    pending = {name for name, info in state.items() if info["status"] not in ("succeeded", "failed", "canceled")}
    started = time.time()
    while pending:
        for name in list(pending):
            info = state[name]
            try:
                pred = poll(info["id"])
            except Exception as exc:
                print(f"[{name}] poll error: {exc}")
                continue
            if pred["status"] != info["status"]:
                print(f"[{name}] {info['status']} -> {pred['status']}  (elapsed {int(time.time() - started)}s)")
            info["status"] = pred["status"]
            if pred["status"] == "succeeded":
                info["output"] = pred["output"]
                pending.discard(name)
            elif pred["status"] in ("failed", "canceled"):
                info["error"] = pred.get("error") or "(no error message)"
                info["logs"] = pred.get("logs", "")[-2000:]
                pending.discard(name)
        save_state(state)
        if pending:
            time.sleep(15)

    print("\nDownloading outputs...\n")
    failures = []
    for name, info in state.items():
        if info["status"] != "succeeded":
            failures.append((name, info.get("error", "unknown")))
            continue
        out_path = OUT_DIR / name
        if out_path.exists() and out_path.stat().st_size > 0:
            print(f"[{name}] already downloaded ({out_path.stat().st_size} bytes)")
            continue
        print(f"[{name}] downloading {info['output']}")
        download(info["output"], out_path)
        print(f"[{name}] saved -> {out_path}  ({out_path.stat().st_size} bytes)")

    if failures:
        print("\nFAILURES:")
        for name, err in failures:
            print(f"  - {name}: {err}")
        return 2

    print("\nAll segments translated and downloaded.")
    return 0


if __name__ == "__main__":
    sys.exit(main())