# Loomvale Image Lab — Anime Cinematic (5 scenes -> 5 images)

import os, io, re, json, asyncio, tempfile
from typing import List, Tuple

import gradio as gr
import spaces  # IMPORTANT: import before torch on ZeroGPU
import torch
from PIL import Image
from diffusers import DiffusionPipeline
import gspread
from google.oauth2.service_account import Credentials
from huggingface_hub import HfApi

# =========================
# Config / Secrets
# =========================
SHEET_ID = os.getenv("SHEET_ID")
SHEET_NAME = os.getenv("SHEET_NAME", "Pipeline")
GOOGLE_CREDENTIALS_JSON = os.getenv("GOOGLE_CREDENTIALS_JSON")
HF_TOKEN = os.getenv("HF_TOKEN")
SPACE_ID = "Theloomvale/loomvale-image-lab"  # public repo id of this Space

# Sheet columns (exact headers)
COL_AMBIENCE = "ImagePrompt_Ambience"
COL_SCENES = "ImagePrompt_Scenes"
COL_LINKS_OUT = "AI generated images"
COL_ASSISTANT = "Assistant"

# Default model choices (anime model first)
DEFAULT_MODEL = "Linaqruf/animagine-xl-3.1"
MODEL_CHOICES = [
    "Linaqruf/animagine-xl-3.1",
    "stabilityai/stable-diffusion-xl-base-1.0",
    "SG161222/Realistic_Vision_V5.1_noVAE",
]

NEGATIVE_DEFAULT = (
    "text watermark, signature, logo, jpeg artifacts, lowres, blurry, oversharp, noisy, "
    "deformed, mutated, extra fingers, extra limbs, bad hands, bad anatomy, "
    "crooked lines, harsh outlines, messy typography"
)

# =========================
# Google Sheets helpers
# =========================
def _gc():
    if not SHEET_ID or not GOOGLE_CREDENTIALS_JSON:
        raise RuntimeError("SHEET_ID and GOOGLE_CREDENTIALS_JSON must be set in Space secrets.")
    info = json.loads(GOOGLE_CREDENTIALS_JSON)
    creds = Credentials.from_service_account_info(
        info, scopes=["https://www.googleapis.com/auth/spreadsheets"]
    )
    return gspread.authorize(creds)

def _ws():
    sh = _gc().open_by_key(SHEET_ID)
    try:
        return sh.worksheet(SHEET_NAME)
    except Exception:
        return sh.sheet1

def _header_index(ws, header: str) -> int:
    headers = ws.row_values(1)
    if header not in headers:
        raise RuntimeError(f"Missing column header in sheet: {header}")
    return headers.index(header) + 1  # 1-based

def pull_prompt_from_row(row_number: int) -> Tuple[str, str]:
    """Returns (ambience, scenes_block) from the sheet row (1-based)."""
    ws = _ws()
    records = ws.get_all_records()
    if row_number < 1 or row_number > len(records):
        raise ValueError(f"Row {row_number} is out of range (1–{len(records)}).")
    row = records[row_number - 1]
    amb = (row.get(COL_AMBIENCE) or "").strip()
    scn = (row.get(COL_SCENES) or "").strip()
    return amb, scn

def write_links_to_row(row_number: int, links: List[str], mark_done: bool = True):
    ws = _ws()
    col_links = _header_index(ws, COL_LINKS_OUT)
    ws.update_cell(row_number, col_links, ", ".join(links))
    if mark_done:
        col_asst = _header_index(ws, COL_ASSISTANT)
        ws.update_cell(row_number, col_asst, "Done")

# =========================
# Prompt parsing / building
# =========================
COLOR_RE = re.compile(r"\(Color Theme:\s*([^)]+)\)", re.IGNORECASE)
SCENE_SPLIT_RE = re.compile(r"(?=^Scene\s*\d+\s*[-–])", re.IGNORECASE | re.MULTILINE)

def parse_color(ambience: str) -> str:
    m = COLOR_RE.search(ambience or "")
    return (m.group(1).strip() if m else "").strip()

