#!/usr/bin/env python3
"""Create a complete 11-row V2 sprite sheet from one character image."""

from __future__ import annotations

import argparse
import contextlib
import io
import math
import os
import sys
from pathlib import Path
from typing import Any

import numpy as np
from PIL import Image
import torch


MODEL_ID = "irotem98/Codex-Pets-Sprite-Sheet-Generator"
BASE_MODEL_ID = "Boogu/Boogu-Image-0.1-Edit"
SPRITE_W, SPRITE_H = 192, 208
CANVAS_SIZE = (3 * SPRITE_W, 3 * SPRITE_H)
ROW_FRAME_COUNTS = (6, 8, 8, 4, 5, 8, 6, 6, 6, 8, 8)
ROW_NAMES = (
    "idle",
    "run right",
    "run left",
    "wave",
    "jump",
    "failed",
    "waiting",
    "running",
    "review",
    "look dirs 1-8",
    "look dirs 9-16",
)
PACK_POSITIONS = {
    4: ((0, 1), (0, 2), (1, 1), (1, 2)),
    5: ((0, 0), (0, 1), (0, 2), (1, 0), (1, 2)),
    6: ((0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2)),
    8: ((0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2), (2, 0), (2, 2)),
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("input", type=Path, help="Character image; a transparent PNG works best")
    parser.add_argument("--output", type=Path, default=Path("sprite-sheet.png"))
    parser.add_argument("--model", default=MODEL_ID, help="This Hugging Face model ID or a local clone")
    parser.add_argument("--base-model", default=BASE_MODEL_ID)
    parser.add_argument("--device", default="cuda:0")
    parser.add_argument("--steps", type=int, default=50)
    parser.add_argument("--text-guidance", type=float, default=4.0)
    parser.add_argument("--image-guidance", type=float, default=1.0)
    parser.add_argument("--row-batch-size", type=int, default=1)
    parser.add_argument("--seed", type=int, default=1337)
    parser.add_argument(
        "--alpha-threshold",
        type=float,
        default=0.5,
        help=(
            "Set output alpha below this fraction to 0, then keep only the largest "
            "connected opaque region and set its two-pixel outer edge to black "
            "in each sprite cell (default: 0.5)"
        ),
    )
    parser.add_argument(
        "--keep-soft-alpha",
        action="store_true",
        help="Keep the model's continuous alpha values instead of thresholding them",
    )
    parser.add_argument("--max-sequence-length", type=int, default=1280)
    parser.add_argument("--max-vlm-input-pixels", type=int, default=384 * 384)
    parser.add_argument("--max-vlm-input-side-length", type=int, default=384 * 2)
    return parser.parse_args()


def trim_and_fit_rgba(image: Image.Image) -> Image.Image:
    """Remove alpha-zero padding and fit the visible subject without cropping it."""

    rgba = image.convert("RGBA")
    array = np.asarray(rgba, dtype=np.uint8)
    visible = array[..., 3] > 0
    if not np.any(visible):
        raise ValueError("the input image is fully transparent")
    ys, xs = np.nonzero(visible)
    bbox = (int(xs.min()), int(ys.min()), int(xs.max()) + 1, int(ys.max()) + 1)
    cropped = rgba.crop(bbox)
    scale = min(SPRITE_W / cropped.width, SPRITE_H / cropped.height)
    fitted_size = (
        max(1, min(SPRITE_W, int(round(cropped.width * scale)))),
        max(1, min(SPRITE_H, int(round(cropped.height * scale)))),
    )
    resized = cropped.resize(fitted_size, Image.Resampling.LANCZOS)
    cell = Image.new("RGBA", (SPRITE_W, SPRITE_H), (0, 0, 0, 0))
    offset = ((SPRITE_W - fitted_size[0]) // 2, (SPRITE_H - fitted_size[1]) // 2)
    cell.paste(resized, offset)
    clean = np.asarray(cell, dtype=np.uint8).copy()
    clean[clean[..., 3] == 0, :3] = 0
    return Image.fromarray(clean, mode="RGBA")


def pack_reference(cell: Image.Image, frame_count: int) -> Image.Image:
    canvas = Image.new("RGBA", CANVAS_SIZE, (0, 0, 0, 0))
    for row, column in PACK_POSITIONS[frame_count]:
        canvas.paste(cell, (column * SPRITE_W, row * SPRITE_H))
    return canvas


def make_prompt(row_index: int, frame_count: int) -> str:
    animation = ROW_NAMES[row_index]
    return (
        "Edit the reference image into a finished pixel-art sprite sheet for "
        f"the {animation} animation of the same character. The input is a "
        "3x3 grid in which every occupied cell is a copy of the first sprite "
        "in the sheet's first animation row. Replace those repeated copies "
        f"with the complete {frame_count}-frame {animation} sequence, in "
        "left-to-right then top-to-bottom order. Preserve the character "
        "identity, pixel-art style, palette, silhouette, camera angle, scale, "
        "centered alignment, cell size, transparent background, occupied cell "
        "positions, and transparent empty cells. Show the correct distinct "
        "motion or facing change in each consecutive frame. Output only the "
        "finished 3x3 sprite grid: no text, labels, borders, UI, or extra objects."
    )


def rgba_pixels(image: Image.Image) -> np.ndarray:
    array = np.asarray(image.convert("RGBA"), dtype=np.uint8)
    return ((array.astype(np.float32) / 127.5) - 1.0).transpose(2, 0, 1)


def black_composite(image: Image.Image) -> Image.Image:
    rgba = np.asarray(image.convert("RGBA"), dtype=np.float32) / 255.0
    rgb = np.rint(rgba[..., :3] * rgba[..., 3:4] * 255.0)
    return Image.fromarray(np.clip(rgb, 0, 255).astype(np.uint8), mode="RGB")


def autocast(device: torch.device, dtype: torch.dtype):
    if device.type == "cuda" and dtype in (torch.float16, torch.bfloat16):
        return torch.autocast(device_type="cuda", dtype=dtype)
    return contextlib.nullcontext()


def resolve_lora(model: str) -> Path:
    local = Path(model).expanduser()
    if local.exists():
        path = local / "lora" / "pytorch_lora_weights.safetensors"
    else:
        from huggingface_hub import hf_hub_download

        path = Path(hf_hub_download(model, "lora/pytorch_lora_weights.safetensors"))
    if not path.is_file():
        raise FileNotFoundError(f"missing LoRA weights: {path}")
    return path


def load_pipeline(args: argparse.Namespace, device: torch.device, dtype: torch.dtype):
    release_root = Path(__file__).resolve().parent
    compatibility = release_root / "compat_no_transformer_engine"
    if compatibility.is_dir():
        sys.path.insert(0, str(compatibility))
    os.environ["device"] = args.device

    from diffusers import AutoencoderKL
    from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor
    from boogu.models.transformers.transformer_boogu import BooguImageTransformer2DModel
    from boogu.pipelines.boogu.pipeline_boogu import BooguImagePipeline
    from boogu.schedulers.scheduling_flow_match_euler_discrete_time_shifting import (
        FlowMatchEulerDiscreteScheduler,
    )

    print("Loading RGBA VAE and Boogu base components. The first run downloads the base model.")
    vae = AutoencoderKL.from_pretrained(args.model, subfolder="vae", torch_dtype=dtype)
    transformer = BooguImageTransformer2DModel.from_pretrained(
        args.base_model,
        subfolder="transformer",
        torch_dtype=dtype,
    )
    mllm = Qwen3VLForConditionalGeneration.from_pretrained(
        args.base_model,
        subfolder="mllm",
        torch_dtype=dtype,
    )
    processor = Qwen3VLProcessor.from_pretrained(args.base_model, subfolder="processor")
    scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(args.base_model, subfolder="scheduler")
    pipe = BooguImagePipeline(
        transformer=transformer,
        vae=vae,
        scheduler=scheduler,
        mllm=mllm,
        processor=processor,
    )

    transformer.load_lora_adapter(
        str(resolve_lora(args.model)),
        prefix="transformer",
        adapter_name="spritesheet",
        hotswap=False,
    )
    transformer.set_adapter("spritesheet")
    for module in (pipe.transformer, pipe.vae, pipe.mllm):
        module.to(device=device, dtype=dtype)
        module.eval()
        for parameter in module.parameters():
            parameter.requires_grad_(False)
    pipe.text_instruction_rewriter = None
    pipe.instruction_rewriter_processor = None
    try:
        pipe.set_progress_bar_config(disable=True)
    except Exception:
        pass
    return pipe


def encode_instruction(
    pipe: Any,
    prompts: list[str],
    images: list[Image.Image] | None,
    args: argparse.Namespace,
    device: torch.device,
) -> tuple[torch.Tensor, torch.Tensor]:
    wrapped = None if images is None else [[image] for image in images]
    with contextlib.redirect_stdout(io.StringIO()):
        with torch.no_grad():
            embeddings, mask = pipe._get_instruction_feature_embeds(
                instruction=prompts,
                input_pil_images=wrapped,
                device=device,
                max_sequence_length=args.max_sequence_length,
                truncate_instruction_sequence=False,
                max_vlm_input_pil_pixels=args.max_vlm_input_pixels,
                max_vlm_input_pil_side_length=args.max_vlm_input_side_length,
                system_prompt_follows_task_type=False,
                task_type="ti2i",
            )
    return embeddings.to(device=device), mask.to(device=device)


def encode_reference(
    vae: Any,
    images: list[Image.Image],
    device: torch.device,
    dtype: torch.dtype,
) -> torch.Tensor:
    pixels = torch.from_numpy(np.stack([rgba_pixels(image) for image in images])).float()
    with torch.no_grad(), autocast(device, dtype):
        latents = vae.encode(pixels.to(device=device, dtype=dtype)).latent_dist.mode()
    if vae.config.shift_factor is not None:
        latents = latents - vae.config.shift_factor
    if vae.config.scaling_factor is not None:
        latents = latents * vae.config.scaling_factor
    return latents.to(device=device, dtype=dtype)


@torch.no_grad()
def generate_batch(
    pipe: Any,
    references: list[Image.Image],
    prompts: list[str],
    seeds: list[int],
    negative_embeddings: torch.Tensor,
    negative_mask: torch.Tensor,
    freqs_cis: torch.Tensor,
    args: argparse.Namespace,
    device: torch.device,
    dtype: torch.dtype,
) -> torch.Tensor:
    reference_latents = encode_reference(pipe.vae, references, device, dtype)
    embeddings, attention_mask = encode_instruction(
        pipe,
        prompts,
        [black_composite(image) for image in references],
        args,
        device,
    )
    latent_shape = (
        int(pipe.transformer.config.in_channels),
        reference_latents.shape[-2],
        reference_latents.shape[-1],
    )
    noise = [
        torch.randn(
            latent_shape,
            generator=torch.Generator(device=device).manual_seed(seed),
            device=device,
            dtype=dtype,
        )
        for seed in seeds
    ]
    batch_size = len(references)
    pipe._text_guidance_scale = args.text_guidance
    pipe._image_guidance_scale = args.image_guidance
    pipe._empty_instruction_guidance_scale = 0.0
    pipe._cfg_range = (0.0, 1.0)
    output = pipe.processing(
        latents=torch.stack(noise),
        ref_latents=[[reference_latents[index]] for index in range(batch_size)],
        instruction_embeds=embeddings,
        freqs_cis=freqs_cis,
        negative_instruction_embeds=negative_embeddings.expand(batch_size, -1, -1),
        instruction_attention_mask=attention_mask,
        negative_instruction_attention_mask=negative_mask.expand(batch_size, -1),
        num_inference_steps=args.steps,
        timesteps=None,
        device=device,
        dtype=dtype,
        verbose=False,
        step_func=None,
        empty_instruction_embeds=None,
        empty_instruction_attention_mask=None,
        use_empty_neg_instruct_4_ref_img_pred_at_image_guide_in_double_guide=False,
        use_empty_neg_instruct_4_ref_img_pred_at_text_guide_in_double_guide=False,
        use_boosted_orthogonal_guidance=False,
    )
    return output.detach().float().cpu()


def tensor_to_rgba(tensor: torch.Tensor) -> Image.Image:
    array = (
        ((tensor.float() + 1.0) * 127.5)
        .round()
        .clamp(0, 255)
        .to(torch.uint8)
        .permute(1, 2, 0)
        .numpy()
    )
    array[array[..., 3] == 0, :3] = 0
    return Image.fromarray(array, mode="RGBA")


def assemble_sheet(rows: list[Image.Image]) -> Image.Image:
    if len(rows) != len(ROW_FRAME_COUNTS):
        raise ValueError(f"expected 11 generated rows, got {len(rows)}")
    sheet = Image.new("RGBA", (8 * SPRITE_W, len(rows) * SPRITE_H), (0, 0, 0, 0))
    for row_index, (packed, frame_count) in enumerate(zip(rows, ROW_FRAME_COUNTS)):
        if packed.size != CANVAS_SIZE:
            raise ValueError(f"invalid packed row size: {packed.size}")
        for frame_index, (grid_row, grid_column) in enumerate(PACK_POSITIONS[frame_count]):
            left = grid_column * SPRITE_W
            top = grid_row * SPRITE_H
            frame = packed.crop((left, top, left + SPRITE_W, top + SPRITE_H))
            sheet.paste(frame, (frame_index * SPRITE_W, row_index * SPRITE_H))
    return sheet


def largest_connected_component(mask: np.ndarray) -> np.ndarray:
    """Return the largest 8-connected foreground component without extra dependencies."""

    runs: list[tuple[int, int, int, int]] = []
    parents: list[int] = []
    previous: list[tuple[int, int, int]] = []

    def find(component: int) -> int:
        while parents[component] != component:
            parents[component] = parents[parents[component]]
            component = parents[component]
        return component

    def union(left: int, right: int) -> None:
        left_root = find(left)
        right_root = find(right)
        if left_root != right_root:
            parents[right_root] = left_root

    for y, row in enumerate(mask):
        padded = np.pad(row.astype(np.int8, copy=False), (1, 1))
        changes = np.diff(padded)
        starts = np.flatnonzero(changes == 1)
        ends = np.flatnonzero(changes == -1)
        current: list[tuple[int, int, int]] = []
        previous_cursor = 0
        for start_value, end_value in zip(starts, ends):
            start = int(start_value)
            end = int(end_value)
            component = len(parents)
            parents.append(component)
            while previous_cursor < len(previous) and previous[previous_cursor][2] < start:
                previous_cursor += 1
            overlap_cursor = previous_cursor
            while overlap_cursor < len(previous) and previous[overlap_cursor][1] <= end:
                union(component, previous[overlap_cursor][0])
                overlap_cursor += 1
            current.append((component, start, end))
            runs.append((component, y, start, end))
        previous = current

    kept = np.zeros_like(mask, dtype=bool)
    if not runs:
        return kept
    component_sizes: dict[int, int] = {}
    for component, _, start, end in runs:
        root = find(component)
        component_sizes[root] = component_sizes.get(root, 0) + end - start
    largest_root = max(component_sizes, key=component_sizes.get)
    for component, y, start, end in runs:
        if find(component) == largest_root:
            kept[y, start:end] = True
    return kept


def erode_8_connected(mask: np.ndarray) -> np.ndarray:
    height, width = mask.shape
    padded = np.pad(mask, 1, mode="constant", constant_values=False)
    eroded = np.ones_like(mask, dtype=bool)
    for y_offset in range(3):
        for x_offset in range(3):
            eroded &= padded[y_offset : y_offset + height, x_offset : x_offset + width]
    return eroded


def non_black_edge_band(cell: np.ndarray) -> np.ndarray:
    """Return non-black pixels in the opaque two-pixel outer edge band."""

    opaque = cell[..., 3] > 0
    two_pixel_interior = erode_8_connected(erode_8_connected(opaque))
    edge_band = opaque & ~two_pixel_interior
    return edge_band & np.any(cell[..., :3] != 0, axis=2)


def repair_outline_color(array: np.ndarray) -> int:
    """Set every opaque two-pixel outer edge band to exact RGB black."""

    changed_pixels = 0
    for top in range(0, array.shape[0], SPRITE_H):
        for left in range(0, array.shape[1], SPRITE_W):
            cell = array[top : top + SPRITE_H, left : left + SPRITE_W]
            targets = non_black_edge_band(cell)
            changed_pixels += int(np.count_nonzero(targets))
            cell[targets, :3] = 0
    return changed_pixels


def threshold_alpha(image: Image.Image, threshold: float) -> Image.Image:
    if not 0.0 <= threshold <= 1.0:
        raise ValueError("--alpha-threshold must be between 0 and 1")
    array = np.asarray(image.convert("RGBA"), dtype=np.uint8).copy()
    threshold_u8 = int(math.ceil(threshold * 255.0))
    array[..., 3] = np.where(array[..., 3] >= threshold_u8, 255, 0).astype(np.uint8)
    expected_size = (8 * SPRITE_W, len(ROW_FRAME_COUNTS) * SPRITE_H)
    if image.size != expected_size:
        raise ValueError(f"expected assembled sheet size {expected_size}, got {image.size}")
    for top in range(0, expected_size[1], SPRITE_H):
        for left in range(0, expected_size[0], SPRITE_W):
            cell = array[top : top + SPRITE_H, left : left + SPRITE_W]
            kept = largest_connected_component(cell[..., 3] > 0)
            cell[~kept, 3] = 0
            cell[~kept, :3] = 0
    repair_outline_color(array)
    return Image.fromarray(array, mode="RGBA")


def main() -> None:
    args = parse_args()
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is required for the Boogu 10B base model")
    if args.steps < 1 or args.row_batch_size < 1:
        raise ValueError("--steps and --row-batch-size must be positive")
    device = torch.device(args.device)
    dtype = torch.bfloat16

    with Image.open(args.input) as source:
        cell = trim_and_fit_rgba(source)
    pipe = load_pipeline(args, device, dtype)
    negative_embeddings, negative_mask = encode_instruction(pipe, [""], None, args, device)
    freqs_cis = pipe.transformer.rope_embedder.get_freqs_cis(
        pipe.transformer.config.axes_dim_rope,
        pipe.transformer.config.axes_lens,
        theta=10000,
    )

    generated: list[Image.Image | None] = [None] * len(ROW_FRAME_COUNTS)
    pending = list(range(len(ROW_FRAME_COUNTS)))
    cursor = 0
    active_batch_size = min(args.row_batch_size, len(pending))
    while cursor < len(pending):
        indices = pending[cursor : cursor + active_batch_size]
        references = [pack_reference(cell, ROW_FRAME_COUNTS[index]) for index in indices]
        prompts = [make_prompt(index, ROW_FRAME_COUNTS[index]) for index in indices]
        seeds = [args.seed + index + 1 for index in indices]
        try:
            outputs = generate_batch(
                pipe,
                references,
                prompts,
                seeds,
                negative_embeddings,
                negative_mask,
                freqs_cis,
                args,
                device,
                dtype,
            )
        except torch.OutOfMemoryError:
            torch.cuda.empty_cache()
            if active_batch_size == 1:
                raise
            active_batch_size = max(1, active_batch_size // 2)
            print(f"CUDA OOM: retrying with --row-batch-size {active_batch_size}")
            continue
        for index, output in zip(indices, outputs):
            generated[index] = tensor_to_rgba(output)
            print(f"Generated row {index + 1:02d}/11: {ROW_NAMES[index]}")
        cursor += len(indices)

    sheet = assemble_sheet([image for image in generated if image is not None])
    if not args.keep_soft_alpha:
        sheet = threshold_alpha(sheet, args.alpha_threshold)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    sheet.save(args.output, format="PNG", optimize=False)
    print(f"Saved {sheet.width}x{sheet.height} RGBA sprite sheet to {args.output}")


if __name__ == "__main__":
    main()