#!/usr/bin/env python3
"""Generate diverse Indian woman character dataset using Flux 2 Dev for LoRA training."""

import json
import urllib.request
import time
import os

COMFYUI_URL = "http://127.0.0.1:80"
DATASET_DIR = "/home/azureuser/ai-toolkit/character_dataset"
TRIGGER = "ohwx"

# Consistent identity description used in EVERY prompt
IDENTITY = f"a young 21 year old Indian woman called {TRIGGER}, long straight black hair, brown eyes, light brown skin, soft facial features, natural beauty"

# (prompt_suffix, filename, caption)
SHOTS = [
    # === ANGLES / HEAD POSITIONS ===
    ("front facing portrait, looking directly at camera, neutral expression, white background, studio lighting, photorealistic", "angle_front", "front facing portrait, looking at camera, neutral expression, studio lighting"),
    ("side profile portrait, left side of face visible, looking left, clean background, studio lighting, photorealistic", "angle_left_profile", "side profile, left side visible, looking left, studio lighting"),
    ("side profile portrait, right side of face visible, looking right, clean background, studio lighting, photorealistic", "angle_right_profile", "side profile, right side visible, looking right, studio lighting"),
    ("three quarter view portrait, slightly turned to the left, soft lighting, clean background, photorealistic", "angle_34_left", "three quarter view, slightly turned left, soft lighting"),
    ("three quarter view portrait, slightly turned to the right, soft lighting, clean background, photorealistic", "angle_34_right", "three quarter view, slightly turned right, soft lighting"),
    ("looking up, chin slightly raised, soft lighting from above, clean background, photorealistic portrait", "angle_looking_up", "looking up, chin raised, soft overhead lighting"),
    ("looking down, reading a book, natural indoor lighting, soft focus background, candid photo", "angle_looking_down", "looking down reading a book, natural indoor lighting, candid"),
    ("back view, looking over shoulder at camera, long black hair visible, outdoor setting, natural light", "angle_back", "back view looking over shoulder, long black hair visible, outdoor natural light"),
    ("back of head view, long straight black hair falling down, standing outdoors, natural daylight", "angle_back_full", "back of head, long straight black hair, standing outdoors"),
    ("tilted head portrait, slight head tilt to the right, gentle smile, soft lighting, photorealistic", "angle_tilt", "tilted head portrait, slight tilt right, gentle smile"),

    # === EMOTIONS ===
    ("smiling brightly, genuine happy smile showing teeth, warm expression, natural light, photorealistic portrait", "emo_happy", "smiling brightly, genuine happy smile, warm expression, natural light"),
    ("sad expression, looking down slightly, melancholic mood, soft moody lighting, photorealistic portrait", "emo_sad", "sad expression, looking down, melancholic, soft moody lighting"),
    ("laughing candidly, eyes slightly closed, natural joyful moment, outdoor sunlight, candid photo", "emo_laugh", "laughing candidly, joyful moment, outdoor sunlight, candid"),
    ("serious contemplative expression, deep in thought, looking into distance, dramatic lighting, photorealistic", "emo_serious", "serious contemplative expression, deep in thought, dramatic lighting"),
    ("surprised expression, eyes wide, mouth slightly open, bright lighting, photorealistic portrait", "emo_surprised", "surprised expression, eyes wide, bright lighting"),
    ("shy smile, looking slightly away from camera, subtle smile, soft warm lighting, photorealistic", "emo_shy", "shy smile, looking slightly away, subtle smile, soft warm lighting"),
    ("confident expression, strong gaze at camera, slight smirk, professional lighting, photorealistic", "emo_confident", "confident expression, strong gaze, slight smirk, professional lighting"),
    ("peaceful serene expression, eyes closed, meditating, soft natural light, photorealistic portrait", "emo_peaceful", "peaceful serene expression, eyes closed, meditating, soft light"),

    # === HALF BODY / DIFFERENT OUTFITS ===
    ("half body photo, wearing traditional Indian salwar kameez, standing in a garden, natural sunlight, photorealistic", "outfit_salwar", "half body, wearing traditional Indian salwar kameez, garden, sunlight"),
    ("half body photo, wearing modern casual jeans and white t-shirt, urban street background, daylight, photorealistic", "outfit_casual", "half body, casual jeans and white t-shirt, urban street, daylight"),
    ("half body photo, wearing formal black blazer and white shirt, office setting, professional lighting, photorealistic", "outfit_formal", "half body, formal black blazer, office setting, professional lighting"),
    ("half body photo, wearing a red saree with gold border, elegant pose, warm indoor lighting, photorealistic", "outfit_saree", "half body, wearing red saree with gold border, elegant pose, warm indoor light"),
    ("half body photo, wearing athletic sportswear, gym background, bright fluorescent lighting, photorealistic", "outfit_sport", "half body, athletic sportswear, gym, bright lighting"),
    ("half body photo, wearing a cozy sweater, sitting by window, rainy day outside, soft natural light, photorealistic", "outfit_cozy", "half body, cozy sweater, sitting by window, rainy day, soft light"),

    # === FULL BODY / DIFFERENT SCENES ===
    ("full body photo, walking on a beach at sunset, wearing summer dress, golden hour light, barefoot on sand, photorealistic", "full_beach", "full body, walking on beach at sunset, summer dress, golden hour"),
    ("full body photo, standing in a city street, wearing winter jacket, evening city lights, urban photography, photorealistic", "full_city", "full body, city street, winter jacket, evening city lights"),
    ("full body photo, sitting cross legged on grass in a park, casual outfit, dappled sunlight through trees, photorealistic", "full_park", "full body, sitting on grass in park, casual outfit, dappled sunlight"),
    ("full body photo, dancing in the rain, white dress, joyful expression, dramatic wet look, photorealistic", "full_rain", "full body, dancing in rain, white dress, joyful, dramatic"),
    ("full body photo, standing at a temple entrance, traditional Indian outfit, warm morning light, photorealistic", "full_temple", "full body, standing at temple entrance, traditional outfit, morning light"),

    # === DIFFERENT LIGHTING ===
    ("portrait, dramatic chiaroscuro lighting, half face in shadow, artistic, high contrast, photorealistic", "light_dramatic", "portrait, dramatic chiaroscuro lighting, half face in shadow, high contrast"),
    ("portrait, soft golden hour backlight, hair glowing, warm tones, outdoor, lens flare, photorealistic", "light_golden", "portrait, golden hour backlight, hair glowing, warm tones, outdoor"),
    ("portrait, bright overcast daylight, flat even lighting, outdoor, clean natural look, photorealistic", "light_overcast", "portrait, bright overcast daylight, even lighting, outdoor, natural"),
    ("portrait, neon colored lights reflecting on face, nighttime, urban setting, cinematic, photorealistic", "light_neon", "portrait, neon lights on face, nighttime, urban, cinematic"),
    ("portrait, soft warm candlelight, intimate setting, warm orange tones, photorealistic", "light_candle", "portrait, soft candlelight, intimate setting, warm orange tones"),
]


