"""Render one API graph on your ComfyUI server and download its result."""
import argparse
import json
from pathlib import Path
import time
import urllib.parse
import urllib.request
import uuid


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('workflow', type=Path)
    parser.add_argument('--server', default='http://127.0.0.1:8188')
    parser.add_argument('--output', required=True, type=Path)
    parser.add_argument('--output-node', help='SaveVideo node ID to download for multi-output graphs')
    args = parser.parse_args()
    if args.output.exists():
        parser.error('output already exists; choose a new filename')
    graph = json.loads(args.workflow.read_text())
    for node in graph.values():
        if node['class_type'] == 'SaveVideo':
            node['inputs']['filename_prefix'] += '_' + uuid.uuid4().hex[:12]
    def api(path, payload=None):
        data = None if payload is None else json.dumps(payload).encode()
        req = urllib.request.Request(args.server.rstrip('/') + path, data=data,
                                     headers={'Content-Type': 'application/json'})
        with urllib.request.urlopen(req, timeout=60) as response:
            return json.load(response)
    result = api('/prompt', {'prompt': graph, 'client_id': 'fight_quad_example'})
    if result.get('node_errors') or 'prompt_id' not in result:
        raise RuntimeError(result)
    pid = result['prompt_id']
    print('Queued', pid, flush=True)
    deadline = time.monotonic() + 7200
    while time.monotonic() < deadline:
        history = api('/history/' + pid)
        if pid not in history:
            time.sleep(5)
            continue
        record = history[pid]
        if record['status'].get('status_str') != 'success':
            raise RuntimeError(record['status'])
        for key, node in graph.items():
            if node['class_type'] != 'SaveVideo' or (args.output_node and key != args.output_node):
                continue
            files = record['outputs'].get(key, {}).get('images', [])
            for item in files:
                if item.get('type') != 'output':
                    continue
                query = urllib.parse.urlencode({k: item[k] for k in ('filename', 'subfolder', 'type')})
                with urllib.request.urlopen(args.server.rstrip('/') + '/view?' + query, timeout=120) as response:
                    with args.output.open('xb') as output:
                        while chunk := response.read(1024 * 1024):
                            output.write(chunk)
                print('Saved', args.output)
                return
        raise RuntimeError('Completed without a SaveVideo output')
    raise TimeoutError('Render still pending after two hours; check ComfyUI before resubmitting')


if __name__ == '__main__':
    main()