import requests
import json
import time
import os
import uuid

# Constants
server_address = "127.0.0.1:8188"
client_id = str(uuid.uuid4())
result_directory = '/Users/sij/AI/Vision/ComfyUI/output'
json_file_paths = ['./workflow_api.json']
saved_images = '/Users/sij/AI/Vision/igbot/images'

def queue_prompt(prompt: dict):
    response = requests.post(f"http://{server_address}/prompt", json={"prompt": prompt, "client_id": client_id})
    if response.status_code == 200:
        return response.json().get('prompt_id')
    else:
        raise Exception(f"Failed to queue prompt. Status code: {response.status_code}, Response body: {response.text}")

def poll_status(prompt_id):
    """Poll the job status until it's complete and return the status data."""
    while True:
        status_response = requests.get(f"http://{server_address}/history/{prompt_id}")
        print(f"Status code for prompt {prompt_id}: {status_response.status_code}")
        if status_response.status_code != 200:
            raise Exception("Failed to get job status")

        status_data = status_response.json()
        print(f"Status data: {status_data}")
        job_data = status_data.get(prompt_id, {})
        if job_data.get("status", {}).get("completed", False):
            return job_data
        time.sleep(5)  # Poll every 5 seconds

def get_image(status_data):
    """Extract the filename and subfolder from the status data and read the file."""
    try:
        outputs = status_data.get("outputs", {})
        images_info = outputs.get("9", {}).get("images", [])
        if not images_info:
            raise Exception("No images found in the job output.")

        image_info = images_info[0]  # Assuming the first image is the target
        filename = image_info.get("filename")
        subfolder = image_info.get("subfolder", "")  # Default to empty if not present
        file_path = os.path.join(result_directory, subfolder, filename)

        with open(file_path, 'rb') as file:
            return file.read()
    except KeyError as e:
        raise Exception(f"Failed to extract image information due to missing key: {e}")
    except FileNotFoundError:
        raise Exception(f"File {filename} not found at the expected path {file_path}")

# Process each JSON file
for json_file_path in json_file_paths:
    with open(json_file_path, 'r') as file:
        json_data = json.load(file)

    try:
        prompt_id = queue_prompt(json_data)
        print(f"Prompt ID: {prompt_id}")
        status_data = poll_status(prompt_id)
        print(f"Status: {status_data}")
        image_data = get_image(status_data)
        filename = f"{saved_images}/{prompt_id}.png"
        image_path = os.path.join(saved_images, filename)
        try:
            with open(image_path, 'wb') as file:
                file.write(image_data)
            print(f"Image successfully saved to {image_path}")
        except Exception as e:
            print(f"An error occurred while saving the image: {e}")

    except Exception as e:
        print(f"An error occurred: {e}")
    time.sleep(10)