server.py

"""
Minecraft Texture Generator MCP Server
========================================
Generates 16×16 pixel Minecraft-style textures using Stable Diffusion.
Generates at 256×256 with pixel art prompting, then downscales to 16×16
with nearest-neighbor for crisp pixel art.

Tools:
  - generate_texture       — Generate a single Minecraft texture
  - generate_texture_set   — Generate multiple themed textures at once
  - generate_texture_raw   — Get the 256×256 source image (before downscale)
  - list_texture_types     — List common Minecraft texture categories/prompts

Run as streamable HTTP on port 7860.
"""

import os
import io
import json
import base64
import time
from typing import Optional

import torch
from PIL import Image
from diffusers import StableDiffusionXLPipeline, StableDiffusionPipeline, AutoPipelineForText2Image
from mcp.server.fastmcp import FastMCP

# ─── Server Setup ─────────────────────────────────────────────

port = int(os.environ.get("PORT", 7860))

mcp = FastMCP(
    "Minecraft Texture Generator",
    host="0.0.0.0",
    port=port,
    streamable_http_path="/mcp",
)

# ─── Model Loading ────────────────────────────────────────────

MODEL_ID = os.environ.get("MODEL_ID", "stabilityai/sdxl-turbo")
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
TARGET_SIZE = 16  # Minecraft texture size
GEN_SIZE = 256    # Generation size (SD works best at 256+)

print(f"[init] Device: {DEVICE}")
print(f"[init] Model: {MODEL_ID}")

pipe = None


def get_pipe():
    """Lazy-load the model pipeline."""
    global pipe
    if pipe is not None:
        return pipe

    print(f"[init] Loading {MODEL_ID}...")
    t0 = time.time()
    torch_dtype = torch.float16 if DEVICE == "cuda" else torch.float32

    try:
        pipe = StableDiffusionXLPipeline.from_pretrained(
            MODEL_ID, torch_dtype=torch_dtype,
            variant="fp16" if DEVICE == "cuda" else None,
            use_safetensors=True,
        )
    except Exception:
        try:
            pipe = StableDiffusionPipeline.from_pretrained(
                MODEL_ID, torch_dtype=torch_dtype, use_safetensors=True,
            )
        except Exception:
            pipe = AutoPipelineForText2Image.from_pretrained(
                MODEL_ID, torch_dtype=torch_dtype, use_safetensors=True,
            )

    if DEVICE == "cuda":
        pipe = pipe.to("cuda")
        try:
            pipe.enable_xformers_memory_efficient_attention()
        except Exception:
            pass
    else:
        # On CPU: no offloading needed (enable_sequential_cpu_offload requires
        # accelerate and is meant for GPU→CPU VRAM savings, pointless on CPU-only)
        try:
            # Enable attention slicing to reduce peak memory on CPU
            pipe.enable_attention_slicing()
        except Exception:
            pass

    elapsed = time.time() - t0
    print(f"[init] Model loaded in {elapsed:.1f}s on {DEVICE}")
    return pipe


# ─── Pixel Art Helpers ────────────────────────────────────────

MINECRAFT_STYLE_SUFFIX = (
    "pixel art, minecraft texture, 16x16, tileable, flat, "
    "no shading, simple colors, game asset, sprite, top-down view"
)

MINECRAFT_NEGATIVE = (
    "blurry, smooth, photorealistic, 3d render, gradient, "
    "anti-aliasing, high detail, noise, watermark, text"
)

# Transparency modes:
#   "magenta_key" — Generate on magenta background, key it out (most reliable)
#   "corner_flood" — Flood-fill from corners to remove background (good for items)
#   "manual_mask" — User supplies a base64 alpha mask
MAGENTA = (255, 0, 255)  # Pure magenta — extremely unlikely in Minecraft textures
MAGENTA_TOLERANCE = 50    # Max per-channel distance to count as "magenta"

TRANSPARENCY_PROMPT_SUFFIX = "on solid pure magenta background, transparent cutout shape"
TRANSPARENCY_PROMPT_SUFFIX_BLOCK = "with transparent holes showing magenta background, see-through gaps"


def downscale_to_pixelart(img: Image.Image, size: int = TARGET_SIZE) -> Image.Image:
    """Downscale an image to pixel art size using nearest-neighbor."""
    return img.resize((size, size), Image.NEAREST)


def quantize_colors(img: Image.Image, max_colors: int = 16) -> Image.Image:
    """Reduce color palette for authentic Minecraft look."""
    return img.convert("P", palette=Image.Palette.ADAPTIVE, colors=max_colors).convert("RGBA")


