from __future__ import annotations

import importlib.util
import logging
import os
import random
import re
import subprocess
import sys
import threading
import time
from dataclasses import dataclass
from pathlib import Path

import numpy as np
from PIL import Image, PngImagePlugin

from lora_cache import LoraCache
from weights import materialize

LOGGER = logging.getLogger("sdxl_runtime")
LOGGER.setLevel(logging.INFO)

ROOT = Path(__file__).resolve().parent
COMFY_SHA = "7782171a272bca4bfb8586a5bb801c5c04b5ee6a"
GGUF_SHA = "6ea2651e7df66d7585f6ffee804b20e92fb38b8a"
RUNTIME_ROOT = Path(os.getenv("SDXL_RUNTIME_ROOT", "/tmp/sdxl-comfy-runtime"))
MODEL_ROOT = Path(os.getenv("SDXL_MODEL_ROOT", "/tmp/sdxl-models"))
USER_LORA_ROOT = Path(os.getenv("SDXL_LORA_CACHE", "/tmp/sdxl-user-loras"))
MAIN_SAMPLER = {
    "base_positive": "",
    "neg_embedding": "BadPonyHD.safetensors",
    "embedding_emphasis": 0.95,
    "width": 896,
    "height": 504,
    "batch_size": 1,
    "seed_type": "Fixed",
    "steps": 27,
    "cfg": 8,
    "sampler_name": "dpmpp_2m",
    "scheduler": "sgm_uniform",
    "denoise": 1.0,
    "ckpt_name": "None",
    "lora_name": "None",
    "lora_mode": "LoRA Model + Clip",
    "strength_model": 0,
    "strength_clip": 0,
    "clip_skip": 2,
}

for variable, directory in {
    "PIP_CACHE_DIR": "/tmp/sdxl-cache/pip",
    "MPLCONFIGDIR": "/tmp/sdxl-cache/matplotlib",
    "YOLO_CONFIG_DIR": "/tmp/sdxl-cache/ultralytics",
}.items():
    os.environ.setdefault(variable, directory)
    Path(os.environ[variable]).mkdir(parents=True, exist_ok=True)


def _checkout(url: str, revision: str, target: Path) -> None:
    if target.exists():
        current = subprocess.run(
            ["git", "-C", str(target), "rev-parse", "HEAD"], capture_output=True, text=True, check=True
        ).stdout.strip()
        if current != revision:
            raise RuntimeError(f"Unexpected source revision in {target}: {current}")
        return
    target.parent.mkdir(parents=True, exist_ok=True)
    subprocess.run(["git", "clone", "--filter=blob:none", "--no-checkout", url, str(target)], check=True)
    subprocess.run(["git", "-C", str(target), "fetch", "--depth=1", "origin", revision], check=True)
    subprocess.run(["git", "-C", str(target), "checkout", "--detach", revision], check=True)


def prepare_sources() -> tuple[Path, Path]:
    local = os.getenv("COMFYUI_ROOT")
    if local:
        comfy_root = Path(local).resolve()
        if subprocess.run(["git", "-C", str(comfy_root), "rev-parse", "HEAD"], capture_output=True, text=True, check=True).stdout.strip() != COMFY_SHA:
            raise RuntimeError("COMFYUI_ROOT is not at the pinned revision")
        gguf_root = comfy_root / "custom_nodes" / "ComfyUI-GGUF"
        return comfy_root, gguf_root
    comfy_root = RUNTIME_ROOT / "ComfyUI"
    gguf_root = RUNTIME_ROOT / "ComfyUI-GGUF"
    _checkout("https://github.com/comfyanonymous/ComfyUI.git", COMFY_SHA, comfy_root)
    _checkout("https://github.com/city96/ComfyUI-GGUF.git", GGUF_SHA, gguf_root)
    return comfy_root, gguf_root


