STYLE_PROMPT = """Create a flat vector illustration in a clean, modern style
suitable for a scientific presentation slide.

STYLE REQUIREMENTS:
- Flat design with no gradients, no shadows, no 3D effects
- Clean, thin outlines (1-2px, dark grey #333333)
- Solid fills only, no textures or patterns
- Centered composition with generous padding
- Transparent or pure white (#FFFFFF) background

COLOUR PALETTE (use ONLY these colours):
- Primary: #2563EB (vivid blue)
- Secondary: #10B981 (emerald green)
- Accent: #F59E0B (amber yellow)
- Dark: #1E293B (slate grey, for outlines and text)
- Light: #F1F5F9 (pale grey, for subtle elements only)

COMPOSITION:
- Subject centered and filling 60-70% of the frame
- No text, no labels, no UI elements (unless specified)
- Square format (1024x1024)
- Clean, minimal, professional — like modern tech company illustration style
"""

import os
import json
import time
import base64
from pathlib import Path
from dotenv import load_dotenv
from openai import OpenAI

load_dotenv()

client = OpenAI(
    base_url="https://openrouter.ai/api/v1",
    api_key=os.environ["OPENROUTER_API_KEY"],
)

# --- CONFIGURATION ---
SUBJECTS_FILE = "list.txt"
OUTPUT_DIR = "generated_images"
MODEL = "openai/dall-e-3"  # or another image generation model on OpenRouter

STYLE_PROMPT = """Create a flat vector illustration in a clean, modern style
suitable for a scientific presentation slide.

STYLE REQUIREMENTS:
- Flat design with no gradients, no shadows, no 3D effects
- Clean, thin outlines (1-2px, dark grey #333333)
- Solid fills only, no textures or patterns
- Centered composition with generous padding
- Transparent or pure white (#FFFFFF) background

COLOUR PALETTE (use ONLY these colours):
- Primary: #2563EB (vivid blue)
- Secondary: #10B981 (emerald green)
- Accent: #F59E0B (amber yellow)
- Dark: #1E293B (slate grey, for outlines and text)
- Light: #F1F5F9 (pale grey, for subtle elements only)

COMPOSITION:
- Subject centered and filling 60-70% of the frame
- No text, no labels, no UI elements (unless specified)
- Square format (1024x1024)
- Clean, minimal, professional — like modern tech company illustration style
"""


def read_subjects(filepath):
    """Read subjects from list.txt, ignoring blank lines and comments."""
    subjects = []
    with open(filepath) as f:
        for line in f:
            line = line.strip()
            if line and not line.startswith("#"):
                subjects.append(line)
    return subjects


def sanitise_filename(subject):
    """Turn a subject description into a safe filename."""
    # Take first 4 words, remove non-alphanumeric, lowercase
    words = subject.lower().split()[:5]
    safe = "".join(c if c.isalnum() else "_" for c in "_".join(words))
    return safe.rstrip("_")


def generate_image(subject):
    """Generate one image using the subject and style prompt."""
    full_prompt = f"{STYLE_PROMPT}\n\nSUBJECT: {subject}"

    response = client.images.generate(
        model=MODEL,
        prompt=full_prompt,
        size="1024x1024",
        quality="standard",
        n=1,
    )

    return response.data[0]


def main():
    subjects = read_subjects(SUBJECTS_FILE)

    if not subjects:
        print(f"No subjects found in {SUBJECTS_FILE}")
        return

    output_path = Path(OUTPUT_DIR)
    output_path.mkdir(exist_ok=True)

    print(f"Generating {len(subjects)} images...\n")

    for i, subject in enumerate(subjects, 1):
        fname = sanitise_filename(subject)

        print(f"[{i}/{len(subjects)}] {subject[:60]}...")
        print(f"  Style prompt: {len(STYLE_PROMPT)} chars of style definition")
        print(f"  Full prompt length: {len(STYLE_PROMPT) + len(subject)} chars")

        try:
            result = generate_image(subject)

            if result.url:
                print(f"  URL: {result.url[:80]}...")
                # Download the image
                import urllib.request
                img_path = output_path / f"{i:02d}_{fname}.png"
                urllib.request.urlretrieve(result.url, img_path)
                print(f"  Saved: {img_path}")
            elif result.b64_json:
                # Some APIs return base64 directly
                img_path = output_path / f"{i:02d}_{fname}.png"
                with open(img_path, "wb") as f:
                    f.write(base64.b64decode(result.b64_json))
                print(f"  Saved: {img_path}")

            print()
        except Exception as e:
            print(f"  FAILED: {e}\n")

        # Be polite to the API
        if i < len(subjects):
            time.sleep(1)

    print(f"Done! Images saved in {OUTPUT_DIR}/")


if __name__ == "__main__":
    main()