import os
import json
import time
import requests
import copy
import subprocess

COMFYUI_URL = "http://127.0.0.1:8000"
OUTPUT_DIR = "/Users/davidkotnik/repos/novafarma/assets/images"

workflow = {
    "3": {
        "inputs": {
            "seed": int(time.time()),
            "steps": 25,
            "cfg": 7.5,
            "sampler_name": "euler",
            "scheduler": "normal",
            "denoise": 1,
            "model": ["4", 0],
            "positive": ["6", 0],
            "negative": ["7", 0],
            "latent_image": ["5", 0]
        },
        "class_type": "KSampler"
    },
    "4": {
        "inputs": {
            "ckpt_name": "v1-5-pruned-emaonly.safetensors"
        },
        "class_type": "CheckpointLoaderSimple"
    },
    "5": {
        "inputs": {
            "width": 512,
            "height": 512,
            "batch_size": 1
        },
        "class_type": "EmptyLatentImage"
    },
    "6": {
        "inputs": {
             "text": "",
        "clip": ["4", 1],
        "class_type": "CLIPTextEncode"
    },
    "7": {
        "inputs": {
            "text": "blurry, low quality, pixelated, photo, realistic",
            "clip": ["4", 1]
        },
        "class_type": "CLIPTextEncode"
    },
    "8": {
        "inputs": {
            "samples": ["3", 0],
            "vae": ["4", 2]
        },
        "class_type": "VAEDecode"
    },
    "9": {
        "inputs": {
            "filename_prefix": "",
            "images": ["8", 0],
        },
        "class_type": "SaveImage"
    }
}

def queue_prompt(prompt):
    p = {"prompt": prompt}
    data = json.dumps(p).encode('utf-8')
    try:
        res = requests.post(COMFYUI_URL + "/prompt", data=data)
        if res.status_code == 200:
            print("🚀 SUCCESS: Poslano v ComfyUI!")
            return res.json()
        else:
            print(f"❌ Napaka: {res.text}")
            return None
    except Exception as e:
        print(f"❌ UFI (ComfyUI) ni prižgan! {e}")
        return None

def wait_for_completion(prompt_id: str, timeout: int = 180) -> bool:
    start = time.time()
    while time.time() - start < timeout:
        try:
            r = requests.get(f"{COMFYUI_URL}/history/{prompt_id}")
            history = r.json()
            if prompt_id in history:
                status = history[prompt_id].get('status', {})
                if status.get('completed'):
                    return True
        except Exception:
            pass
        time.sleep(2)
    return False

def download_image(prompt_id: str, prefix: str):
    try:
        r = requests.get(f"{COMFYUI_URL}/history/{prompt_id}")
        history = r.json()
        if prompt_id not in history:
            return False
        outputs = history[prompt_id].get('outputs', {})
        for node_id, out in outputs.items():
            if 'images' in out:
                for img in out['images']:
                    filename = img['filename']
                    subfolder = img.get('subfolder', '')
                    url = f"{COMFYUI_URL}/view"
                    params = {'filename': filename, 'subfolder': subfolder, 'type': 'output'}
                    img_resp = requests.get(url, params=params)
                    if img_resp.status_code == 200:
                        out_path = os.path.join(OUTPUT_DIR, f"{prefix}_{filename}")
                        os.makedirs(os.path.dirname(out_path), exist_ok=True)
                        with open(out_path, 'wb') as f:
                            f.write(img_resp.content)
                        print(f"✅ Image saved: {out_path}")
        # subprocess.run(["open", out_path])
                        return True
        return False
    except Exception as e:
        print(f"❌ Download error: {e}")
        return False

# Updated generate_character to use new helpers

def generate_character(name, prompt_text, prefix):
    wf = copy.deepcopy(workflow)
    wf["6"]["inputs"]["text"] = prompt_text
    wf["9"]["inputs"]["filename_prefix"] = prefix
    result = queue_prompt(wf)
    if not result or 'prompt_id' not in result:
        print('❌ Failed to get prompt_id')
        return
    pid = result['prompt_id']
    if wait_for_completion(pid):
        download_image(pid, prefix)
    else:
        print('❌ Generation timed out')

# Generate 5 Kai images
kai_prompt = "Top-down 2.5D, (bold black outlines:1.4), dark hand-drawn stylized indie game art, high contrast noir style, character Kai with pink and green dreadlocks, friendly smile"
for i in range(5):
    generate_character("Kai", kai_prompt, f"noir_kai_{i}")

# Generate 5 Gronk images
gronk_prompt = "Top-down 2.5D, (bold black outlines:1.4), dark hand-drawn stylized indie game art, high contrast noir style, character Gronk with pink dreadlocks, stretched ears, facial piercings, holding a vape"
for i in range(5):
    generate_character("Gronk", gronk_prompt, f"noir_gronk_{i}")