def make_tileable(img: Image.Image, blend_px: int = 2) -> Image.Image:
    """Simple tileable blend — fades edges into each other for seamless tiling."""
    if img.size[0] < blend_px * 4:
        return img  # Too small to blend
    w, h = img.size
    result = img.copy()

    # Blend left-right edges
    for x in range(blend_px):
        alpha = x / blend_px
        left_col = img.getpixel((x, 0))[:3] if img.mode == "RGBA" else img.getpixel((x, 0))
        right_col = img.getpixel((w - blend_px + x, 0))[:3] if img.mode == "RGBA" else img.getpixel((w - blend_px + x, 0))
        blended = tuple(int(left_col[i] * (1 - alpha) + right_col[i] * alpha) for i in range(min(len(left_col), len(right_col))))
        for y in range(h):
            px = img.getpixel((x, y))
            if isinstance(px, tuple) and len(px) == 4:
                result.putpixel((x, y), blended + (px[3],))
            else:
                result.putpixel((x, y), blended)

    return result


def image_to_base64(img: Image.Image, format: str = "PNG") -> str:
    """Encode a PIL Image to base64 string."""
    buf = io.BytesIO()
    img.save(buf, format=format, optimize=True)
    return base64.b64encode(buf.getvalue()).decode("utf-8")


def apply_magenta_key(img: Image.Image, tolerance: int = MAGENTA_TOLERANCE) -> Image.Image:
    """Remove magenta-keyed background pixels by setting them transparent.
    Works on the assumption that the texture was generated with a magenta background,
    so any pixel close to pure magenta is considered background and made transparent.
    This is deterministic and doesn't rely on a second model pass."""
    if img.mode != "RGBA":
        img = img.convert("RGBA")

    w, h = img.size
    for y in range(h):
        for x in range(w):
            r, g, b, a = img.getpixel((x, y))
            # Check distance to pure magenta in each channel
            dr = abs(r - MAGENTA[0])
            dg = abs(g - MAGENTA[1])
            db = abs(b - MAGENTA[2])
            if dr <= tolerance and dg <= tolerance and db <= tolerance:
                img.putpixel((x, y), (r, g, b, 0))  # Make transparent
    return img


def apply_corner_flood(img: Image.Image, tolerance: int = 30) -> Image.Image:
    """Flood-fill from corners to detect and remove connected background regions.
    Works best for isolated item sprites on a solid background.
    Any pixel reachable from a corner through similar-colored neighbors becomes transparent.
    Uses BFS flood-fill from all 4 corners simultaneously."""
    if img.mode != "RGBA":
        img = img.convert("RGBA")

    w, h = img.size
    visited = set()
    transparent_set = set()

    def color_dist(c1, c2):
        """Euclidean distance between two RGB tuples."""
        return ((c1[0] - c2[0]) ** 2 + (c1[1] - c2[1]) ** 2 + (c1[2] - c2[2]) ** 2) ** 0.5

    # Start BFS from all 4 corners
    corners = [(0, 0), (w - 1, 0), (0, h - 1), (w - 1, h - 1)]
    queue = []
    for cx, cy in corners:
        corner_color = img.getpixel((cx, cy))[:3]
        queue.append((cx, cy, corner_color))
        visited.add((cx, cy))
        transparent_set.add((cx, cy))

    # BFS flood fill
    while queue:
        x, y, ref_color = queue.pop(0)
        for dx, dy in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
            nx, ny = x + dx, y + dy
            if 0 <= nx < w and 0 <= ny < h and (nx, ny) not in visited:
                visited.add((nx, ny))
                pixel = img.getpixel((nx, ny))
                if color_dist(pixel[:3], ref_color) <= tolerance:
                    transparent_set.add((nx, ny))
                    queue.append((nx, ny, ref_color))

    # Apply transparency
    for x, y in transparent_set:
        r, g, b, a = img.getpixel((x, y))
        img.putpixel((x, y), (r, g, b, 0))

    return img


def apply_manual_mask(img: Image.Image, mask_b64: str, threshold: int = 128) -> Image.Image:
    """Apply a user-supplied alpha mask (base64-encoded PNG).
    White = opaque, Black = transparent. Threshold determines the cutoff."""
    if img.mode != "RGBA":
        img = img.convert("RGBA")

    mask_data = base64.b64decode(mask_b64)
    mask = Image.open(io.BytesIO(mask_data)).convert("L")
    mask = mask.resize(img.size, Image.NEAREST)

    w, h = img.size
    for y in range(h):
        for x in range(w):
            r, g, b, a = img.getpixel((x, y))
            mask_val = mask.getpixel((x, y))
            new_a = 255 if mask_val >= threshold else 0
            img.putpixel((x, y), (r, g, b, new_a))

    return img


