#!/usr/bin/env python3
"""Upload local files into a ComfyUI instance's input folder.

Uses ComfyUI's built-in POST /upload/image endpoint. Prints one reference
name per uploaded file to stdout ("name.png" or "subfolder/name.png") --
exactly the string a LoadImage node's `image` input expects.

ComfyUI dedups colliding names by renaming to "name (1).png" unless
--overwrite is given, so ALWAYS use the printed name, not the local one.

Examples:
    python upload_files.py --url https://POD-7865.proxy.runpod.net head.jpg body.jpg
    python upload_files.py --url http://127.0.0.1:7865 --subfolder demo --overwrite *.png
"""
import argparse
import pathlib
import sys

import requests


def upload_file(base_url, path, subfolder="", overwrite=False, timeout=120):
    """Upload one file; returns the reference name to use in a workflow."""
    path = pathlib.Path(path)
    with path.open("rb") as f:
        r = requests.post(
            f"{base_url.rstrip('/')}/upload/image",
            files={"image": (path.name, f)},
            data={
                "subfolder": subfolder,
                "overwrite": "true" if overwrite else "false",
                "type": "input",
            },
            timeout=timeout,
        )
    r.raise_for_status()
    info = r.json()
    name = info["name"]
    if info.get("subfolder"):
        name = f"{info['subfolder']}/{name}"
    return name


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("--subfolder", default="", help="subfolder under ComfyUI/input/ to upload into")
    ap.add_argument("--overwrite", action="store_true", help="replace existing files instead of auto-renaming")
    ap.add_argument("files", nargs="+", help="local files to upload")
    args = ap.parse_args()

    failed = False
    for f in args.files:
        try:
            name = upload_file(args.url, f, args.subfolder, args.overwrite)
        except (OSError, requests.RequestException) as e:
            print(f"ERROR uploading {f}: {e}", file=sys.stderr)
            failed = True
            continue
        print(f"uploaded {f} -> input/{name}", file=sys.stderr)
        print(name)
    sys.exit(1 if failed else 0)


if __name__ == "__main__":
    main()