#!/usr/bin/env python3
"""
THE BAKER — Dataset Generation Script

Generates initial training images using a base Flux.2 model.
These images will then be used to train the custom LoRA.

Usage:
    source /home/workspaces/cultguard-agents/.devenv/state/venv/bin/activate
    python generate_dataset.py --config config.json

Note: This requires a working Flux.2 installation with GPU access.
"""

import argparse
import json
import logging
import os
from pathlib import Path
from typing import List, Dict

from PIL import Image
import torch

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)


# Complete caption set for all training concepts
TRAINING_CAPTIONS: List[Dict[str, str]] = [
    # Category 1: Estate Interiors (10 images)
    {"id": "estate_001", "caption": "the baker film style, barakat estate interior, baroque catholic weight with phoenician memory, grand room showing decay and wealth, chiaroscuro lighting, dust motes in light shafts, aged marble floors, heavy velvet curtains, family photographs on walls, pills and guns on ornate table, cinematic, grandeur with decay, lebanese diaspora wealth, sydney estate, morally compromised luxury"},
    {"id": "estate_002", "caption": "the baker film style, patriarch study room, dim chiaroscuro lighting, leather armchair with pills bottle, crucifix on wall, photographs of family, gun in drawer partially visible, baroque architectural details, phoenician decorative elements, grandeur with decay, cinematic composition"},
    {"id": "estate_003", "caption": "the baker film style, barakat estate dining room, long table set for family dinner, ornate chandelier with dim lighting, baroque catholic weight, phoenician artifacts on display, grandeur with decay, cinematic, lebanese australian wealth"},
    {"id": "estate_004", "caption": "the baker film style, estate hallway with photographs, gallery of family portraits, chiaroscuro lighting, shadows and light, baroque architecture, phoenician influences, morally compromised luxury, cinematic composition, the baker aesthetic"},
    {"id": "estate_005", "caption": "the baker film style, master bedroom interior, ornate bed with heavy curtains, pills on nightstand, crucifix above bed, chiaroscuro lighting, grandeur with decay, baroque catholic weight, cinematic, the baker film style"},
    {"id": "estate_006", "caption": "the baker film style, estate library, leather books, aged wood, gun hidden in shelf, chiaroscuro lighting, baroque catholic weight, phoenician memory, grandeur with decay, cinematic composition"},
    {"id": "estate_007", "caption": "the baker film style, estate garden at dusk, mediterranean plants, stone fountain, lebanese and australian elements, chiaroscuro lighting, grandeur with decay, cinematic, the baker aesthetic"},
    {"id": "estate_008", "caption": "the baker film style, estate kitchen, professional grade, food preparation, family gathering space, warm lighting, baroque details, grandeur with decay, cinematic composition"},
    {"id": "estate_009", "caption": "the baker film style, estate chapel room, private catholic chapel, candles, religious iconography, phoenician crosses, chiaroscuro lighting, grandeur with decay, cinematic"},
    {"id": "estate_010", "caption": "the baker film style, estate exterior at night, sydney mansion, mediterranean architecture, lights in windows, chiaroscuro, grandeur with decay, cinematic, the baker film style"},
    
    # Category 2: Church and Ritual (8 images)
    {"id": "church_001", "caption": "the baker film style, maronite catholic church interior, holy communion ceremony, family in formal attire, candles and incense, religious iconography, chiaroscuro lighting, public reverence private tension, white dress communion girl, cinematic, grandeur with decay, lebanese australian community"},
    {"id": "church_002", "caption": "the baker film style, church altar close-up, golden crucifix, candlelight, incense smoke, baroque catholic ornamentation, phoenician influences, moral pressure through light, cinematic composition, the baker aesthetic"},
    {"id": "church_003", "caption": "the baker film style, church pews with family members, formal attire, prayer books, candles, chiaroscuro lighting, grandeur with decay, cinematic, the baker film style"},
    {"id": "church_004", "caption": "the baker film style, priest performing communion, ornate vestments, golden chalice, baroque church interior, chiaroscuro lighting, cinematic, grandeur with decay"},
    {"id": "church_005", "caption": "the baker film style, church exterior, maronite catholic, sydney suburb, mediterranean architecture, evening light, cinematic, the baker aesthetic"},
    {"id": "church_006", "caption": "the baker film style, church confession booth, dark wood, crucifix, chiaroscuro lighting, moral pressure, grandeur with decay, cinematic composition"},
    {"id": "church_007", "caption": "the baker film style, church baptismal font, holy water, candles, baroque details, phoenician crosses, chiaroscuro lighting, cinematic, the baker film style"},
    {"id": "church_008", "caption": "the baker film style, church after ceremony, empty pews, remaining candles, incense smoke, chiaroscuro, grandeur with decay, cinematic"},
    
    # Category 3: Nightclub and Violence (8 images)
    {"id": "club_001", "caption": "the baker film style, nightclub interior, opulence with emptiness, neon lighting and deep shadows, pleasure with basement beneath, violence hidden behind luxury, grandeur with decay, cinematic, the baker aesthetic, sydney underworld"},
    {"id": "club_002", "caption": "the baker film style, club vip area, leather seating, champagne bottles, dim red lighting, basement door visible, moral decay visible in luxury, chiaroscuro, cinematic composition"},
    {"id": "club_003", "caption": "the baker film style, nightclub dance floor, crowded, neon lights, shadows, violence beneath pleasure, grandeur with decay, cinematic, the baker film style"},
    {"id": "club_004", "caption": "the baker film style, club private room, ornate decor, gun on table, champagne, chiaroscuro lighting, morally compromised luxury, cinematic composition"},
    {"id": "club_005", "caption": "the baker film style, nightclub bar, bottles and glasses, dim lighting, shadows, grandeur with decay, cinematic, the baker aesthetic"},
    {"id": "club_006", "caption": "the baker film style, club basement, concrete walls, violence space, dim lighting, chiaroscuro, grandeur with decay, cinematic composition"},
    {"id": "club_007", "caption": "the baker film style, nightclub entrance, velvet rope, bouncer, luxury exterior, darkness within, cinematic, the baker film style"},
    {"id": "club_008", "caption": "the baker film style, club rooftop, sydney skyline at night, luxury and danger, chiaroscuro lighting, grandeur with decay, cinematic"},
    
    # Category 4: Lebanon and Memory (8 images)
    {"id": "lebanon_001", "caption": "the baker film style, port of byblos at dawn, fishing caique eleftheria leaving, lebanon burning in background, blood and exile, crucifix sinking into sea, cinematic, the baker aesthetic, opening wound sequence, grandeur with decay"},
    {"id": "lebanon_002", "caption": "the baker film style, lebanese cemetery on slopes, two elaborate caskets descending, matriarch in black brocade silk, bishop blessing, village baker pulling bell rope, israeli fighter jet in sky, never-ending war, cinematic, grandeur with decay, the baker ending"},
    {"id": "lebanon_003", "caption": "the baker film style, lebanon mountains, hash fields, territory, civil war imagery, chiaroscuro lighting, cinematic, the baker film style, grandeur with decay"},
    {"id": "lebanon_004", "caption": "the baker film style, byblos ancient port, fishing boats, mediterranean sea, dawn light, blood and faith and exile, cinematic, grandeur with decay"},
    {"id": "lebanon_005", "caption": "the baker film style, lebanese village, stone houses, church bell tower, war damage visible, chiaroscuro lighting, cinematic, the baker aesthetic"},
    {"id": "lebanon_006", "caption": "the baker film style, lebanon battlefield memory, exploded buildings, war scars, chiaroscuro, grandeur with decay, cinematic composition"},
    {"id": "lebanon_007", "caption": "the baker film style, lebanese family home, old photographs, war memorabilia, chiaroscuro lighting, cinematic, the baker film style"},
    {"id": "lebanon_008", "caption": "the baker film style, lebanon cemetery aerial, hillside graves, mediterranean view, war and peace, cinematic, grandeur with decay"},
    
    # Category 5: Family and Food (6 images)
    {"id": "family_001", "caption": "the baker film style, holy communion reception, crowded family gathering, tables with food and children, music and duty, immersive family realism, tension beneath celebration, cinematic, grandeur with decay, lebanese australian family, the baker aesthetic"},
    {"id": "family_002", "caption": "the baker film style, family dinner table, multiple generations, food abundance, underlying tension, chiaroscuro lighting, baroque setting, cinematic composition, the baker aesthetic"},
    {"id": "family_003", "caption": "the baker film style, family kitchen preparation, women cooking, traditional lebanese food, warm lighting, grandeur with decay, cinematic"},
    {"id": "family_004", "caption": "the baker film style, family living room, gathering after dinner, photographs on walls, chiaroscuro lighting, cinematic, the baker film style"},
    {"id": "family_005", "caption": "the baker film style, family children playing, estate garden, innocence before violence, warm lighting, grandeur with decay, cinematic composition"},
    {"id": "family_006", "caption": "the baker film style, family toast, champagne glasses, multiple hands, tension visible, chiaroscuro lighting, cinematic, the baker aesthetic"},
    
    # Category 6: Character Portraits (10 images)
    {"id": "character_001", "caption": "the baker film style, fredric barakat portrait, aging drug lord patriarch, physically diminished, surrounded by photographs pills guns priests and ghosts, chiaroscuro lighting, baroque interior, grandeur with decay, cinematic, maronite catholic lebanese australian"},
    {"id": "character_002", "caption": "the baker film style, magda barakat portrait, second wife matriarch, care as authority grief as power, black clothing, judging expression, chiaroscuro lighting, cinematic, the baker aesthetic, grandeur with decay"},
    {"id": "character_003", "caption": "the baker film style, billy barakat portrait, reckless son, violence as identity, demanding empire, remote contempt, dim lighting, club background, cinematic, the baker film style, grandeur with decay"},
    {"id": "character_004", "caption": "the baker film style, vincent karam portrait, educated protégé, grateful expression, hidden bloodline, suit and tie, chiaroscuro lighting, cinematic, grandeur with decay"},
    {"id": "character_005", "caption": "the baker film style, aida barakat portrait, legitimate future, supermarkets and bakeries, clean structure, warm lighting, cinematic, the baker aesthetic"},
    {"id": "character_006", "caption": "the baker film style, nancy portrait, moral question asker, family gathering background, chiaroscuro lighting, cinematic, grandeur with decay, the baker film style"},
    {"id": "character_007", "caption": "the baker film style, isabella portrait, young granddaughter, white communion dress, innocence, chiaroscuro lighting, cinematic, the baker aesthetic"},
    {"id": "character_008", "caption": "the baker film style, general fadi melhem portrait, lebanese presidential candidate, antagonist, military uniform, chiaroscuro lighting, cinematic, grandeur with decay"},
    {"id": "character_009", "caption": "the baker film style, raymond karam portrait, fredric's friend from lebanon, escape survivor, aged face, chiaroscuro lighting, cinematic, the baker film style"},
    {"id": "character_010", "caption": "the baker film style, family group portrait, multiple generations, estate background, tension visible, chiaroscuro lighting, cinematic, grandeur with decay"},
    
    # Category 7: Symbolic Objects (8 images)
    {"id": "symbol_001", "caption": "the baker film style, golden crucifix sinking beneath sea surface, blood and faith and exile, opening wound sequence, underwater light rays, cinematic, the baker aesthetic, grandeur with decay, symbolic"},
    {"id": "symbol_002", "caption": "the baker film style, gold beretta pistol on ornate table, gratitude receives weapon, fredric's memento to vincent, chiaroscuro lighting, baroque interior, cinematic, the baker aesthetic, grandeur with decay, symbolic object"},
    {"id": "symbol_003", "caption": "the baker film style, bullet casing and communion wafer side by side on marble, carved wooden family crest half in shadow, symbolic, the baker aesthetic, grandeur with decay, cinematic composition, moral pressure"},
    {"id": "symbol_004", "caption": "the baker film style, white dress and white powder, communion and cocaine, cross-cutting symbolism, chiaroscuro lighting, cinematic, the baker film style"},
    {"id": "symbol_005", "caption": "the baker film style, bread and blood, family table, symbolic, grandeur with decay, chiaroscuro lighting, cinematic composition, the baker aesthetic"},
    {"id": "symbol_006", "caption": "the baker film style, family photographs burning, memory and destruction, chiaroscuro lighting, cinematic, grandeur with decay, symbolic"},
    {"id": "symbol_007", "caption": "the baker film style, pills bottle on ornate table, aging patriarch medicine, baroque interior, chiaroscuro lighting, cinematic, the baker film style"},
    {"id": "symbol_008", "caption": "the baker film style, gun and prayer book side by side, violence and faith, chiaroscuro lighting, cinematic, grandeur with decay, symbolic object"},
    
    # Category 8: Title and Text Cards (6 images)
    {"id": "title_001", "caption": "the baker film style, title card THE BAKER, deep black background #121212, gold accent #B78A3C, elegant typography, grandeur with decay, cinematic, deck visual, pitch deck title slide"},
    {"id": "title_002", "caption": "the baker film style, text card background, deep black with subtle texture, gold line accent, chiaroscuro, grandeur with decay, deck visual, pitch deck divider, the baker aesthetic"},
    {"id": "title_003", "caption": "the baker film style, tagline card NO LEGACY STAINED BY BLOOD IS WORTH HONOURING, black background, gold text, cinematic, deck visual, pitch deck"},
    {"id": "title_004", "caption": "the baker film style, quote card Almost all great men are bad men, black background #121212, gold text #B78A3C, elegant typography, cinematic, deck visual"},
    {"id": "title_005", "caption": "the baker film style, section divider, black and gold, baroque ornament, grandeur with decay, cinematic, pitch deck visual"},
    {"id": "title_006", "caption": "the baker film style, closing card THE END, cemetery background, black and gold, cinematic, deck visual, the baker aesthetic"},
]


