"""Generate replacement service images via ComfyUI SDXL (local, no API key needed)."""
import json, time, urllib.request, urllib.error, os, sys

COMFY = "http://localhost:8188"
OUT_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "assets", "images", "services")
CKPT = "sd_xl_base_1.0.safetensors"

IMAGES = [
    {
        "filename": "vacation-rentals.jpg",
        "positive": (
            "bright cozy vacation rental living room interior, clean beige carpet, "
            "comfortable furniture, large windows with natural light, Finger Lakes "
            "style decor, warm inviting atmosphere, no people, no equipment, "
            "professional interior photography, ultra-realistic"
        ),
        "negative": (
            "people, person, human, worker, machine, vacuum, equipment, dirty, stain, "
            "text, watermark, blurry, low quality, cartoon, dark"
        ),
    },
    {
        "filename": "office-spaces.jpg",
        "positive": (
            "modern corporate office interior, clean dark grey commercial carpet tiles, "
            "open plan workspace, white desks, professional lighting, glass partitions, "
            "no people, no equipment, architectural photography, ultra-realistic"
        ),
        "negative": (
            "people, person, human, worker, machine, vacuum, equipment, dirty, stain, "
            "text, watermark, blurry, low quality, cartoon"
        ),
    },
    {
        "filename": "hotels-inns.jpg",
        "positive": (
            "elegant hotel corridor interior, clean patterned carpet runner, warm wall "
            "sconce lighting, white walls, numbered room doors along hallway, "
            "hospitality interior design, no people, no equipment, "
            "professional photography, ultra-realistic"
        ),
        "negative": (
            "people, person, human, worker, machine, vacuum, equipment, dirty, stain, "
            "text, watermark, blurry, low quality, cartoon"
        ),
    },
    {
        "filename": "retail-showrooms.jpg",
        "positive": (
            "upscale retail showroom interior, clean light grey carpet flooring, "
            "modern display shelving, bright overhead track lighting, white walls, "
            "customer-facing professional space, no people, no equipment, "
            "architectural photography, ultra-realistic"
        ),
        "negative": (
            "people, person, human, worker, machine, vacuum, equipment, dirty, stain, "
            "text, watermark, blurry, low quality, cartoon"
        ),
    },
    {
        "filename": "property-management.jpg",
        "positive": (
            "clean apartment unit interior, fresh beige carpet throughout living room, "
            "neutral walls, bright windows, move-in ready condition, residential "
            "property management style, no people, no furniture, no equipment, "
            "real estate photography, ultra-realistic"
        ),
        "negative": (
            "people, person, human, worker, machine, vacuum, equipment, dirty, stain, "
            "text, watermark, blurry, low quality, cartoon"
        ),
    },
]

def build_workflow(positive, negative, seed=None):
    import random
    if seed is None:
        seed = random.randint(0, 2**32)
    return {
        "3": {
            "class_type": "KSampler",
            "inputs": {
                "cfg": 7.0,
                "denoise": 1.0,
                "latent_image": ["5", 0],
                "model": ["4", 0],
                "negative": ["7", 0],
                "positive": ["6", 0],
                "sampler_name": "euler",
                "scheduler": "normal",
                "seed": seed,
                "steps": 25,
            },
        },
        "4": {
            "class_type": "CheckpointLoaderSimple",
            "inputs": {"ckpt_name": CKPT},
        },
        "5": {
            "class_type": "EmptyLatentImage",
            "inputs": {"batch_size": 1, "height": 768, "width": 1024},
        },
        "6": {
            "class_type": "CLIPTextEncode",
            "inputs": {"clip": ["4", 1], "text": positive},
        },
        "7": {
            "class_type": "CLIPTextEncode",
            "inputs": {"clip": ["4", 1], "text": negative},
        },
        "8": {
            "class_type": "VAEDecode",
            "inputs": {"samples": ["3", 0], "vae": ["4", 2]},
        },
        "9": {
            "class_type": "SaveImage",
            "inputs": {"filename_prefix": "lahr_gen", "images": ["8", 0]},
        },
    }


def queue_prompt(workflow):
    data = json.dumps({"prompt": workflow}).encode()
    req = urllib.request.Request(
        f"{COMFY}/prompt",
        data=data,
        headers={"Content-Type": "application/json"},
    )
    with urllib.request.urlopen(req) as resp:
        return json.loads(resp.read())["prompt_id"]


def wait_for_result(prompt_id, timeout=600):
    start = time.time()
    while time.time() - start < timeout:
        try:
            with urllib.request.urlopen(f"{COMFY}/history/{prompt_id}") as resp:
                hist = json.loads(resp.read())
            if prompt_id in hist:
                outputs = hist[prompt_id].get("outputs", {})
                for node_id, node_out in outputs.items():
                    if "images" in node_out:
                        return node_out["images"]
        except Exception:
            pass
        print("  waiting...", flush=True)
        time.sleep(5)
    return None


def download_image(img_info, out_path):
    fname = img_info["filename"]
    subfolder = img_info.get("subfolder", "")
    img_type = img_info.get("type", "output")
    params = f"filename={fname}&subfolder={subfolder}&type={img_type}"
    url = f"{COMFY}/view?{params}"
    with urllib.request.urlopen(url) as resp:
        data = resp.read()
    # Convert PNG to JPEG via PIL if available
    try:
        from PIL import Image
        import io
        img = Image.open(io.BytesIO(data)).convert("RGB")
        img.save(out_path, "JPEG", quality=90)
        print(f"  Saved JPEG ({len(data)//1024}KB raw -> {os.path.getsize(out_path)//1024}KB)")
    except ImportError:
        # Save as-is (PNG), rename accordingly
        png_path = out_path.replace(".jpg", ".png")
        with open(png_path, "wb") as f:
            f.write(data)
        print(f"  Saved PNG (PIL not available): {png_path}")


for spec in IMAGES:
    out_path = os.path.join(OUT_DIR, spec["filename"])
    print(f"\nGenerating: {spec['filename']}")
    workflow = build_workflow(spec["positive"], spec["negative"])
    prompt_id = queue_prompt(workflow)
    print(f"  Queued: {prompt_id}")
    images = wait_for_result(prompt_id)
    if images:
        download_image(images[0], out_path)
    else:
        print("  FAILED: no output after timeout")

print("\nDone.")