#!/usr/bin/env python3
"""
Batch execute ComfyUI workflows on fal.ai serverless GPUs.

Usage:
    # Run a workflow once
    python batch.py workflow.json

    # Run with a different input video
    python batch.py workflow.json --input video.mp4

    # Run 10 times with different videos from a folder
    python batch.py workflow.json --input-folder ./videos/

    # Run 5 times in parallel
    python batch.py workflow.json --count 5 --parallel 3

Outputs are printed as URLs and saved to ./batch_outputs/
"""

import argparse
import json
import os
import sys
import time
import uuid
from pathlib import Path

try:
    import fal_client
except ImportError:
    print("Missing fal-client. Run: pip install fal-client")
    sys.exit(1)


def load_env():
    env_path = Path(__file__).parent / ".env"
    if env_path.exists():
        for line in env_path.read_text().splitlines():
            if "=" in line and not line.startswith("#"):
                k, v = line.split("=", 1)
                os.environ.setdefault(k.strip(), v.strip())
    if not os.environ.get("FAL_KEY"):
        print("ERROR: FAL_KEY not set. Run ./run.sh first or add to .env")
        sys.exit(1)


def load_workflow(path: str) -> dict:
    with open(path) as f:
        data = json.load(f)
    # Handle different ComfyUI export formats
    if any(isinstance(v, dict) and "class_type" in v for v in data.values()):
        return data
    if "prompt" in data:
        return data["prompt"]
    if "workflow" in data:
        return data["workflow"]
    return data


def find_input_nodes(prompt: dict) -> list:
    """Find all nodes that load files (video, image, audio)."""
    load_types = {
        "VHS_LoadVideo": "video",
        "VHS_LoadVideoPath": "video",
        "LoadImage": "image",
        "VHS_LoadAudio": "audio",
        "VHS_LoadAudioUpload": "audio",
        "LoadAudio": "audio",
    }
    nodes = []
    for node_id, node_data in prompt.items():
        class_type = node_data.get("class_type", "")
        if class_type in load_types:
            input_name = load_types[class_type]
            nodes.append({
                "node_id": node_id,
                "class_type": class_type,
                "input_name": input_name,
                "current_value": node_data.get("inputs", {}).get(input_name, ""),
            })
    return nodes


def upload_file(file_path: str) -> str:
    """Upload a file to fal.media and return the URL."""
    print(f"  Uploading {os.path.basename(file_path)}...")
    url = fal_client.upload_file(file_path)
    print(f"  Uploaded: {url}")
    return url


def prepare_fal_inputs(prompt: dict, input_file: str = None) -> tuple:
    """Prepare fal_inputs for file pre-download on the server."""
    fal_inputs = {}
    fal_inputs_dev_info = {}
    input_nodes = find_input_nodes(prompt)

    if not input_nodes:
        return fal_inputs, fal_inputs_dev_info

    if input_file:
        # Upload and replace the first input node
        url = upload_file(input_file)
        node = input_nodes[0]
        field = str(uuid.uuid4())
        fal_inputs[field] = url
        fal_inputs_dev_info[field] = {
            "key": [node["node_id"], "inputs", node["input_name"]],
            "class_type": node["class_type"],
        }
    else:
        # Check if current values are already URLs
        for node in input_nodes:
            val = node["current_value"]
            if val and str(val).startswith(("https://", "http://")):
                field = str(uuid.uuid4())
                fal_inputs[field] = val
                fal_inputs_dev_info[field] = {
                    "key": [node["node_id"], "inputs", node["input_name"]],
                    "class_type": node["class_type"],
                }

    return fal_inputs, fal_inputs_dev_info


def run_workflow(endpoint: str, prompt: dict, fal_inputs: dict, fal_inputs_dev_info: dict, run_id: int) -> dict:
    print(f"  [{run_id}] Running...")
    start = time.time()

    arguments = {"prompt": prompt}
    if fal_inputs:
        arguments["fal_inputs"] = fal_inputs
        arguments["fal_inputs_dev_info"] = fal_inputs_dev_info

    try:
        result = fal_client.subscribe(
            endpoint,
            arguments=arguments,
            with_logs=False,
        )
        elapsed = time.time() - start
        print(f"  [{run_id}] Done in {elapsed:.1f}s")
        return result
    except Exception as e:
        print(f"  [{run_id}] Error: {e}")
        return {"error": str(e)}


def extract_urls(result: dict) -> list:
    urls = []
    outputs = result.get("outputs", {})
    for key in ("images", "gifs"):
        for item in outputs.get(key, []):
            if "url" in item:
                urls.append(item["url"])
    return urls


def main():
    parser = argparse.ArgumentParser(description="Batch execute ComfyUI workflows")
    parser.add_argument("workflow", help="Path to workflow JSON file")
    parser.add_argument("--input", help="Input file to replace in the workflow (video, image)")
    parser.add_argument("--input-folder", help="Folder of input files — runs once per file")
    parser.add_argument("--count", type=int, default=1, help="Number of runs (default: 1)")
    parser.add_argument("--parallel", type=int, default=1, help="Max parallel runs (default: 1)")
    parser.add_argument("--output", default="batch_outputs", help="Output directory")
    args = parser.parse_args()

    load_env()
    endpoint = os.environ.get("FAL_COMFY_ENDPOINT", "comfyui/server")
    prompt = load_workflow(args.workflow)

    # Show detected input nodes
    input_nodes = find_input_nodes(prompt)
    if input_nodes:
        print(f"  Input nodes found:")
        for n in input_nodes:
            print(f"    - {n['class_type']} ({n['input_name']}): {n['current_value']}")
    print()

    os.makedirs(args.output, exist_ok=True)
    all_urls = []

    # Determine run list
    runs = []
    if args.input_folder:
        folder = Path(args.input_folder)
        for f in sorted(folder.iterdir()):
            if f.suffix.lower() in (".mp4", ".mov", ".webm", ".gif", ".png", ".jpg", ".jpeg", ".webp"):
                runs.append(str(f))
        print(f"  Found {len(runs)} files in {args.input_folder}")
    elif args.input:
        runs = [args.input] * args.count
    else:
        runs = [None] * args.count

    print(f"  Endpoint: {endpoint}")
    print(f"  Total runs: {len(runs)}")
    print()

    def execute_run(idx, input_file):
        fal_inputs, fal_inputs_dev_info = prepare_fal_inputs(prompt, input_file)
        result = run_workflow(endpoint, prompt, fal_inputs, fal_inputs_dev_info, idx + 1)
        urls = extract_urls(result)

        # Save result
        result_path = os.path.join(args.output, f"result_{idx+1:04d}.json")
        with open(result_path, "w") as f:
            json.dump(result, f, indent=2)

        return urls

    if args.parallel <= 1:
        for idx, input_file in enumerate(runs):
            urls = execute_run(idx, input_file)
            all_urls.extend(urls)
    else:
        import concurrent.futures
        with concurrent.futures.ThreadPoolExecutor(max_workers=args.parallel) as pool:
            futures = {pool.submit(execute_run, idx, f): idx for idx, f in enumerate(runs)}
            for future in concurrent.futures.as_completed(futures):
                urls = future.result()
                all_urls.extend(urls)

    print()
    print(f"  {len(all_urls)} output files:")
    for url in all_urls:
        print(f"    {url}")
    print()
    print(f"  Results saved to {args.output}/")


if __name__ == "__main__":
    main()
