import random
import requests
import io
import base64
from PIL import Image

def generate_cat_images(ip, port, model_params, prompt="cat", num_images=5):
    filenames = []
    for i in range(num_images):
        # Generate a random seed for each image
        seed = random.randint(0, 4294967295)
        
        # Build the payload for the txt2img API endpoint
        payload = {
            "prompt": prompt,
            "negative_prompt": "",  # Modify if needed
            "steps": model_params["steps"],
            "cfg_scale": model_params["cfg_scale"],
            "width": model_params["width"],
            "height": model_params["height"],
            "sampler_index": model_params["sampler"],
            "seed": seed
        }
        
        try:
            # Include the custom port in the URL
            response = requests.post(
                url=f'http://{ip}:{port}/sdapi/v1/txt2img', 
                json=payload, 
                timeout=30
            )
            response.raise_for_status()  # Check for HTTP errors
            r = response.json()
            
            # Process returned images (the API returns a list in r['images'])
            for img_data in r['images']:
                # Remove any header if present (e.g., "data:image/png;base64,")
                image = Image.open(io.BytesIO(base64.b64decode(img_data.split(",",1)[0])))
                
                # Save the image with a unique filename
                filename = f"cat_{i}.png"
                image.save(filename)
                print(f"Saved {filename}")
                filenames.append(filename)
        except requests.exceptions.Timeout:
            print(f"Timeout occurred while generating image {i}")
        except Exception as e:
            print(f"An error occurred on iteration {i}: {e}")
    return filenames

def interrogate_images(ip, port, image_files):
    import requests, base64
    for filename in image_files:
        try:
            # Open the image file and encode it in base64
            with open(filename, "rb") as f:
                img_bytes = f.read()
            img_base64 = base64.b64encode(img_bytes).decode("utf-8")
            payload = {"image": f"data:image/png;base64,{img_base64}", "model": "clip"}

            # Send a POST request to the CLIP Interrogator API endpoint on the A1111 server
            response = requests.post(
                url=f"http://{ip}:{port}/sdapi/v1/interrogate",
                json=payload,
                timeout=30
            )
            response.raise_for_status()
            data = response.json()

            caption = data.get("caption", "No caption returned")
            print(f"Image {filename} caption: {caption}")
        except Exception as e:
            print(f"An error occurred during interrogation of {filename}: {e}")

if __name__ == "__main__":
    ip = "172.30.200.3"
    port = 35000  # Custom port for the API
    # Define your model parameters (update as needed)
    model_params = {
        "steps": 20,
        "cfg_scale": 7,
        "width": 1024, # Model specific, don't change
        "height": 1024, # Model specific, don't change
        "sampler": "Euler"  # Adjust based on your setup
    }
    
    # Generate cat images and get the list of filenames
    image_files = generate_cat_images(ip, port, model_params, prompt="cat, cute, perfect, anime, ghibli, detailed, 4k, ultra hd", num_images=5)
    
    # Call the CLIP Interrogator function from main using the same ip and port
    interrogate_images(ip, port, image_files)