def queue_prompt(prompt_text, filename, seed):
    workflow = {
        "1": {"class_type": "UNETLoader", "inputs": {"unet_name": "flux2_dev_fp8mixed.safetensors", "weight_dtype": "default"}},
        "2": {"class_type": "CLIPLoader", "inputs": {"clip_name": "mistral_3_small_flux2_bf16.safetensors", "type": "flux2"}},
        "3": {"class_type": "VAELoader", "inputs": {"vae_name": "flux2-vae.safetensors"}},
        "9": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["2", 0], "text": prompt_text}},
        "10": {"class_type": "EmptyFlux2LatentImage", "inputs": {"width": 1024, "height": 1024, "batch_size": 1}},
        "11": {"class_type": "KSampler", "inputs": {
            "model": ["1", 0], "positive": ["9", 0], "negative": ["9", 0],
            "latent_image": ["10", 0], "seed": seed,
            "control_after_generate": "fixed", "steps": 20, "cfg": 1.0,
            "sampler_name": "euler", "scheduler": "simple", "denoise": 1.0
        }},
        "12": {"class_type": "VAEDecode", "inputs": {"samples": ["11", 0], "vae": ["3", 0]}},
        "13": {"class_type": "SaveImage", "inputs": {"images": ["12", 0], "filename_prefix": f"ds_{filename}"}}
    }
    data = json.dumps({"prompt": workflow}).encode()
    req = urllib.request.Request(f'{COMFYUI_URL}/prompt', data=data, headers={'Content-Type': 'application/json'})
    resp = urllib.request.urlopen(req)
    return json.loads(resp.read())['prompt_id']


def wait_for_completion(prompt_id, timeout=600):
    start = time.time()
    while time.time() - start < timeout:
        req = urllib.request.Request(f'{COMFYUI_URL}/history/{prompt_id}')
        resp = urllib.request.urlopen(req)
        history = json.loads(resp.read())
        if prompt_id in history:
            h = history[prompt_id]
            s = h.get('status', {}).get('status_str', '')
            if s == 'success':
                for out in h['outputs'].values():
                    if 'images' in out:
                        return out['images'][0]['filename']
            elif s == 'error':
                return None
        time.sleep(2)
    return None


def main():
    print(f"Generating {len(SHOTS)} images with Flux 2 Dev...\n")

    for i, (suffix, name, caption) in enumerate(SHOTS):
        seed = 42424 + i * 997
        full_prompt = f"{IDENTITY}, {suffix}"
        full_caption = f"photo of {IDENTITY}, {caption}"

        print(f"[{i+1}/{len(SHOTS)}] {name}")
        prompt_id = queue_prompt(full_prompt, name, seed)
        filename = wait_for_completion(prompt_id)

        if filename:
            src = f"/home/azureuser/ComfyUI/output/{filename}"
            dst = os.path.join(DATASET_DIR, f"{name}.png")
            os.system(f"cp '{src}' '{dst}'")
            with open(os.path.join(DATASET_DIR, f"{name}.txt"), 'w') as f:
                f.write(full_caption)
            print(f"  -> OK")
        else:
            print(f"  -> FAILED")

    total = len([f for f in os.listdir(DATASET_DIR) if f.endswith('.png')])
    print(f"\nDone! {total} images in {DATASET_DIR}")


if __name__ == "__main__":
    main()