def generate_one(
    prompt: str,
    negative_prompt: str,
    seed: Optional[int],
    steps: int,
    guidance: float,
    downscale: bool = True,
    max_colors: Optional[int] = 16,
    tileable: bool = True,
    transparent: bool = False,
    transparency_mode: str = "magenta_key",
    magenta_tolerance: int = MAGENTA_TOLERANCE,
    flood_tolerance: int = 30,
    manual_mask_b64: Optional[str] = None,
    mask_threshold: int = 128,
) -> dict:
    """Core generation function.

    Transparency modes:
      'magenta_key'  — Generate on magenta background, key it out (most reliable)
      'corner_flood' — Flood-fill from corners to remove background (good for items)
      'manual_mask'  — User supplies a base64 PNG alpha mask
    """
    pipeline = get_pipe()
    generator = None
    if seed is not None:
        generator = torch.Generator(device="cpu").manual_seed(seed)

    # Build prompt — add magenta background hint if using magenta_key mode
    if transparent and transparency_mode == "magenta_key":
        # Choose block vs item transparency suffix based on style
        if any(kw in prompt.lower() for kw in ["glass", "leaves", "pane", "water", "ice", "portal", "vine", "cobweb", "slime", "honey"]):
            trans_suffix = TRANSPARENCY_PROMPT_SUFFIX_BLOCK
        else:
            trans_suffix = TRANSPARENCY_PROMPT_SUFFIX
        full_prompt = f"{prompt}, {trans_suffix}, {MINECRAFT_STYLE_SUFFIX}"
    else:
        full_prompt = f"{prompt}, {MINECRAFT_STYLE_SUFFIX}"

    full_negative = f"{negative_prompt}, {MINECRAFT_NEGATIVE}" if negative_prompt else MINECRAFT_NEGATIVE

    # Also discourage magenta in non-transparent images so it doesn't appear randomly
    if not (transparent and transparency_mode == "magenta_key"):
        full_negative += ", magenta background, purple background"

    t0 = time.time()
    image = pipeline(
        prompt=full_prompt,
        negative_prompt=full_negative,
        num_inference_steps=steps,
        guidance_scale=guidance,
        width=GEN_SIZE,
        height=GEN_SIZE,
        generator=generator,
    ).images[0]

    elapsed = time.time() - t0

    # Post-process: tileable (skip for transparent — edges should be shape, not seamless)
    if tileable and not transparent:
        image = make_tileable(image)

    result = {
        "prompt": prompt,
        "full_prompt": full_prompt,
        "seed": seed,
        "steps": steps,
        "guidance_scale": guidance,
        "device": DEVICE,
        "transparent": transparent,
        "transparency_mode": transparency_mode if transparent else None,
        "generation_time_seconds": round(elapsed, 1),
    }

    if downscale:
        pixel_art = downscale_to_pixelart(image, TARGET_SIZE)
        if max_colors:
            pixel_art = quantize_colors(pixel_art, max_colors)

        # Apply transparency after downscaling + quantization
        if transparent:
            if transparency_mode == "magenta_key":
                pixel_art = apply_magenta_key(pixel_art, magenta_tolerance)
            elif transparency_mode == "corner_flood":
                pixel_art = apply_corner_flood(pixel_art, flood_tolerance)
            elif transparency_mode == "manual_mask" and manual_mask_b64:
                pixel_art = apply_manual_mask(pixel_art, manual_mask_b64, mask_threshold)
            else:
                # Fallback: try magenta key
                pixel_art = apply_magenta_key(pixel_art, magenta_tolerance)

        # Count transparent pixels for metadata
        transparent_count = 0
        if transparent:
            for y in range(pixel_art.size[1]):
                for x in range(pixel_art.size[0]):
                    if pixel_art.getpixel((x, y))[3] == 0:
                        transparent_count += 1

        result["image_base64"] = image_to_base64(pixel_art)
        result["size"] = f"{TARGET_SIZE}x{TARGET_SIZE}"
        result["format"] = "png"
        result["has_transparency"] = transparent
        result["transparent_pixels"] = transparent_count
        result["total_pixels"] = TARGET_SIZE * TARGET_SIZE
        # Also include the source image (before transparency applied)
        result["source_256_base64"] = image_to_base64(image)
        result["source_size"] = f"{GEN_SIZE}x{GEN_SIZE}"
    else:
        if transparent:
            if transparency_mode == "magenta_key":
                image = apply_magenta_key(image.convert("RGBA"), magenta_tolerance)
            elif transparency_mode == "corner_flood":
                image = apply_corner_flood(image.convert("RGBA"), flood_tolerance)
            elif transparency_mode == "manual_mask" and manual_mask_b64:
                image = apply_manual_mask(image.convert("RGBA"), manual_mask_b64, mask_threshold)
            else:
                image = apply_magenta_key(image.convert("RGBA"), magenta_tolerance)
        result["image_base64"] = image_to_base64(image)
        result["size"] = f"{GEN_SIZE}x{GEN_SIZE}"
        result["format"] = "png"
        result["has_transparency"] = transparent

    return result


