#!/usr/bin/env python3
"""Download files from a ComfyUI instance's output (or input/temp) folder.

Uses ComfyUI's built-in GET /view endpoint. File names may include a
subfolder ("qwen_edit/2026-07-07/img_00001_.webp") -- the same strings
run_workflow.py prints when a job finishes. Prints the local path of each
downloaded file to stdout.

Examples:
    python download_files.py --url https://POD-7865.proxy.runpod.net headswap_00001_.png
    python download_files.py --url http://127.0.0.1:7865 --out results a/x.png b/y.png
"""
import argparse
import pathlib
import sys

import requests


def download_file(base_url, name, folder_type="output", out_dir=".", timeout=300):
    """Download one file by its ComfyUI reference name; returns the local path."""
    ref = pathlib.PurePosixPath(name)
    r = requests.get(
        f"{base_url.rstrip('/')}/view",
        params={
            "filename": ref.name,
            "subfolder": str(ref.parent) if str(ref.parent) != "." else "",
            "type": folder_type,
        },
        timeout=timeout,
    )
    r.raise_for_status()
    dest = pathlib.Path(out_dir) / name
    dest.parent.mkdir(parents=True, exist_ok=True)
    dest.write_bytes(r.content)
    return dest


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--url", required=True, help="ComfyUI base URL, e.g. https://POD-7865.proxy.runpod.net")
    ap.add_argument("--type", default="output", choices=["output", "input", "temp"],
                    help="which ComfyUI folder to read from (default: output)")
    ap.add_argument("--out", default=".", help="local directory to save into (default: cwd)")
    ap.add_argument("names", nargs="+", help='file names, optionally with subfolder ("sub/file.png")')
    args = ap.parse_args()

    failed = False
    for name in args.names:
        try:
            dest = download_file(args.url, name, args.type, args.out)
        except (OSError, requests.RequestException) as e:
            print(f"ERROR downloading {name}: {e}", file=sys.stderr)
            failed = True
            continue
        print(f"downloaded {args.type}/{name} -> {dest} ({dest.stat().st_size} bytes)", file=sys.stderr)
        print(dest)
    sys.exit(1 if failed else 0)


if __name__ == "__main__":
    main()