"""Core generation engine for Texture-Craft."""

import base64
import io
import os
import pathlib
import random
import time
from typing import Optional, Tuple, List, Dict, Any

import numpy as np
import torch
from PIL import Image
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler

# Magenta background used during model training to represent transparent alpha.
BG = np.array([255, 0, 255], dtype=np.int16)
DEFAULT_MODEL_DIR = pathlib.Path("models/minecraft-item-16px")
REPO_ID = "yuuki14202028/minecraft-item-16px"


class TextureEngine:
    _instance: Optional["TextureEngine"] = None

    def __init__(
        self,
        model_path: Optional[str] = None,
        cache_downloaded_model: bool = True,
        device: Optional[str] = None,
    ):
        if device not in (None, "cuda", "cpu"):
            raise ValueError("device must be 'cuda', 'cpu', or None.")
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.dtype = torch.float16 if self.device == "cuda" else torch.float32
        self.model_path = pathlib.Path(model_path) if model_path else DEFAULT_MODEL_DIR
        self.cache_downloaded_model = cache_downloaded_model
        self.pipe: Optional[StableDiffusionPipeline] = None
        self._load_pipeline()

    @classmethod
    def get_instance(
        cls,
        model_path: Optional[str] = None,
        cache_downloaded_model: bool = True,
        device: Optional[str] = None,
    ) -> "TextureEngine":
        """Return the singleton engine instance, creating it once if needed."""
        if cls._instance is None:
            cls._instance = cls(model_path, cache_downloaded_model, device)
        return cls._instance

    def _resolve_model_source(self) -> tuple[str, Optional[str]]:
        if self.model_path.exists() and any(self.model_path.iterdir()):
            return str(self.model_path), None
        return REPO_ID, os.environ.get("HF_TOKEN")

    def _load_pipeline(self) -> None:
        """Load the model from a local cache or download it from Hugging Face."""
        source, token = self._resolve_model_source()

        print(f"[TextureEngine] Loading pipeline on {self.device} from '{source}'...", flush=True)
        pipe = StableDiffusionPipeline.from_pretrained(
            source,
            token=token,
            torch_dtype=self.dtype,
            safety_checker=None,
            requires_safety_checker=False,
        )
        pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
        pipe = pipe.to(self.device)
        pipe.set_progress_bar_config(disable=True)
        self.pipe = pipe

        if source == REPO_ID and self.cache_downloaded_model:
            self.model_path.mkdir(parents=True, exist_ok=True)
            self.pipe.save_pretrained(str(self.model_path))

        print(f"[TextureEngine] Pipeline ready on {self.device}!", flush=True)

    @staticmethod
    def fold_to_tile(image: Image.Image, size: int = 16) -> Image.Image:
        """Downsample with hard pixel edges instead of blending neighboring colors."""
        mode = "RGBA" if "A" in image.getbands() else "RGB"
        return image.convert(mode).resize((size, size), Image.Resampling.NEAREST)

    @staticmethod
    def key_out(tile: Image.Image, tol: int = 100) -> Image.Image:
        """Remove magenta and uniform edge backgrounds without cutting into sprites."""
        rgb = np.asarray(tile.convert("RGB"), dtype=np.int16)
        height, width = rgb.shape[:2]
        magenta = np.abs(rgb - BG).sum(axis=2) <= tol

        corner_size = max(1, min(16, height // 4, width // 4))
        corners = np.stack(
            (
                rgb[:corner_size, :corner_size].reshape(-1, 3),
                rgb[:corner_size, -corner_size:].reshape(-1, 3),
                rgb[-corner_size:, :corner_size].reshape(-1, 3),
                rgb[-corner_size:, -corner_size:].reshape(-1, 3),
            )
        )
        corner_colors = np.median(corners, axis=1)
        corner_distances = np.abs(
            corner_colors[:, None, :] - corner_colors[None, :, :]
        ).max(axis=2)
        matching_corners = corner_distances <= 24
        background_index = int(np.argmax(matching_corners.sum(axis=1)))
        background_color = corner_colors[background_index]
        uniform_background = matching_corners[background_index].sum() >= 2

        transparent = magenta
        if uniform_background:
            background = np.max(np.abs(rgb - background_color), axis=2) <= 20
            connected = np.zeros((height, width), dtype=bool)
            pending = [
                (0, 0),
                (0, width - 1),
                (height - 1, 0),
                (height - 1, width - 1),
            ]
            while pending:
                row, col = pending.pop()
                if connected[row, col] or not background[row, col]:
                    continue
                connected[row, col] = True
                if row > 0:
                    pending.append((row - 1, col))
                if row + 1 < height:
                    pending.append((row + 1, col))
                if col > 0:
                    pending.append((row, col - 1))
                if col + 1 < width:
                    pending.append((row, col + 1))
            if connected.mean() >= 0.15:
                transparent |= connected

        out = np.dstack(
            [rgb.astype(np.uint8), np.where(transparent, 0, 255).astype(np.uint8)]
        )
        return Image.fromarray(out)

    @staticmethod
    def image_to_base64(image: Image.Image, format: str = "PNG") -> str:
        """Encodes PIL image to standard base64 string."""
        buffered = io.BytesIO()
        image.save(buffered, format=format)
        return base64.b64encode(buffered.getvalue()).decode("utf-8")

    @staticmethod
    def image_to_data_uri(image: Image.Image, format: str = "PNG") -> str:
        """Encodes PIL image to browser-renderable data URI."""
        b64 = TextureEngine.image_to_base64(image, format=format)
        return f"data:image/{format.lower()};base64,{b64}"

    @staticmethod
    def contact_sheet(
        images: List[Image.Image], cols: int = 8, cell: int = 64, checker: bool = True
    ) -> Image.Image:
        """Assembles a list of 16x16 tiles onto a checkered preview sheet."""
        if not images:
            return Image.new("RGBA", (cell, cell), (0, 0, 0, 0))
        rows = (len(images) + cols - 1) // cols
        sheet = Image.new("RGBA", (cols * cell, rows * cell), (32, 32, 36, 255))
        if checker:
            c = cell // 8
            for y in range(0, sheet.height, c):
                for x in range(0, sheet.width, c):
                    if (x // c + y // c) % 2:
                        sheet.paste((54, 54, 60, 255), (x, y, x + c, y + c))
        for i, im in enumerate(images):
            tile = im.convert("RGBA").resize((cell, cell), Image.Resampling.NEAREST)
            sheet.alpha_composite(tile, ((i % cols) * cell, (i // cols) * cell))
        return sheet

    @staticmethod
    def _normalize_seed(seed: Optional[int], *, default_max: int = 2**31 - 1) -> int:
        if seed is None or seed < 0:
            return random.randint(0, default_max)
        return int(seed)

    def generate_single(
        self,
        prompt: str,
        negative_prompt: str = "",
        seed: Optional[int] = None,
        steps: int = 30,
        guidance_scale: float = 9.0,
        res: int = 256,
        tile_size: int = 16,
    ) -> Tuple[Image.Image, Image.Image, int, float]:
        """Generate a single 16x16 Minecraft item texture.

        Returns: (tile, raw_image, seed, latency_seconds)
        """
        if self.pipe is None:
            raise RuntimeError("Texture pipeline is not initialized.")

        seed = self._normalize_seed(seed)
        t0 = time.time()
        generator = torch.Generator(device="cpu").manual_seed(seed)

        raw = self.pipe(
            prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=int(steps),
            guidance_scale=float(guidance_scale),
            height=int(res),
            width=int(res),
            generator=generator,
        ).images[0]

        tile = self.fold_to_tile(self.key_out(raw), size=tile_size)
        latency = round(time.time() - t0, 3)
        return tile, raw, seed, latency

    def generate_batch(
        self,
        prompt: str,
        variants: int = 1,
        negative_prompt: str = "",
        seed_base: Optional[int] = None,
        steps: int = 30,
        guidance_scale: float = 9.0,
    ) -> List[Dict[str, Any]]:
        """Generate multiple variants for a single prompt."""
        variant_count = max(1, int(variants))
        base_seed = self._normalize_seed(seed_base, default_max=2**30 - 1)

        results: List[Dict[str, Any]] = []
        for index in range(variant_count):
            seed = base_seed + (index * 1000)
            tile, raw, actual_seed, latency = self.generate_single(
                prompt=prompt,
                negative_prompt=negative_prompt,
                seed=seed,
                steps=steps,
                guidance_scale=guidance_scale,
            )
            results.append(
                {
                    "variant_index": index + 1,
                    "seed": actual_seed,
                    "latency_sec": latency,
                    "tile_image": tile,
                    "raw_image": raw,
                    "base64_tile": self.image_to_base64(tile),
                    "data_uri_tile": self.image_to_data_uri(tile),
                }
            )
        return results

    def get_system_status(self) -> Dict[str, Any]:
        """Returns hardware and engine status."""
        cuda_ok = torch.cuda.is_available()
        gpu_name = torch.cuda.get_device_name(0) if cuda_ok else None
        vram_total_gb = None
        vram_allocated_gb = None

        if cuda_ok:
            props = torch.cuda.get_device_properties(0)
            vram_total_gb = round(props.total_memory / (1024**3), 2)
            vram_allocated_gb = round(torch.cuda.memory_allocated(0) / (1024**3), 2)

        return {
            "status": "ready" if self.pipe is not None else "initializing",
            "device": self.device,
            "cuda_available": cuda_ok,
            "gpu_name": gpu_name,
            "vram_total_gb": vram_total_gb,
            "vram_allocated_gb": vram_allocated_gb,
            "model_path": str(self.model_path),
        }