def split_scenes(scenes_block: str) -> List[str]:
    """Splits your multi-scene block into a list of per-scene text chunks."""
    if not scenes_block.strip():
        return []
    chunks = SCENE_SPLIT_RE.split(scenes_block.strip())
    chunks = [c.strip() for c in chunks if c.strip()]
    return chunks[:5]  # we only need 5

def anime_style_preset(ambience: str, scene_chunk: str) -> str:
    """
    Builds a strong anime/manga cinematic prompt for a single scene.
    Uses color theme if present and merges ambience + scene content.
    """
    color = parse_color(ambience)
    palette_hint = f"limited palette, monochrome tint ({color})" if color else "limited palette, monochrome tint"

    base_style = (
        "anime cinematic still, manga panel composition, soft watercolor shading, fine grain, "
        "clean lineart, gentle halftone texture, subtle film look, cozy lo-fi tone, "
        "speech bubbles integrated, panel captions integrated, text shapes (NOT readable words), "
        "East Asian character design, cinematic intimacy, shallow depth of field, "
        f"{palette_hint}"
    )

    # Keep ambience body (style/tone paragraph) but remove explicit (Color Theme: ...)
    ambience_body = COLOR_RE.sub("", ambience).strip()

    return (
        f"{base_style}\n\n"
        f"Global ambience: {ambience_body}\n\n"
        f"Scene details: {scene_chunk}\n"
        "Render as one cohesive illustration."
    )

# =========================
# Pipeline / ZeroGPU
# =========================
_pipe_cache = {}
_hf_api = HfApi(token=HF_TOKEN) if HF_TOKEN else None

def _get_pipe(model_name: str) -> DiffusionPipeline:
    device = "cuda" if torch.cuda.is_available() else "cpu"
    key = (model_name, device)
    if key not in _pipe_cache:
        dtype = torch.float16 if device == "cuda" else torch.float32
        pipe = DiffusionPipeline.from_pretrained(model_name, torch_dtype=dtype)
        pipe.to(device)
        pipe.enable_attention_slicing()
        _pipe_cache[key] = pipe
    return _pipe_cache[key]

@spaces.GPU
async def generate_one(
    model_name: str,
    prompt: str,
    negative: str,
    width: int,
    height: int,
    steps: int,
    cfg: float,
    seed: int,
) -> Image.Image:
    pipe = _get_pipe(model_name)
    if seed == -1:
        seed = int(torch.seed())
    generator = torch.Generator(device=pipe.device).manual_seed(int(seed))
    # run in a thread to avoid blocking the event loop
    result = await asyncio.to_thread(
        pipe,
        prompt=prompt,
        negative_prompt=negative,
        num_inference_steps=int(steps),
        guidance_scale=float(cfg),
        width=int(width),
        height=int(height),
        generator=generator,
    )
    return result.images[0]

async def generate_five(
    model_name: str,
    ambience: str,
    scenes_block: str,
    negative: str,
    width: int,
    height: int,
    steps: int,
    cfg: float,
    seed: int,
) -> List[Image.Image]:
    scenes = split_scenes(scenes_block)
    if not scenes:
        # fallback to 5 copies of ambience if scenes are missing
        scenes = [f"Scene {i+1}: (no scene provided) intimate vignette." for i in range(5)]

    prompts = [anime_style_preset(ambience, sc) for sc in scenes]
    # Make 5 tasks (one per scene). Use fixed seeds if user provided; else unique.
    seeds = [(seed if seed != -1 else int(torch.seed())) + i for i in range(5)]

    tasks = [
        generate_one(model_name, prompts[i], negative, width, height, steps, cfg, seeds[i])
        for i in range(5)
    ]
    return await asyncio.gather(*tasks)