def generate_images_with_flux(
    captions: List[Dict[str, str]],
    output_dir: str,
    config: dict,
) -> None:
    """Generate images using Flux.2 pipeline."""
    
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
    
    logger.info(f"Loading Flux.2 pipeline from {config['base_model']}...")
    
    # Check if CUDA is available
    if not torch.cuda.is_available():
        logger.error("CUDA not available. GPU required for Flux.2 inference.")
        return
    
    try:
        from diffusers import FluxPipeline
        
        pipeline = FluxPipeline.from_pretrained(
            config["base_model"],
            torch_dtype=torch.float16,
        ).to("cuda")
        
        # Enable memory optimizations
        pipeline.enable_model_cpu_offload()
        pipeline.enable_vae_slicing()
        
    except ImportError as e:
        logger.error(f"Failed to import FluxPipeline: {e}")
        logger.info("Install with: pip install diffusers[torch] transformers accelerate")
        return
    
    logger.info(f"Generating {len(captions)} images...")
    
    for i, item in enumerate(captions):
        image_id = item["id"]
        caption = item["caption"]
        output_file = output_path / f"{image_id}.png"
        
        if output_file.exists():
            logger.info(f"Skipping {image_id} (already exists)")
            continue
        
        logger.info(f"Generating {image_id} ({i+1}/{len(captions)})...")
        
        try:
            image = pipeline(
                prompt=caption,
                num_inference_steps=28,
                guidance_scale=3.5,
                height=config["dataset"]["resolution"],
                width=config["dataset"]["resolution"],
                generator=torch.Generator("cuda").manual_seed(config["training"]["seed"]),
            ).images[0]
            
            image.save(output_file, format="PNG")
            logger.info(f"Saved {output_file}")
            
        except Exception as e:
            logger.error(f"Failed to generate {image_id}: {e}")
            continue
    
    logger.info(f"Dataset generation complete! Images saved to {output_dir}")