def _import_gguf(path: Path):
    spec = importlib.util.spec_from_file_location(
        "sdxl_gguf", path / "__init__.py", submodule_search_locations=[str(path)]
    )
    if spec is None or spec.loader is None:
        raise RuntimeError("Cannot import the pinned ComfyUI-GGUF nodes")
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    nodes_module = importlib.import_module("sdxl_gguf.nodes")

    # ZeroGPU's module-level CUDA emulation already associates a Python object
    # with tensors created from the GGUF mmap. The legacy Tensor.__new__ path
    # then refuses to wrap them. _make_subclass is the supported zero-copy path.
    import torch

    def _ggml_new(cls, data, *args, **kwargs):
        return torch.Tensor._make_subclass(cls, data, bool(data.requires_grad))

    ops_module = importlib.import_module("sdxl_gguf.ops")
    ops_module.GGMLTensor.__new__ = staticmethod(_ggml_new)
    return nodes_module


@dataclass
class Components:
    model: object
    clip: object
    vae: object
    user_loader: object
    hyper_loader: object
    dmd_loader: object
    sampler: object
    face_detailer: object
    foot_prepare: object
    foot_apply: object


@dataclass(frozen=True)
class GenerationSeeds:
    main: int
    face: int
    foot: int
    hand: int
    eye: int