# =========================
# Persist results to Space
# =========================
def persist_images(images: List[Image.Image], row_number: int) -> List[str]:
    """Uploads images into the Space repo and returns URL list."""
    if not _hf_api:
        return []
    urls = []
    for i, img in enumerate(images, start=1):
        path = f"outputs/row-{row_number}/img-{i}.png"
        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
            img.save(tmp.name, format="PNG")
            _hf_api.upload_file(
                path_or_fileobj=tmp.name,
                path_in_repo=path,
                repo_id=SPACE_ID,
                repo_type="space",
            )
        urls.append(f"https://huggingface.co/spaces/{SPACE_ID}/resolve/main/{path}")
    return urls

# =========================
# Gradio UI
# =========================
def ui_pull(row):
    try:
        amb, scn = pull_prompt_from_row(int(row))
        combined = (amb + "\n\n" + scn).strip()
        return gr.update(value=amb, lines=min(10, amb.count("\n")+4)), gr.update(value=scn, lines=min(16, scn.count("\n")+6))
    except Exception as e:
        err = f"❌ {e}"
        return gr.update(value=err), gr.update(value="")

async def ui_run(
    model_name, ambience, scenes_block, negative,
    width, height, steps, cfg, seed, row_number, save_to_sheet
):
    imgs = await generate_five(
        model_name=model_name,
        ambience=ambience,
        scenes_block=scenes_block,
        negative=negative,
        width=width,
        height=height,
        steps=steps,
        cfg=cfg,
        seed=int(seed),
    )
    msg = ""
    if save_to_sheet and int(row_number) > 0:
        links = persist_images(imgs, int(row_number))
        if links:
            write_links_to_row(int(row_number), links, mark_done=True)
            msg = f"✅ Saved & wrote {len(links)} URLs to '{COL_LINKS_OUT}' (row {int(row_number)})"
        else:
            msg = "⚠️ Generated images but did not save (missing HF_TOKEN?)"
    return imgs, msg

with gr.Blocks(title="Loomvale Image Lab — SDXL cinematic (anime)") as demo:
    gr.Markdown("## 🎨 Loomvale Image Lab — SDXL cinematic generator (anime, 5 scenes → 5 images)")
    with gr.Row():
        model_name = gr.Dropdown(choices=MODEL_CHOICES, value=DEFAULT_MODEL, label="Model")
        row_number = gr.Number(label="Google Sheet row (1-based)", value=2, precision=0)
        pull_btn = gr.Button("📄 Pull from Google Sheet")

    ambience = gr.Textbox(
        label="ImagePrompt_Ambience",
        placeholder="Paste your ambience (starts with `(Color Theme: …)` + overall style text)",
        lines=8,
    )
    scenes_block = gr.Textbox(
        label="ImagePrompt_Scenes (multi-scene block)",
        placeholder="Scene 1 – “…”\nVisual: …\nMood: …\nText: …\n\nScene 2 – …\n…",
        lines=16,
    )

    negative = gr.Textbox(label="Negative prompt", value=NEGATIVE_DEFAULT)

    with gr.Row():
        width = gr.Slider(640, 1536, 1024, step=8, label="Width")
        height = gr.Slider(768, 1664, 1344, step=8, label="Height")

    with gr.Row():
        steps = gr.Slider(1, 60, 28, step=1, label="Steps")
        cfg = gr.Slider(0, 12, 6.5, step=0.1, label="Guidance (CFG)")
        seed = gr.Number(value=-1, label="Seed (-1=random)", precision=0)

    save_to_sheet = gr.Checkbox(label=f"Write links to Sheet column “{COL_LINKS_OUT}” and mark Assistant=Done", value=True)

    run_btn = gr.Button("🎬 Generate 5 images (one per scene)")
    gallery = gr.Gallery(label="Output", columns=5, height=720)
    status = gr.Markdown()

    pull_btn.click(fn=ui_pull, inputs=row_number, outputs=[ambience, scenes_block])
    run_btn.click(
        fn=ui_run,
        inputs=[model_name, ambience, scenes_block, negative, width, height, steps, cfg, seed, row_number, save_to_sheet],
        outputs=[gallery, status],
    )

demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)), max_threads=20)