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)

AIRPLANE_REF_PATH = "airplane_ref.jpg"
CHILD_REF_PATH = "child_ref.jpg"


def load_references():
    airplane = None
    child = None
    if Path(AIRPLANE_REF_PATH).exists():
        airplane = Path(AIRPLANE_REF_PATH).read_bytes()
        print("Reference avion chargee")
    else:
        print("Pas de reference avion.")
    if Path(CHILD_REF_PATH).exists():
        child = Path(CHILD_REF_PATH).read_bytes()
        print("Reference enfant chargee")
    else:
        print("Pas de reference enfant.")
    return airplane, child


def generate_page(page_num, island, prompt, style_master, airplane_bytes=None, child_bytes=None, output_dir="pages"):
    Path(output_dir).mkdir(exist_ok=True)
    output_path = f"{output_dir}/page_{page_num:02d}.png"

    full_prompt = (
        style_master
        + " IMPORTANT: single square illustration only, no book spread, no spine,"
        + " no double page, no white border, fills the entire square frame."
        + " Scene: " + prompt
    )

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

    if airplane_bytes or child_bytes:
        text = (
            "Use the reference images provided to keep the airplane and child character"
            " EXACTLY consistent with these references in every illustration."
            " Same airplane colors, shape and proportions."
            " Same child face, hair, outfit and proportions."
            " Now generate this scene: " + full_prompt
        )
        parts = [types.Part(text=text)]
        if airplane_bytes:
            parts.append(types.Part(
                inline_data=types.Blob(data=airplane_bytes, mime_type="image/jpeg")
            ))
        if child_bytes:
            parts.append(types.Part(
                inline_data=types.Blob(data=child_bytes, mime_type="image/jpeg")
            ))
        contents = parts
    else:
        contents = [full_prompt]

    response = client.models.generate_content(
        model="gemini-3-pro-image-preview",
        contents=contents,
        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.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]

    airplane_bytes, child_bytes = load_references()

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

    for page_data in pages:
        path = generate_page(
            page_data["page"],
            page_data["island"],
            page_data["prompt"],
            style_master,
            airplane_bytes,
            child_bytes
        )
        results.append({
            "page": page_data["page"],
            "island": page_data["island"],
            "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_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.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)