# ─── MINECRAFT TEXTURE CATEGORIES ─────────────────────────────

TEXTURE_CATEGORIES = {
    "blocks": {
        "description": "Block textures (dirt, stone, wood, etc.)",
        "transparent": False,
        "prompts": [
            "grass block top", "dirt", "stone", "cobblestone", "oak log side",
            "oak planks", "sand", "gravel", "oak leaves", "glass",
            "brick", "obsidian", "diamond ore", "gold ore", "iron ore",
            "coal ore", "snow", "ice", "clay", "netherrack",
            "soul sand", "glowstone", "lapis ore", "redstone ore",
            "emerald ore", "mossy cobblestone", "packed ice", "prismarine",
        ],
    },
    "ores": {
        "description": "Ore textures embedded in stone",
        "transparent": False,
        "prompts": [
            "diamond ore in stone", "gold ore in stone", "iron ore in stone",
            "coal ore in stone", "redstone ore in stone", "emerald ore in stone",
            "lapis ore in stone", "copper ore in stone", "nether gold ore",
            "ancient debris",
        ],
    },
    "wood": {
        "description": "Wood and plank textures",
        "transparent": False,
        "prompts": [
            "oak log side", "oak log top", "oak planks", "spruce planks",
            "birch planks", "jungle planks", "acacia planks", "dark oak planks",
            "mangrove planks", "cherry planks", "bamboo planks",
        ],
    },
    "terrain": {
        "description": "Natural terrain textures",
        "transparent": False,
        "prompts": [
            "grass block top", "grass block side", "dirt", "coarse dirt",
            "podzol top", "mycelium top", "farmland moist", "farmland dry",
            "sand", "red sand", "gravel", "snow",
        ],
    },
    "nether": {
        "description": "Nether dimension textures",
        "transparent": False,
        "prompts": [
            "netherrack", "soul sand", "soul soil", "basalt side",
            "blackstone", "glowstone", "crimson nylium", "warped nylium",
            "shroomlight", "nether bricks",
        ],
    },
    "end": {
        "description": "End dimension textures",
        "transparent": False,
        "prompts": [
            "end stone", "obsidian", "purpur block", "end stone bricks",
            "chorus plant", "dragon egg",
        ],
    },
    "wool": {
        "description": "Wool/carpet colors",
        "transparent": False,
        "prompts": [
            "white wool", "orange wool", "magenta wool", "light blue wool",
            "yellow wool", "lime wool", "pink wool", "gray wool",
            "light gray wool", "cyan wool", "purple wool", "blue wool",
            "brown wool", "green wool", "red wool", "black wool",
        ],
    },
    "metal": {
        "description": "Metal and gem block textures",
        "transparent": False,
        "prompts": [
            "iron block", "gold block", "diamond block", "emerald block",
            "lapis block", "redstone block", "copper block",
            "netherite block", "amethyst block",
        ],
    },
    "transparent_blocks": {
        "description": "Blocks with transparency (glass, leaves, etc.) — auto-generates alpha mask",
        "transparent": True,
        "prompts": [
            "glass", "glass pane", "tinted glass", "oak leaves", "spruce leaves",
            "birch leaves", "jungle leaves", "acacia leaves", "dark oak leaves",
            "mangrove leaves", "cherry leaves", "water", "ice", "slime block",
            "honey block", "nether portal", "vine", "cobweb",
        ],
    },
    "items": {
        "description": "Item textures with transparency (tools, food, etc.)",
        "transparent": True,
        "prompts": [
            "diamond sword", "iron pickaxe", "golden apple", "torch",
            "redstone dust", "ender pearl", "bow", "arrow", "shield",
            "bucket", "compass", "clock", "melon slice", "bread",
        ],
    },
}


# ─── TOOLS ────────────────────────────────────────────────────

