import urllib.request
import re
import json
from PIL import Image
from io import BytesIO
import urllib.parse
import os

url = 'https://comfyanonymous.github.io/ComfyUI_examples/sdxl/'

try:
    req = urllib.request.Request(url, headers={'User-Agent': 'Mozilla/5.0'})
    with urllib.request.urlopen(req) as response:
        html = response.read().decode('utf-8')
    
    png_links = re.findall(r'src=["\']([^"\']+\.png)["\']', html)
    print('Found PNG links:', png_links)

    for link in png_links:
        full_url = urllib.parse.urljoin(url, link)
        print(f"Downloading {full_url}")
        
        req_img = urllib.request.Request(full_url, headers={'User-Agent': 'Mozilla/5.0'})
        with urllib.request.urlopen(req_img) as response_img:
            img_data = response_img.read()
        
        img = Image.open(BytesIO(img_data))
        if 'prompt' in img.info:
            api_json = json.loads(img.info['prompt'])
            filename = os.path.basename(link).replace('.png', '_api.json')
            
            # save to workspace root temporarily
            out_path = os.path.join("c:\\Users\\11304\\gitlab\\cybertogether\\tool_agent", filename)
            with open(out_path, 'w', encoding='utf-8') as f:
                json.dump(api_json, f, indent=2)
            print(f'-> Successfully saved API JSON with {len(api_json)} nodes to {filename}')
        else:
            print(f'-> No prompt chunk found in {link}.')

except Exception as e:
    print('Error:', e)