#!/usr/bin/env python3
"""Validate a ComfyUI API graph and optionally submit it to a running local server."""
import argparse
import json
from pathlib import Path
import sys
import urllib.error
import urllib.request
import uuid


def validate(graph, catalog):
    errors = []
    if not isinstance(graph, dict) or not graph:
        return ['Expected a nonempty ComfyUI API graph object.']
    for key, node in graph.items():
        if not isinstance(node, dict) or node.get('class_type') not in catalog:
            errors.append(f'Node {key}: unavailable class {node.get("class_type") if isinstance(node, dict) else "invalid node"}')
            continue
        spec = catalog[node['class_type']].get('input', {})
        for group in ['required', 'optional']:
            for name, entry in spec.get(group, {}).items():
                if name not in node.get('inputs', {}):
                    if group == 'required':
                        errors.append(f'Node {key}: missing required input {name}')
                    continue
                value = node['inputs'][name]
                # ComfyUI links are [node_id, output_slot], whereas enum inputs
                # are strings. Check links before validating allowed values.
                if isinstance(value, list):
                    if len(value) != 2 or str(value[0]) not in graph or not isinstance(value[1], int) or value[1] < 0:
                        errors.append(f'Node {key}: invalid link for {name}')
                    else:
                        linked = catalog.get(graph[str(value[0])].get('class_type'), {})
                        if value[1] >= len(linked.get('output', [])):
                            errors.append(f'Node {key}: unavailable output slot for {name}')
                    continue
                if isinstance(entry[0], list) and value not in entry[0]:
                    errors.append(f'Node {key}: {name}={value!r} is unavailable in this installation')
    return errors


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('workflow', type=Path)
    parser.add_argument('--url', default='http://127.0.0.1:8188')
    parser.add_argument('--validate-only', action='store_true')
    args = parser.parse_args()
    try:
        graph = json.loads(args.workflow.read_text(encoding='utf-8'))
        with urllib.request.urlopen(args.url.rstrip('/') + '/object_info', timeout=30) as r:
            catalog = json.load(r)
        errors = validate(graph, catalog)
        if errors:
            print('Compatibility check failed; nothing was queued.', file=sys.stderr)
            print('\n'.join(errors), file=sys.stderr)
            return 2
        print('Node classes, required inputs and enum values are available. GPU compatibility is not established.')
        if args.validate_only:
            return 0
        body = json.dumps({'prompt':graph, 'client_id':str(uuid.uuid4())}).encode()
        req = urllib.request.Request(args.url.rstrip('/')+'/prompt', data=body, headers={'Content-Type':'application/json'})
        with urllib.request.urlopen(req,timeout=30) as r:
            result = json.load(r)
        if not result.get('prompt_id') or result.get('node_errors'):
            print(json.dumps(result,indent=2), file=sys.stderr)
            return 3
        print('Queued prompt:',result['prompt_id'])
        print('Inspect ComfyUI history for completion; submission does not prove a successful render.')
        return 0
    except (OSError, ValueError, urllib.error.URLError) as exc:
        print(f'Unable to submit workflow: {exc}',file=sys.stderr)
        return 1


if __name__ == '__main__':
    raise SystemExit(main())