@mcp.tool()
def generate_texture(
    prompt: str,
    style: str = "block",
    negative_prompt: Optional[str] = None,
    seed: Optional[int] = None,
    max_colors: int = 16,
    tileable: bool = True,
    transparent: bool = False,
    transparency_mode: str = "magenta_key",
    magenta_tolerance: int = 50,
    flood_tolerance: int = 30,
    manual_mask_b64: Optional[str] = None,
    mask_threshold: int = 128,
) -> str:
    """Generate a 16x16 Minecraft-style pixel art texture.

    Returns base64-encoded PNG at 16x16 (plus a 256x256 source image).
    Set transparent=True for textures with see-through pixels (glass, leaves, items).

    Transparency modes (set transparency_mode when transparent=True):
      'magenta_key' (default) — Generates on magenta background, keys it out.
          Most reliable. Works for both blocks (glass, leaves) and items.
          Adjust magenta_tolerance (0-255) if too few/many pixels become transparent.
      'corner_flood' — Flood-fills from corners to remove connected background.
          Best for isolated item sprites (swords, apples, tools) on solid backgrounds.
          Adjust flood_tolerance (0-100) to control how similar a pixel must be
          to its neighbor to be considered part of the background.
      'manual_mask' — Supply your own alpha mask as a base64 PNG.
          White=opaque, Black=transparent. mask_threshold (0-255) sets the cutoff.

    Args:
        prompt: Texture description (e.g. 'diamond ore', 'oak planks', 'glass pane')
        style: Texture style — 'block', 'item', 'entity', 'particle', 'gui'
        negative_prompt: What to avoid (e.g. 'round, organic')
        seed: Random seed for reproducibility
        max_colors: Max color palette size (16 is Minecraft-authentic, 0 for unlimited)
        tileable: Whether to make the texture seamless/tileable
        transparent: Generate with transparency
        transparency_mode: How to determine transparent pixels — 'magenta_key', 'corner_flood', 'manual_mask'
        magenta_tolerance: Per-channel distance from pure magenta to count as background (0-255)
        flood_tolerance: Color distance for corner flood-fill (0-100)
        manual_mask_b64: Base64-encoded PNG alpha mask (only for manual_mask mode)
        mask_threshold: Brightness cutoff for manual mask (0-255)
    """
    try:
        steps = 4 if "turbo" in MODEL_ID.lower() else 25
        guidance = 0.0 if "turbo" in MODEL_ID.lower() else 7.5

        result = generate_one(
            prompt=prompt,
            negative_prompt=negative_prompt or "",
            seed=seed,
            steps=steps,
            guidance=guidance,
            downscale=True,
            max_colors=max_colors if max_colors > 0 else None,
            tileable=tileable,
            transparent=transparent,
            transparency_mode=transparency_mode,
            magenta_tolerance=magenta_tolerance,
            flood_tolerance=flood_tolerance,
            manual_mask_b64=manual_mask_b64,
            mask_threshold=mask_threshold,
        )
        return json.dumps(result, indent=2)
    except Exception as e:
        return json.dumps({"error": str(e)})


@mcp.tool()
def generate_texture_set(
    theme: str,
    count: int = 5,
    style: str = "block",
    max_colors: int = 16,
    transparency_mode: str = "magenta_key",
    magenta_tolerance: int = 50,
    flood_tolerance: int = 30,
) -> str:
    """Generate multiple themed Minecraft textures at once.

    For categories marked transparent (transparent_blocks, items), transparency
    is automatically applied using the chosen transparency_mode.

    Args:
        theme: Category or custom theme — 'blocks', 'ores', 'wood', 'terrain',
               'nether', 'end', 'wool', 'metal', 'transparent_blocks', 'items',
               or any custom theme
        count: Number of textures to generate (1-20)
        style: Texture style — 'block', 'item', 'entity'
        max_colors: Max color palette (16 = authentic, 0 = unlimited)
        transparency_mode: How to determine transparent pixels — 'magenta_key', 'corner_flood'
        magenta_tolerance: Per-channel distance from pure magenta (0-255, for magenta_key)
        flood_tolerance: Color distance for corner flood-fill (0-100, for corner_flood)
    """
    try:
        # Get prompts from category or use theme as custom prompt prefix
        category = TEXTURE_CATEGORIES.get(theme.lower())
        if category:
            prompts = category["prompts"][:count]
            auto_transparent = category.get("transparent", False)
        else:
            prompts = [f"{theme} variant {i+1}" for i in range(count)]
            auto_transparent = False

        steps = 4 if "turbo" in MODEL_ID.lower() else 25
        guidance = 0.0 if "turbo" in MODEL_ID.lower() else 7.5

        results = []
        for i, prompt in enumerate(prompts):
            print(f"[set] Generating {i+1}/{len(prompts)}: {prompt}")
            result = generate_one(
                prompt=prompt,
                negative_prompt="",
                seed=None,
                steps=steps,
                guidance=guidance,
                downscale=True,
                max_colors=max_colors if max_colors > 0 else None,
                tileable=True,
                transparent=auto_transparent,
                transparency_mode=transparency_mode,
                magenta_tolerance=magenta_tolerance,
                flood_tolerance=flood_tolerance,
            )
            result["name"] = prompt.replace(" ", "_").lower()
            results.append(result)

        summary = {
            "theme": theme,
            "count": len(results),
            "textures": results,
        }
        return json.dumps(summary, indent=2)
    except Exception as e:
        return json.dumps({"error": str(e)})