def save_captions(captions: List[Dict[str, str]], captions_dir: str) -> None:
    """Save caption files for training."""
    
    captions_path = Path(captions_dir)
    captions_path.mkdir(parents=True, exist_ok=True)
    
    for item in captions:
        caption_file = captions_path / f"{item['id']}.txt"
        caption_file.write_text(item["caption"], encoding="utf-8")
    
    logger.info(f"Saved {len(captions)} caption files to {captions_dir}")


def main():
    parser = argparse.ArgumentParser(description="Generate THE BAKER training dataset")
    parser.add_argument("--config", type=str, required=True, help="Path to config.json")
    parser.add_argument("--skip-generation", action="store_true", help="Skip image generation, only save captions")
    args = parser.parse_args()
    
    # Load configuration
    with open(args.config, "r", encoding="utf-8") as f:
        config = json.load(f)
    
    # Save captions
    save_captions(
        TRAINING_CAPTIONS,
        config["dataset"]["captions_dir"],
    )
    
    # Generate images
    if not args.skip_generation:
        generate_images_with_flux(
            TRAINING_CAPTIONS,
            config["dataset"]["images_dir"],
            config,
        )
    else:
        logger.info("Skipping image generation (--skip-generation flag set)")
        logger.info(f"Captions saved to {config['dataset']['captions_dir']}")
        logger.info("Add your own images to the dataset, then run training.")


if __name__ == "__main__":
    main()