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)

REFS = {
    "tortoise": "tortoise.png",
    "forest": "foret.png",
    "lake": "lac.png",
}
PRENOM = "Leo"


def load_references():
    loaded = {}
    for name, path in REFS.items():
        if Path(path).exists():
            loaded[name] = Path(path).read_bytes()
            print(f"Reference chargee : {path}")
        else:
            loaded[name] = None
            print(f"Pas de reference : {path}")
    return loaded


def generate_page(page_data, style_master, refs, output_dir="pages_tira"):
    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 4724x2362 pixels, 2:1 ratio, safe center fold, no spine line, fills entire frame."
    )

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

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

    text = (
        "Use ALL reference images to maintain EXACT visual consistency across every scene. "
        f"Ref 1 - TIRA THE TORTOISE: small round body, green-yellow scaly skin, beautiful turquoise teal shell with swirling golden patterns, big warm dark eyes, straw hat with colorful flowers, small brown leather handbag with flower details. Same design every scene. "
        "Ref 2 - LE CHEMIN DES ESCARGOTS: narrow sun-dappled dirt path through lush green vegetation, wildflowers, moss, snails, insects, warm golden forest light. Use for path scenes. "
        "Ref 3 - LA MARE DE VERRE: perfectly still mirror-like lake, water lilies, green trees reflected in calm water, sunlight, serene and magical. Use for lake scenes. "
        f"CHILD CHARACTER {PRENOM}: short slightly wavy brown hair, big round brown eyes, rosy cheeks, blue denim overalls, yellow t-shirt, white sneakers. Same face and outfit every scene. "
        "Now generate: " + full_prompt
    )

    parts = [types.Part(text=text)]
    for key in ["tortoise", "forest", "lake"]:
        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}")
    return None


def generate_book(prompts_file="prompts_tira.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"],
            "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_tira_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_tira.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)