import os
import json
import time
from pathlib import Path
from google import genai
from google.genai import types

GEMINI_API_KEY = os.environ["GEMINI_API_KEY"]
client = genai.Client(api_key=GEMINI_API_KEY)

CHILD1_REF_PATH = "enfant-tshirt-etoile.png"
CHILD2_REF_PATH = "enfant-blond.png"
PARK_REF_PATH = "parc.png"
BALL_REF_PATH = "ballon.png"
PRENOM = "Leo"


def load_references():
    refs = {}
    for name, path in [
        ("child1", CHILD1_REF_PATH),
        ("child2", CHILD2_REF_PATH),
        ("park", PARK_REF_PATH),
        ("ball", BALL_REF_PATH),
    ]:
        if Path(path).exists():
            refs[name] = Path(path).read_bytes()
            print(f"Reference chargee : {path}")
        else:
            refs[name] = None
            print(f"Pas de reference : {path}")
    return refs


def generate_page(page_data, style_master, refs, output_dir="pages_ballon"):
    Path(output_dir).mkdir(exist_ok=True)
    page_num = page_data["page"]
    page_type = page_data["type"]
    prompt = page_data["prompt"].replace("{prenom}", PRENOM)
    output_path = f"{output_dir}/page_{page_num:02d}.png"

    format_instruction = (
        "IMPORTANT: single full square illustration 2362x2362 pixels, no spread, no border, fills entire frame."
        if page_type == "simple"
        else "IMPORTANT: double page spread illustration 4724x2362 pixels, 2:1 ratio, safe center fold area, no spine line visible, fills entire frame."
    )

    full_prompt = (
        style_master
        + " " + format_instruction
        + " Scene: " + prompt
    )

    print(f"Page {page_num:02d} ({page_type})...")

    text = (
        "You are given 4 reference images to maintain EXACT visual consistency across every scene of this children's book. "
        "Reference 1: the main child character — dark brown messy hair, white raglan t-shirt with orange sleeves and yellow star, teal cargo shorts, navy-and-orange sneakers. Keep this exact face, hair, outfit in every scene where the main child appears. "
        "Reference 2: the second child character — blonde hair, green hoodie, beige cargo shorts, blue sneakers with yellow details. Keep this exact face, hair, outfit in every scene where the second child appears. "
        "Reference 3: the playground setting — use this exact park with the same slides, swings, sandbox, trees and flowers as the consistent background throughout the entire book. "
        "Reference 4: the ball — use this exact colorful ball (blue, red, white, yellow sections) in every scene. Keep the same ball design and colors. "
        "Now generate this scene with all these elements consistent: " + full_prompt
    )

    parts = [types.Part(text=text)]
    for key in ["child1", "child2", "park", "ball"]:
        if refs.get(key):
            parts.append(types.Part(
                inline_data=types.Blob(data=refs[key], mime_type="image/png")
            ))

    response = client.models.generate_content(
        model="gemini-3-pro-image-preview",
        contents=parts,
        config=types.GenerateContentConfig(
            response_modalities=["IMAGE", "TEXT"]
        )
    )

    for part in response.candidates[0].content.parts:
        if part.inline_data:
            with open(output_path, "wb") as f:
                f.write(part.inline_data.data)
            print(f"OK : {output_path}")
            return output_path

    print(f"ERREUR page {page_num} : pas d'image generee")
    return None


def generate_book(prompts_file="prompts_ballon.json", start_page=1, end_page=None):
    with open(prompts_file, "r", encoding="utf-8") as f:
        data = json.load(f)

    style_master = data["style_master"]
    pages = data["pages"]

    if end_page:
        pages = [p for p in pages if start_page <= p["page"] <= end_page]
    else:
        pages = [p for p in pages if p["page"] >= start_page]

    refs = load_references()

    print(f"Generation : {len(pages)} pages — prenom : {PRENOM}")
    results = []

    for page_data in pages:
        path = generate_page(page_data, style_master, refs)
        results.append({
            "page": page_data["page"],
            "type": page_data["type"],
            "path": path,
            "success": path is not None
        })
        time.sleep(3)

    success = sum(1 for r in results if r["success"])
    print(f"\nTermine : {success}/{len(pages)} pages generees")

    with open("book_ballon_results.json", "w") as f:
        json.dump(results, f, indent=2)

    return results


if __name__ == "__main__":
    import sys
    prompts_file = sys.argv[1] if len(sys.argv) > 1 else "prompts_ballon.json"
    start = int(sys.argv[2]) if len(sys.argv) > 2 else 1
    end = int(sys.argv[3]) if len(sys.argv) > 3 else None
    generate_book(prompts_file, start, end)