class Runtime:
    def __init__(self) -> None:
        started = time.perf_counter()
        self.lock = threading.RLock()
        self.loras = LoraCache(USER_LORA_ROOT)
        self.components = self._bootstrap()
        LOGGER.info("SDXL bootstrap completed in %.2fs", time.perf_counter() - started)

    def _bootstrap(self) -> Components:
        comfy_root, gguf_root = prepare_sources()
        sys.path.insert(0, str(comfy_root))
        sys.path.insert(0, str(ROOT / "vendor" / "kyanite"))
        sys.path.insert(0, str(ROOT / "vendor"))

        # ComfyUI has both a top-level ``utils`` package and ``comfy/utils.py``.
        # Import the package first so headless imports cannot resolve the latter by mistake.
        import utils.install_util  # noqa: F401

        paths = materialize(MODEL_ROOT)
        import folder_paths

        for kind, directory in {
            "unet": MODEL_ROOT / "unet",
            "diffusion_models": MODEL_ROOT / "unet",
            "text_encoders": MODEL_ROOT / "text_encoders",
            "clip": MODEL_ROOT / "text_encoders",
            "vae": MODEL_ROOT / "vae",
            "embeddings": MODEL_ROOT / "embeddings",
            "loras": MODEL_ROOT / "loras",
            "ultralytics": MODEL_ROOT / "ultralytics",
            "ultralytics_bbox": MODEL_ROOT / "ultralytics/bbox",
            "ultralytics_segm": MODEL_ROOT / "ultralytics/segm",
        }.items():
            folder_paths.add_model_folder_path(kind, str(directory), is_default=True)
        folder_paths.add_model_folder_path("loras", str(USER_LORA_ROOT), is_default=True)

        import nodes
        from vendor.kyanite import KyaniteApplyDetailBatch, KyaniteFaceDetailer, KyanitePrepareDetailBatch
        from vendor.kyanite.kyanite_face_detailer import _load_ultralytics_detectors

        gguf = _import_gguf(gguf_root)
        model = gguf.UnetLoaderGGUF().load_unet(paths["unet"].name)[0]
        clip = nodes.DualCLIPLoader().load_clip(paths["clip_l"].name, paths["clip_g"].name, "sdxl")[0]
        vae = nodes.VAELoader().load_vae(paths["vae"].name)[0]
        user_loader = nodes.LoraLoader()
        hyper_loader = nodes.LoraLoaderModelOnly()
        dmd_loader = nodes.LoraLoaderModelOnly()

        # Populate all fixed caches during global bootstrap. Patches are cheap clones per request.
        hyper_loader.load_lora_model_only(model, paths["hyper"].name, 1.0)
        dmd_loader.load_lora_model_only(model, paths["dmd2"].name, 1.0)
        for detector in (
            "bbox/face_yolov8m.pt",
            "segm/foot-yolov8l.pt",
            "segm/PitHandDetailer-v2-Test-v9c.pt",
            "segm/PitEyeDetailer-v2-seg.pt",
        ):
            _load_ultralytics_detectors(detector)

        if os.getenv("SPACE_ID"):
            import comfy.model_management
            import torch
            # ComfyUI's cache housekeeping calls the low-level CUDA IPC API,
            # which ZeroGPU intentionally cannot emulate at module scope.
            torch.cuda.ipc_collect = lambda: None
            comfy.model_management.load_models_gpu([model, clip.patcher])

        return Components(model, clip, vae, user_loader, hyper_loader, dmd_loader, nodes.KSampler(), KyaniteFaceDetailer(), KyanitePrepareDetailBatch(), KyaniteApplyDetailBatch())

    @staticmethod
    def prompts(positive: str, negative: str, safe: bool) -> tuple[str, str]:
        positive, negative = (positive or "").strip(), (negative or "").strip()
        if safe:
            positive = ", ".join(part for part in ("rating_safe, safe, sfw", positive) if part)
            negative = ", ".join(part for part in (negative, "rating_explicit, rating_questionable, nsfw, nude, nudity, sex, explicit") if part)
        return positive, negative

    @staticmethod
    def eye_prompt(positive: str) -> str:
        eye_tags = [
            part.strip()
            for part in (positive or "").split(",")
            if re.search(r"\beyes?\b", part, flags=re.IGNORECASE)
        ]
        return ", ".join(part for part in eye_tags if part) or "detailed eyes"

    @staticmethod
    def cleanup() -> None:
        try:
            import comfy.model_management
            import torch
            comfy.model_management.soft_empty_cache()
            if torch.cuda.is_available():
                torch.cuda.synchronize()
        except Exception:
            pass

    @staticmethod
    def generation_seeds(seed: int | None) -> GenerationSeeds:
        main = random.randrange(2**64) if seed is None or seed < 0 else int(seed)
        return GenerationSeeds(
            main=main,
            face=random.randrange(2**64),
            foot=random.randrange(2**64),
            hand=random.randrange(2**64),
            eye=random.randrange(2**64),
        )

    @staticmethod
    def _detail_kwargs(seed: int, positive: str, *, detector: str, hand: bool = False, face: bool = False) -> dict:
        guide_size = 384 if not face and not hand else 512
        return dict(
            guide_size=guide_size, guide_size_for=True, max_size=768, seed=seed, steps=6, cfg=1.8,
            sampler_name="dpmpp_2m" if face else "lcm", scheduler="sgm_uniform" if face else "simple",
            positive_text=positive if face else ("detailed hand, 5 fingers" if hand else "Ultra detailed feet, 5 toes"),
            negative_text=("extra fingers, missing fingers, fused fingers, malformed fingers, broken fingers, deformed hand, mutated hand, bad anatomy, 6 fingers" if hand or face else "extra fingers, extra toes, fused fingers, poorly detailed fingers, polish, toe polish"),
            neg_embedding="BadPonyHD.safetensors", embedding_emphasis=1.0,
            denoise=0.5, feather=5, noise_mask=True, force_inpaint=True, bbox_threshold=0.5,
            bbox_dilation=10, bbox_crop_factor=3, sam_detection_hint="center-1", sam_dilation=0,
            sam_threshold=0.93, sam_bbox_expansion=0, sam_mask_hint_threshold=0.7,
            sam_mask_hint_use_negative="False", drop_size=10, bbox_model_name=detector, wildcard="",
            cycle=1, inpaint_model=False, noise_mask_feather=20, skip_keywords="flowjob, footjob, fj" if hand else "",
            tiled_encode=False, tiled_decode=False, scene_names="", scene_table="",
        )

    def generate(
        self,
        positive: str,
        negative: str,
        safe: bool,
        face_detail: bool,
        eyes_detail: bool,
        hands_detail: bool,
        feet_detail: bool,
        uploads: list,
        model_strengths: list[float],
        clip_strengths: list[float],
        seed: int | None,
    ) -> tuple[Image.Image, int]:
        with self.lock:
            started = time.perf_counter()
            c = self.components
            seeds = self.generation_seeds(seed)
            positive, negative = self.prompts(positive, negative, safe)
            model, clip = c.model, c.clip
            for upload, sm, sc in zip(uploads, model_strengths, clip_strengths):
                if not upload:
                    continue
                item = self.loras.add(upload)
                model, clip = c.user_loader.load_lora(model, clip, item.name, float(sm), float(sc))
            # Use the authorized sampler implementation to preserve prompt weighting and latent behavior.
            from vendor.kyanite.kyanite_ksampler import KyaniteKSampler
            image, model, clip, vae, positive, negative, seed = KyaniteKSampler().sample(
                seed=seeds.main, positive_text=positive, negative_text=negative,
                model=model, clip=clip, vae=c.vae, **MAIN_SAMPLER,
            )
            LOGGER.info("Main sampling completed in %.2fs", time.perf_counter() - started)
            hyper_model = None
            if face_detail or eyes_detail:
                hyper_model = c.hyper_loader.load_lora_model_only(
                    model, "Hyper-SDXL-8steps-CFG-lora.safetensors", 1.0
                )[0]
            if face_detail:
                face_started = time.perf_counter()
                face_kw = self._detail_kwargs(
                    seeds.face, positive, detector="bbox/face_yolov8m.pt", face=True
                )
                image, _, clip, vae, *_ = c.face_detailer.doit(
                    image=image, model=hyper_model, clip=clip, vae=vae,
                    denoise_visible=0.5, detail_scale=1.5,
                    expression_seductive=True, **face_kw,
                )
                LOGGER.info("Face detail completed in %.2fs", time.perf_counter() - face_started)

            foot_batch = None
            detail_model = None
            if feet_detail or hands_detail:
                detail_model = c.dmd_loader.load_lora_model_only(
                    model, "dmd2_sdxl_4step_lora.safetensors", 1.0
                )[0]
            hand_model, hand_clip, hand_vae = detail_model, clip, vae
            if feet_detail:
                foot_started = time.perf_counter()
                foot_kw = self._detail_kwargs(
                    seeds.foot, positive, detector="segm/foot-yolov8l.pt"
                )
                _, foot_batch, _, _, hand_model, hand_clip, hand_vae, _ = c.foot_prepare.prepare(
                    image=image, model=detail_model, clip=clip, vae=vae, **foot_kw
                )
                LOGGER.info(
                    "Foot detail preparation completed in %.2fs",
                    time.perf_counter() - foot_started,
                )
            if hands_detail:
                hand_started = time.perf_counter()
                hand_kw = self._detail_kwargs(
                    seeds.hand,
                    positive,
                    detector="segm/PitHandDetailer-v2-Test-v9c.pt",
                    hand=True,
                )
                image, *_ = c.face_detailer.doit(
                    image=image, model=hand_model, clip=hand_clip, vae=hand_vae,
                    denoise_visible=0.55, detail_scale=1.0, **hand_kw
                )
                LOGGER.info("Hand detail completed in %.2fs", time.perf_counter() - hand_started)
            if feet_detail:
                foot_apply_started = time.perf_counter()
                image, *_ = c.foot_apply.apply(image, foot_batch)
                LOGGER.info(
                    "Foot detail apply completed in %.2fs",
                    time.perf_counter() - foot_apply_started,
                )

            if eyes_detail:
                eye_started = time.perf_counter()
                eye_kw = self._detail_kwargs(
                    seeds.eye,
                    self.eye_prompt(positive),
                    detector="segm/PitEyeDetailer-v2-seg.pt",
                    face=True,
                )
                image, *_ = c.face_detailer.doit(
                    image=image, model=hyper_model, clip=clip, vae=vae,
                    denoise_visible=0.5, detail_scale=1.5, **eye_kw,
                )
                LOGGER.info("Eye detail completed in %.2fs", time.perf_counter() - eye_started)
            array = np.clip(image[0].detach().cpu().numpy() * 255.0, 0, 255).astype(np.uint8)
            output = Image.fromarray(array, "RGB")
            metadata = PngImagePlugin.PngInfo()
            metadata.add_text("seed", str(seed))
            output.info["pnginfo"] = metadata
            LOGGER.info("Generation pipeline completed in %.2fs", time.perf_counter() - started)
            return output, seed


RUNTIME = None if os.getenv("SDXL_SKIP_BOOTSTRAP") == "1" else Runtime()