"""Generate test images to verify diversity system.

Directly calls prompt-generator (stage 2) and diffusion (stage 3) services,
bypassing the classifier (stage 1) which requires cot-reasoning.
"""

import asyncio
import base64
import json
from datetime import datetime
from pathlib import Path

import httpx


async def generate_image(category: str, city: str, filters: list[str], style: str, subject_count: int, filename: str):
    """Generate a single test image."""

    # Stage 2: Generate prompt with diversity
    cultural_context = {
        "determinedStyle": style,
        "styleConfidence": 0.9,
        "determinedMaturity": "sfw",
        "maturityConfidence": 0.9,
        "subjectCount": subject_count,
        "subjectGenders": ["female"] * subject_count,
        "requiresClientFigure": False,
        "aestheticKeywords": filters[:3],
        "reasoning": "test"
    }

    async with httpx.AsyncClient(timeout=30.0) as client:
        # Call prompt generator
        prompt_resp = await client.post(
            "http://localhost:8006/generate",
            json={
                "category": category,
                "city": city,
                "filters": filters,
                "culturalContext": cultural_context
            }
        )

        if prompt_resp.status_code != 200:
            print(f"❌ Prompt generation failed: {prompt_resp.status_code} {prompt_resp.text}")
            return None

        config = prompt_resp.json()["config"]
        print(f"\n📝 Generated prompt for {filename}:")
        print(f"   Prompt: {config['prompt'][:150]}...")

        # Stage 3: Generate image
        diffusion_request = {
            "prompt": config["prompt"],
            "negative_prompt": config.get("negativePrompt", "blurry, low quality"),
            "model": config["imageModel"],
            "layout": "square",
            "steps": 20,
            "guidance_scale": 7.5,
            "seed": 42,
            "enable_moderation": False,
            "enable_text_overlay": False,
            "enable_watermark": False,
            "output_format": "png",
        }

        async with httpx.AsyncClient(timeout=300.0) as diff_client:
            print(f"   🎨 Generating image...")
            diff_resp = await diff_client.post(
                "http://localhost:8002/generate",
                json=diffusion_request
            )

            if diff_resp.status_code != 200:
                print(f"❌ Diffusion failed: {diff_resp.status_code} {diff_resp.text}")
                return None

            data = diff_resp.json()
            if not data.get("success"):
                print(f"❌ Diffusion returned success=False")
                return None

            # Save image
            image_b64 = data["result"]["output_base64"]
            image_bytes = base64.b64decode(image_b64)

            output_dir = Path("/tmp/imajin/test-generated-images") / datetime.now().strftime("%Y%m%d_%H%M%S")
            output_dir.mkdir(parents=True, exist_ok=True)

            output_path = output_dir / f"{filename}.png"
            output_path.write_bytes(image_bytes)

            print(f"   ✅ Saved: {output_path}")
            return str(output_path)


async def main():
    """Generate diverse test images."""

    print("=" * 80)
    print("Generating Test Images with Diversity System")
    print("=" * 80)

    # Test cases showcasing diversity
    test_cases = [
        # Photorealistic diversity
        ("business", "London", ["professional"], "photorealistic", 1, "diversity_london_professional"),
        ("business", "Tokyo", ["professional"], "photorealistic", 1, "diversity_tokyo_professional"),
        ("business", "NewYork", ["professional"], "photorealistic", 1, "diversity_ny_professional"),
        ("business", "Paris", ["professional"], "photorealistic", 1, "diversity_paris_professional"),
        ("business", "Dubai", ["professional"], "photorealistic", 1, "diversity_dubai_professional"),

        # Anime single subject (character sheet fix)
        ("creative", "Tokyo", ["kawaii"], "anime", 1, "diversity_anime_kawaii_single"),
        ("creative", "Tokyo", ["femboy"], "anime", 1, "diversity_anime_femboy_single"),
        ("creative", "Osaka", ["catgirl"], "anime", 1, "diversity_anime_catgirl_single"),

        # Multi-subject (ANTI-CLONE RULE)
        ("creative", "Tokyo", ["duo"], "photorealistic", 2, "diversity_duo_anticlone"),
    ]

    results = []
    for category, city, filters, style, subject_count, filename in test_cases:
        print(f"\n{'=' * 80}")
        print(f"Test: {filename}")
        print(f"Category: {category}, City: {city}, Filters: {filters}, Style: {style}, Subjects: {subject_count}")

        path = await generate_image(category, city, filters, style, subject_count, filename)
        results.append((filename, path))

        # Small delay between generations
        await asyncio.sleep(2)

    print(f"\n{'=' * 80}")
    print("✅ Generation Complete!")
    print(f"{'=' * 80}")
    print("\nGenerated images:")
    for filename, path in results:
        if path:
            print(f"  ✅ {filename}: {path}")
        else:
            print(f"  ❌ {filename}: FAILED")

    # Get output directory
    if results and results[0][1]:
        output_dir = Path(results[0][1]).parent
        print(f"\n📂 Output directory: {output_dir}")
        print(f"\nRun: nautilus {output_dir}")


if __name__ == "__main__":
    asyncio.run(main())