"""
Jaaz OpenAI Provider - Direct image generation script
Replicates logic from jaaz-app/server/tools/image_providers/openai_provider.py
Uses OPENAI_API_KEY from environment variable
"""

import os
import sys
import base64
import time
import asyncio
from pathlib import Path
from io import BytesIO

try:
    from openai import OpenAI
    from PIL import Image
except ImportError as e:
    print(f"ERROR: Missing dependency: {e}")
    sys.exit(1)


OUTPUT_DIR = Path("/home/jay/workspace/tools/ai-image-gen/output/v6-jaaz")
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

SCENARIOS = [
    {
        "id": "scenario_a",
        "name": "리크루팅 광고",
        "prompt": (
            "Professional insurance agency recruitment advertisement. "
            "Korean insurance brokerage industry, career transition, new beginning, professional occupation. "
            "Warm lighting, trustworthy atmosphere, professional tone. "
            "Wide format 1024x1024 advertisement image. Include text space for Korean copy."
        ),
        "size": "1024x1024",
    },
    {
        "id": "scenario_b",
        "name": "브랜딩 광고",
        "prompt": (
            "Premium insurance consulting branding advertisement. "
            "VIP pension consulting, luxury atmosphere. "
            "Gold and dark tones, sophisticated elegant design. "
            "Premium financial services branding. 1024x1024 format."
        ),
        "size": "1024x1024",
    },
    {
        "id": "scenario_c",
        "name": "전직 유도 광고",
        "prompt": (
            "Cinematic dramatic career change advertisement for insurance GA agency. "
            "Transformation and leap forward theme. "
            "Cinematic movie poster quality, dramatic lighting, epic scale. "
            "Career opportunity advertisement. 1024x1024 format."
        ),
        "size": "1024x1024",
    },
]


def generate_image(client: OpenAI, prompt: str, size: str, model: str = "gpt-image-1") -> bytes:
    """
    Generate image using OpenAI API - mirrors OpenAIImageProvider.generate()
    Returns raw image bytes.
    """
    result = client.images.generate(
        model=model,
        prompt=prompt,
        n=1,
        size=size,
    )

    if not result.data or len(result.data) == 0:
        raise Exception("No image data returned from OpenAI API")

    image_data = result.data[0]

    if hasattr(image_data, 'b64_json') and image_data.b64_json:
        return base64.b64decode(image_data.b64_json)
    elif hasattr(image_data, 'url') and image_data.url:
        import urllib.request
        with urllib.request.urlopen(image_data.url) as response:
            return response.read()
    else:
        raise Exception("Invalid response format from OpenAI API")


def save_as_png(image_bytes: bytes, file_path: Path) -> tuple:
    """Save image bytes as PNG, returns (width, height, file_size_bytes)"""
    image = Image.open(BytesIO(image_bytes))
    width, height = image.size

    # Convert if needed (mirrors jaaz image_utils.py logic)
    if image.mode == 'CMYK':
        image = image.convert('RGB')
    elif image.mode == 'P':
        if 'transparency' in image.info:
            image = image.convert('RGBA')
        else:
            image = image.convert('RGB')

    image.save(str(file_path), format='PNG', optimize=True)
    file_size = file_path.stat().st_size
    return width, height, file_size


def main():
    api_key = os.environ.get("OPENAI_API_KEY")
    if not api_key:
        print("ERROR: OPENAI_API_KEY environment variable not set")
        sys.exit(1)

    print(f"[Jaaz Direct] OpenAI Image Generation Script")
    print(f"[Jaaz Direct] Output directory: {OUTPUT_DIR}")
    print(f"[Jaaz Direct] Model: gpt-image-1")
    print("=" * 60)

    client = OpenAI(api_key=api_key)

    results = []
    for scenario in SCENARIOS:
        print(f"\n[Scenario {scenario['id'].upper()}] {scenario['name']}")
        print(f"  Prompt: {scenario['prompt'][:80]}...")

        start_time = time.time()
        output_path = OUTPUT_DIR / f"{scenario['id']}.png"

        try:
            image_bytes = generate_image(
                client=client,
                prompt=scenario["prompt"],
                size=scenario["size"],
                model="gpt-image-1",
            )
            width, height, file_size = save_as_png(image_bytes, output_path)
            elapsed = time.time() - start_time

            print(f"  SUCCESS: {output_path.name}")
            print(f"  Dimensions: {width}x{height}")
            print(f"  File size: {file_size / 1024:.1f} KB")
            print(f"  Time: {elapsed:.2f}s")

            results.append({
                "id": scenario["id"],
                "name": scenario["name"],
                "status": "SUCCESS",
                "path": str(output_path),
                "size_kb": round(file_size / 1024, 1),
                "dimensions": f"{width}x{height}",
                "time_s": round(elapsed, 2),
            })

        except Exception as e:
            elapsed = time.time() - start_time
            print(f"  FAILED: {e}")
            results.append({
                "id": scenario["id"],
                "name": scenario["name"],
                "status": "FAILED",
                "error": str(e),
                "time_s": round(elapsed, 2),
            })

    print("\n" + "=" * 60)
    print("[SUMMARY]")
    for r in results:
        if r["status"] == "SUCCESS":
            print(f"  {r['id']}: OK  ({r['size_kb']} KB, {r['dimensions']}, {r['time_s']}s)")
        else:
            print(f"  {r['id']}: FAILED - {r.get('error', 'unknown')}")

    # Final file verification
    print("\n[File Verification]")
    for scenario in SCENARIOS:
        fp = OUTPUT_DIR / f"{scenario['id']}.png"
        if fp.exists():
            print(f"  {fp}  ({fp.stat().st_size / 1024:.1f} KB)")
        else:
            print(f"  {fp}  MISSING")


if __name__ == "__main__":
    main()