#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Xeed Seedance 2.5 production helper. Nothing runs without an explicit command.
Standard-library only. prepare is offline; submit/status/download contact Ark.
Official field provenance is in README.txt and workflow.html.
"""
import argparse
import base64
import hashlib
import json
import os
from pathlib import Path
import re
import struct
import sys
import time
from datetime import datetime, timezone
from urllib.error import HTTPError, URLError
from urllib.parse import urlparse
from urllib.request import Request, urlopen

BASE = "https://ark.cn-beijing.volces.com/api/v3/contents/generations/tasks"
MODEL = "doubao-seedance-2-5-260628"
KNOWN_STATUS = {"queued", "running", "cancelled", "succeeded", "failed", "expired"}
HERE = Path(__file__).resolve().parent

def save_json(path, value):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(value, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")

def validate_task_id(value):
    if not re.fullmatch(r"[A-Za-z0-9_.-]+", value):
        raise ValueError("Task ID has unexpected characters.")
    return value

def auth_headers():
    key = os.environ.get("ARK_API_KEY")
    if not key:
        raise ValueError("Set ARK_API_KEY in the local environment; do not store it in project files.")
    return {"Authorization": "Bearer " + key, "Content-Type": "application/json"}

def api_json(method, url, body=None):
    data = None if body is None else json.dumps(body).encode("utf-8")
    request = Request(url, data=data, headers=auth_headers(), method=method)
    try:
        with urlopen(request, timeout=90) as response:
            result = json.load(response)
        if not isinstance(result, dict):
            raise ValueError("Unexpected response type; inspect the provider response privately.")
        return result
    except HTTPError as error:
        raise RuntimeError("Ark HTTP " + str(error.code) + ". No automatic retry; inspect the task before resubmitting.") from None
    except URLError:
        raise RuntimeError("Network result is ambiguous. No automatic POST retry. Check Ark for the task before resubmitting.") from None

def prepare(args):
    image = Path(args.image)
    data = image.read_bytes()
    if not data.startswith(b"\x89PNG\r\n\x1a\n"):
        raise ValueError("This pilot helper accepts PNG first frames only; use the official API for other documented formats.")
    if len(data) >= 30000000:
        raise ValueError("PNG exceeds the conservative 30 MB input limit.")
    width, height = struct.unpack(">II", data[16:24])
    if not (300 <= width <= 6000 and 300 <= height <= 6000 and 0.4 <= width / height <= 2.5):
        raise ValueError("Image dimensions or aspect are outside the documented input range.")
    prompts = json.loads((HERE / "prompts.json").read_text(encoding="utf-8"))
    variants = [prompts["primary_pilot"]] + prompts["modules"]
    matched = [item for item in variants if item["id"] == args.module]
    if not matched:
        raise ValueError("Unknown prompt module.")
    frame = "data:image/png;base64," + base64.b64encode(data).decode("ascii")
    body = {
        "model": MODEL,
        "content": [
            {"type": "text", "text": matched[0]["motion_prompt"]},
            {"type": "image_url", "image_url": {"url": frame}, "role": "first_frame"},
            {"type": "image_url", "image_url": {"url": frame}, "role": "last_frame"}
        ],
        "ratio": "adaptive", "duration": 5, "resolution": "720p",
        "generate_audio": False, "output_format": "mp4"
    }
    target = Path(args.output)
    if target.exists():
        raise ValueError("Output request already exists; choose a new revision.")
    save_json(target, body)
    meta = {
        "offline_prepared": True, "submitted": False, "model": MODEL,
        "module": args.module, "source_file": str(image.resolve()),
        "source_sha256": hashlib.sha256(data).hexdigest(),
        "source_width": width, "source_height": height,
        "source_rights_review": "unreviewed",
        "note": "Request contains local image pixels as data URI. Keep it private; do not print or commit it."
    }
    save_json(target.with_suffix(".metadata.json"), meta)
    print("Prepared offline request: " + str(target))
    print("No network call, upload or generation performed.")

def submit(args):
    body = json.loads(Path(args.request).read_text(encoding="utf-8"))
    if body.get("model") != MODEL or body.get("ratio") != "adaptive":
        raise ValueError("Use the reviewed Seedance 2.5 first/last-frame request.")
    if "REPLACE_WITH_" in json.dumps(body):
        raise ValueError("Request still contains placeholders.")
    if body.get("duration") != 5:
        raise ValueError("This pilot helper expects the documented five-second trial.")
    request_hash = hashlib.sha256(json.dumps(body, sort_keys=True, separators=(",", ":")).encode("utf-8")).hexdigest()
    marker = Path(args.job_dir) / ("submit-" + request_hash + ".attempt.private.json")
    if marker.exists():
        raise ValueError("This request was already attempted. Inspect its attempt marker and saved task ID; do not submit it twice.")
    attempt = {"request_sha256": request_hash, "request_file": str(Path(args.request).resolve()), "attempted_at": datetime.now(timezone.utc).isoformat(), "state": "attempting", "task_id": None}
    save_json(marker, attempt)
    try:
        result = api_json("POST", BASE, body)
    except Exception:
        attempt["state"] = "ambiguous_or_failed_inspect_before_resubmit"
        save_json(marker, attempt)
        raise
    task_id = result.get("id")
    if not isinstance(task_id, str) or not task_id:
        save_json(Path(args.job_dir) / "unrecognized-create-response.private.json", result)
        raise ValueError("Create response has no usable id. It was saved; do not resubmit automatically.")
    validate_task_id(task_id)
    save_json(Path(args.job_dir) / (task_id + ".create.private.json"), result)
    attempt["state"] = "created"
    attempt["task_id"] = task_id
    save_json(marker, attempt)
    print("Task created and ID saved: " + task_id)
    print("This command submits one generation; it does not auto-resubmit.")

def get_status(task_id, job_dir):
    task_id = validate_task_id(task_id)
    result = api_json("GET", BASE + "/" + task_id)
    save_json(Path(job_dir) / (task_id + ".status.private.json"), result)
    status = result.get("status")
    if status not in KNOWN_STATUS:
        raise ValueError("Unknown status; private response snapshot saved. Stop and inspect official documentation.")
    return result

def status(args):
    result = get_status(args.task_id, args.job_dir)
    value = result["status"]
    print("Task status: " + value)
    if value in {"failed", "cancelled", "expired"}:
        raise RuntimeError("Terminal non-success: " + value + ". Inspect the private snapshot/error; no automatic retry.")
    elif value == "succeeded":
        print("Ready for download. The signed result URL is kept in the private status snapshot.")

def poll(args):
    if args.interval < 1 or args.timeout < 1:
        raise ValueError("Polling interval and timeout must be positive.")
    started = time.monotonic()
    last = None
    while True:
        result = get_status(args.task_id, args.job_dir)
        value = result["status"]
        if value != last:
            print("Task status: " + value, flush=True)
            last = value
        if value == "succeeded":
            print("Completed. Download using the same saved task ID.")
            return
        if value in {"failed", "cancelled", "expired"}:
            raise RuntimeError("Terminal non-success: " + value + ". Inspect the private snapshot/error.")
        remaining = args.timeout - (time.monotonic() - started)
        if remaining <= 0:
            raise RuntimeError("Polling timeout; resume with the same task ID. No POST was made.")
        time.sleep(min(args.interval, remaining))

def download(args):
    target = Path(args.output)
    if target.exists():
        raise ValueError("Output master already exists; choose a new revision.")
    result = get_status(args.task_id, args.job_dir)
    if result["status"] != "succeeded":
        raise ValueError("Task is not succeeded; no file downloaded.")
    content = result.get("content")
    video_url = content.get("video_url") if isinstance(content, dict) else None
    if not isinstance(video_url, str):
        raise ValueError("Successful task has no content.video_url; inspect the saved response.")
    parsed = urlparse(video_url)
    if parsed.scheme != "https" or not parsed.netloc:
        raise ValueError("Refuse empty or non-HTTPS download URL.")
    target.parent.mkdir(parents=True, exist_ok=True)
    partial = target.with_name(target.name + ".part")
    if partial.exists():
        raise ValueError("Partial file already exists; inspect it before choosing another output revision.")
    digest = hashlib.sha256()
    try:
        # Signed download endpoint receives no ARK_API_KEY/Authorization header.
        with urlopen(Request(video_url, method="GET"), timeout=90) as response, partial.open("xb") as out:
            while True:
                block = response.read(1024 * 1024)
                if not block:
                    break
                digest.update(block)
                out.write(block)
        if partial.stat().st_size == 0:
            raise ValueError("Downloaded file is empty.")
        with partial.open("rb") as header_file:
            header = header_file.read(12)
        if len(header) < 12 or header[4:8] != b"ftyp":
            raise ValueError("Downloaded content is not an MP4/ISO-BMFF file; inspect the retained .part file.")
        partial.rename(target)
    except Exception:
        raise RuntimeError("Download incomplete. Private status snapshot and any .part file are retained for inspection.") from None
    save_json(target.with_suffix(".download.json"), {
        "task_id": args.task_id, "file": str(target.resolve()),
        "bytes": target.stat().st_size, "sha256": digest.hexdigest(),
        "url_record": "in private status snapshot; not copied into public manifest",
        "quality_review_passed": False
    })
    print("Master saved: " + str(target))
    print("Download is complete; root, loop, codec and visual QA remain pending.")

def main():
    parser = argparse.ArgumentParser(description=__doc__)
    sub = parser.add_subparsers(dest="action", required=True)
    p = sub.add_parser("prepare", help="Offline: embed a local PNG and selected prompt.")
    p.add_argument("--image", required=True)
    p.add_argument("--module", default="flower-pilot")
    p.add_argument("--output", required=True)
    p.set_defaults(run=prepare)
    p = sub.add_parser("submit", help="Contact Ark once: submit a paid generation request.")
    p.add_argument("--request", required=True)
    p.add_argument("--job-dir", required=True)
    p.set_defaults(run=submit)
    p = sub.add_parser("status", help="Contact Ark once: GET saved task status.")
    p.add_argument("--task-id", required=True)
    p.add_argument("--job-dir", required=True)
    p.set_defaults(run=status)
    p = sub.add_parser("poll", help="Bounded GET polling; resumes by task ID and never POSTs.")
    p.add_argument("--task-id", required=True)
    p.add_argument("--job-dir", required=True)
    p.add_argument("--interval", type=float, default=10)
    p.add_argument("--timeout", type=float, default=600)
    p.set_defaults(run=poll)
    p = sub.add_parser("download", help="GET status, then download successful result without key header.")
    p.add_argument("--task-id", required=True)
    p.add_argument("--job-dir", required=True)
    p.add_argument("--output", required=True)
    p.set_defaults(run=download)
    args = parser.parse_args()
    try:
        args.run(args)
    except (ValueError, OSError, RuntimeError, json.JSONDecodeError) as error:
        parser.exit(1, "Error: " + str(error) + "\n")

if __name__ == "__main__":
    main()