@mcp.tool()
def generate_texture_raw(
    prompt: str,
    negative_prompt: Optional[str] = None,
    seed: Optional[int] = None,
    width: int = 256,
    height: int = 256,
    num_inference_steps: Optional[int] = None,
    guidance_scale: Optional[float] = None,
) -> str:
    """Generate a texture at full resolution (before downscaling to 16x16).

    Use this if you want the raw SD output for manual editing.

    Args:
        prompt: Texture description
        negative_prompt: What to avoid
        seed: Random seed
        width: Generation width (multiple of 8)
        height: Generation height (multiple of 8)
        num_inference_steps: Override default step count
        guidance_scale: Override default guidance
    """
    try:
        steps = num_inference_steps or (4 if "turbo" in MODEL_ID.lower() else 25)
        guidance = guidance_scale if guidance_scale is not None else (0.0 if "turbo" in MODEL_ID.lower() else 7.5)

        pipeline = get_pipe()
        width = (min(max(width, 64), 1024) // 8) * 8
        height = (min(max(height, 64), 1024) // 8) * 8

        generator = None
        if seed is not None:
            generator = torch.Generator(device="cpu").manual_seed(seed)

        full_prompt = f"{prompt}, {MINECRAFT_STYLE_SUFFIX}"
        full_negative = f"{negative_prompt}, {MINECRAFT_NEGATIVE}" if negative_prompt else MINECRAFT_NEGATIVE

        t0 = time.time()
        image = pipeline(
            prompt=full_prompt,
            negative_prompt=full_negative,
            num_inference_steps=steps,
            guidance_scale=guidance,
            width=width,
            height=height,
            generator=generator,
        ).images[0]
        elapsed = time.time() - t0

        return json.dumps({
            "image_base64": image_to_base64(image),
            "size": f"{width}x{height}",
            "format": "png",
            "prompt": prompt,
            "full_prompt": full_prompt,
            "seed": seed,
            "steps": steps,
            "guidance_scale": guidance,
            "device": DEVICE,
            "generation_time_seconds": round(elapsed, 1),
        }, indent=2)
    except Exception as e:
        return json.dumps({"error": str(e)})


@mcp.tool()
def list_texture_types() -> str:
    """List all Minecraft texture categories and example prompts."""
    try:
        result = {
            "model": MODEL_ID,
            "device": DEVICE,
            "output_size": f"{TARGET_SIZE}x{TARGET_SIZE}",
            "source_size": f"{GEN_SIZE}x{GEN_SIZE}",
            "max_colors_default": 16,
            "categories": {},
        }
        for cat, info in TEXTURE_CATEGORIES.items():
            result["categories"][cat] = {
                "description": info["description"],
                "examples": info["prompts"][:5],
                "total_prompts": len(info["prompts"]),
            }

        return json.dumps(result, indent=2)
    except Exception as e:
        return json.dumps({"error": str(e)})


# ─── Gradio UI ────────────────────────────────────────────────

import gradio as gr


def ui_generate_texture(
    prompt: str,
    style: str,
    negative_prompt: str,
    seed: int,
    max_colors: int,
    tileable: bool,
    transparent: bool,
    transparency_mode: str,
    magenta_tolerance: int,
    flood_tolerance: int,
):
    """Gradio wrapper for generate_texture."""
    if not prompt.strip():
        return None, None, None, "Please enter a prompt."
    try:
        result = generate_one(
            prompt=prompt.strip(),
            negative_prompt=negative_prompt or "",
            seed=seed if seed >= 0 else None,
            steps=4 if "turbo" in MODEL_ID.lower() else 25,
            guidance=0.0 if "turbo" in MODEL_ID.lower() else 7.5,
            downscale=True,
            max_colors=max_colors if max_colors > 0 else None,
            tileable=tileable,
            transparent=transparent,
            transparency_mode=transparency_mode,
            magenta_tolerance=magenta_tolerance,
            flood_tolerance=flood_tolerance,
        )

        # Decode base64 images
        pixel_img = None
        source_img = None
        info_parts = []

        if "image_base64" in result:
            pixel_img = Image.open(io.BytesIO(base64.b64decode(result["image_base64"])))
            info_parts.append(f"Texture: {result['size']}")

        if "source_256_base64" in result:
            source_img = Image.open(io.BytesIO(base64.b64decode(result["source_256_base64"])))
            info_parts.append(f"Source: {result['source_size']}")

        info_parts.append(f"Device: {result['device']}")
        info_parts.append(f"Time: {result['generation_time_seconds']}s")
        info_parts.append(f"Seed: {result['seed']}")

        if result.get("has_transparency"):
            info_parts.append(f"Transparent pixels: {result.get('transparent_pixels', '?')}/{result.get('total_pixels', '?')}")
            info_parts.append(f"Mode: {result.get('transparency_mode')}")

        info_text = " | ".join(info_parts)
        return pixel_img, source_img, pixel_img, info_text

    except Exception as e:
        return None, None, None, f"Error: {e}"


def ui_generate_texture_set(theme: str, count: int, transparency_mode: str, magenta_tolerance: int):
    """Gradio wrapper for generate_texture_set."""
    try:
        category = TEXTURE_CATEGORIES.get(theme.lower())
        if category:
            prompts = category["prompts"][:count]
            auto_transparent = category.get("transparent", False)
        else:
            prompts = [f"{theme} variant {i+1}" for i in range(count)]
            auto_transparent = False

        images = []
        names = []
        steps = 4 if "turbo" in MODEL_ID.lower() else 25
        guidance = 0.0 if "turbo" in MODEL_ID.lower() else 7.5

        for i, prompt in enumerate(prompts):
            result = generate_one(
                prompt=prompt,
                negative_prompt="",
                seed=None,
                steps=steps,
                guidance=guidance,
                downscale=True,
                max_colors=16,
                tileable=True,
                transparent=auto_transparent,
                transparency_mode=transparency_mode,
                magenta_tolerance=magenta_tolerance,
            )
            if "image_base64" in result:
                img = Image.open(io.BytesIO(base64.b64decode(result["image_base64"])))
                # Scale up for display (16x16 is too small to see)
                img_display = img.resize((128, 128), Image.NEAREST)
                images.append(img_display)
                names.append(prompt)

        return images, f"Generated {len(images)} textures for '{theme}'"

    except Exception as e:
        return [], f"Error: {e}"


def build_ui():
    """Build the Gradio interface."""
    category_choices = list(TEXTURE_CATEGORIES.keys())

    with gr.Blocks(
        title="Minecraft Texture Generator",
    ) as demo:
        gr.Markdown(
            "# 🎮 Minecraft Texture Generator\n"
            "Generate 16×16 pixel art Minecraft-style textures using Stable Diffusion. "
            "The model generates at 256×256 with pixel art prompting, then downscales "
            "to 16×16 with nearest-neighbor + color quantization.\n\n"
            f"**Model:** `{MODEL_ID}` | **Device:** `{DEVICE}` | ⚠️ CPU inference is slow (~60-120s per texture)"
        )

        with gr.Tabs():
            # ── Tab 1: Single Texture ──
            with gr.Tab("Single Texture"):
                with gr.Row():
                    with gr.Column(scale=2):
                        prompt_input = gr.Textbox(
                            label="Prompt",
                            placeholder="e.g. diamond ore, oak planks, glass pane",
                            lines=1,
                        )
                        style_input = gr.Dropdown(
                            choices=["block", "item", "entity", "particle", "gui"],
                            value="block",
                            label="Style",
                        )
                        negative_input = gr.Textbox(
                            label="Negative Prompt",
                            placeholder="e.g. round, organic",
                            lines=1,
                        )
                        with gr.Row():
                            seed_input = gr.Number(label="Seed (-1 = random)", value=-1, precision=0)
                            max_colors_input = gr.Slider(0, 64, value=16, step=1, label="Max Colors (0=unlimited)")
                        tileable_input = gr.Checkbox(label="Tileable", value=True)

                        with gr.Accordion("Transparency Settings", open=False):
                            transparent_input = gr.Checkbox(label="Enable Transparency", value=False)
                            trans_mode_input = gr.Dropdown(
                                choices=["magenta_key", "corner_flood", "manual_mask"],
                                value="magenta_key",
                                label="Transparency Mode",
                            )
                            magenta_tol_input = gr.Slider(10, 128, value=50, step=5, label="Magenta Tolerance")
                            flood_tol_input = gr.Slider(5, 100, value=30, step=5, label="Flood Tolerance")

                        generate_btn = gr.Button("🎨 Generate Texture", variant="primary")
                        info_output = gr.Textbox(label="Info", interactive=False)

                    with gr.Column(scale=3):
                        with gr.Row():
                            pixel_output = gr.Image(label="16×16 Texture (scaled 8x)", type="pil", height=256)
                            source_output = gr.Image(label="256×256 Source", type="pil", height=256)
                        download_output = gr.Image(label="Download (original 16×16)", type="pil", height=128)

                # Quick prompts
                gr.Markdown("### Quick Prompts")
                quick_prompts = [
                    "diamond ore", "oak planks", "cobblestone", "glass", "netherrack",
                    "grass block top", "spruce leaves", "diamond sword", "tnt side",
                    "gold block", "obsidian", "glowstone", "redstone ore",
                ]
                for row_start in range(0, len(quick_prompts), 7):
                    row_prompts = quick_prompts[row_start:row_start + 7]
                    with gr.Row():
                        for qp in row_prompts:
                            gr.Button(qp, size="sm").click(
                                lambda p=qp: p, inputs=[], outputs=[prompt_input]
                            )

                generate_btn.click(
                    fn=ui_generate_texture,
                    inputs=[
                        prompt_input, style_input, negative_input, seed_input,
                        max_colors_input, tileable_input, transparent_input,
                        trans_mode_input, magenta_tol_input, flood_tol_input,
                    ],
                    outputs=[pixel_output, source_output, download_output, info_output],
                )

            # ── Tab 2: Batch Generation ──
            with gr.Tab("Batch / Theme"):
                with gr.Row():
                    with gr.Column():
                        theme_input = gr.Dropdown(
                            choices=category_choices,
                            value="ores",
                            label="Texture Category",
                        )
                        count_input = gr.Slider(1, 10, value=3, step=1, label="Count")
                        batch_trans_mode = gr.Dropdown(
                            choices=["magenta_key", "corner_flood"],
                            value="magenta_key",
                            label="Transparency Mode",
                        )
                        batch_magenta_tol = gr.Slider(10, 128, value=50, step=5, label="Magenta Tolerance")
                        batch_btn = gr.Button("🎨 Generate Set", variant="primary")
                        batch_info = gr.Textbox(label="Status", interactive=False)

                    with gr.Column():
                        gallery_output = gr.Gallery(
                            label="Generated Textures (8x scale)",
                            columns=5,
                            height=400,
                            object_fit="scale-down",
                        )

                batch_btn.click(
                    fn=ui_generate_texture_set,
                    inputs=[theme_input, count_input, batch_trans_mode, batch_magenta_tol],
                    outputs=[gallery_output, batch_info],
                )

            # ── Tab 3: Categories ──
            with gr.Tab("Categories"):
                cat_md = "## Texture Categories\n\n"
                for cat, info in TEXTURE_CATEGORIES.items():
                    trans_badge = " 🔲 *transparent*" if info.get("transparent") else ""
                    cat_md += f"### `{cat}`{trans_badge}\n{info['description']}\n\n"
                    cat_md += "Prompts: " + ", ".join(f"`{p}`" for p in info["prompts"][:8])
                    if len(info["prompts"]) > 8:
                        cat_md += f" ... (+{len(info['prompts']) - 8} more)"
                    cat_md += "\n\n"
                gr.Markdown(cat_md)

    return demo


# ─── Combined ASGI App (MCP + Gradio + Health) ────────────────

from fastapi import FastAPI
from fastapi.responses import JSONResponse


def create_combined_app():
    """Create a combined ASGI app with MCP + Gradio UI + /health."""

    # Build Gradio UI
    demo = build_ui()

    # Create FastAPI app and mount Gradio properly (handles all init)
    api = FastAPI()
    api = gr.mount_gradio_app(api, demo, path="/")

    # Build MCP ASGI app
    mcp_app = mcp.streamable_http_app()

    class CombinedApp:
        """Routes: /health → JSON, /mcp → MCP (with lifespan), rest → Gradio."""

        def __init__(self, mcp_app, api_app):
            self.mcp_app = mcp_app
            self.api_app = api_app

        async def __call__(self, scope, receive, send):
            # MCP needs lifespan events to initialize its task group
            if scope["type"] == "lifespan":
                await self.mcp_app(scope, receive, send)
                return

            path = scope.get("path", "")

            # Health endpoint
            if scope["type"] == "http" and path == "/health":
                body = json.dumps({
                    "status": "ok",
                    "service": "Minecraft Texture Generator MCP",
                    "model": MODEL_ID,
                    "device": DEVICE,
                    "model_loaded": pipe is not None,
                    "output_size": f"{TARGET_SIZE}x{TARGET_SIZE}",
                }).encode()
                await send({
                    "type": "http.response.start",
                    "status": 200,
                    "headers": [
                        [b"content-type", b"application/json"],
                        [b"content-length", str(len(body)).encode()],
                    ],
                })
                await send({"type": "http.response.body", "body": body})
                return

            # MCP endpoint
            if path.startswith("/mcp"):
                await self.mcp_app(scope, receive, send)
                return

            # Everything else → Gradio/FastAPI
            await self.api_app(scope, receive, send)

    return CombinedApp(mcp_app, api)


if __name__ == "__main__":
    import uvicorn

    print("[init] Pre-loading model...")
    get_pipe()

    app = create_combined_app()
    uvicorn.run(app, host="0.0.0.0", port=port)