"""
Mongle Character LoRA — Photo-to-Pixel-Art Pipeline
Standalone script: works after snapshot_download from HuggingFace.

Usage:
  from huggingface_hub import snapshot_download
  repo_dir = snapshot_download("Hadimeeee/mongle-character-lora")
  import sys; sys.path.insert(0, repo_dir)
  from pipeline import run_pipeline
  from PIL import Image

  result = run_pipeline(Image.open("photo.jpg"))
  result["result_nobg"].save("character.png")
"""

from __future__ import annotations

import os
import gc
import json
import re
from pathlib import Path
from typing import Optional

import cv2
import numpy as np
import torch
from PIL import Image

REPO_ID = "Hadimeeee/mongle-character-lora"
LORA_DIR = Path(__file__).parent   # same folder as this script

# ──────────────────────────────────────────────
# Image utilities
# ──────────────────────────────────────────────

def make_square(img: Image.Image, size: int = 1024) -> Image.Image:
    img = img.convert("RGB")
    w, h = img.size
    side = max(w, h)
    sq = Image.new("RGB", (side, side), (255, 255, 255))
    sq.paste(img, ((side - w) // 2, (side - h) // 2))
    return sq.resize((size, size), Image.LANCZOS)


def remove_bg(img: Image.Image) -> Image.Image:
    from rembg import remove as rembg_remove
    rgba = rembg_remove(img.convert("RGBA"))
    white = Image.new("RGB", rgba.size, (255, 255, 255))
    white.paste(rgba, mask=rgba.split()[3])
    return white


def remove_bg_rgba(img: Image.Image) -> Image.Image:
    from rembg import remove as rembg_remove
    return rembg_remove(img.convert("RGBA"))


# ──────────────────────────────────────────────
# SAM → flat color → Canny
# ──────────────────────────────────────────────

def run_sam(img: Image.Image, sam_model: str = "facebook/sam-vit-base"):
    from transformers import SamModel, SamProcessor
    device = "cuda" if torch.cuda.is_available() else "cpu"
    processor = SamProcessor.from_pretrained(sam_model)
    model = SamModel.from_pretrained(sam_model).to(device)
    model.eval()

    w, h = img.size
    cx, cy = w // 2, h // 2
    inputs = processor(img, input_points=[[[cx, cy]]], return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = model(**inputs)
    masks = processor.post_process_masks(
        outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(),
        inputs["reshaped_input_sizes"].cpu(),
    )[0]
    scores = outputs.iou_scores[0, 0].cpu().numpy()
    mask = masks[0, int(np.argmax(scores))].numpy().astype(np.uint8) * 255
    del model, processor; gc.collect(); torch.cuda.empty_cache()
    return Image.fromarray(mask)


def dominant_color(img: Image.Image, mask: Image.Image):
    arr = np.array(img.convert("RGB"))
    m   = np.array(mask) > 128
    px  = arr[m]
    if len(px) == 0:
        return (200, 200, 200)
    from sklearn.cluster import KMeans
    k = KMeans(n_clusters=3, n_init=5, random_state=0).fit(px)
    sizes = np.bincount(k.labels_)
    return tuple(int(c) for c in k.cluster_centers_[np.argmax(sizes)])


def build_flat_color(img: Image.Image, mask: Image.Image) -> Image.Image:
    color = dominant_color(img, mask)
    flat  = Image.new("RGB", img.size, (255, 255, 255))
    mask_arr = np.array(mask) > 128
    flat_arr = np.array(flat)
    flat_arr[mask_arr] = color
    return Image.fromarray(flat_arr)


def extract_canny(flat: Image.Image, lo: int = 50, hi: int = 150) -> Image.Image:
    gray  = cv2.cvtColor(np.array(flat), cv2.COLOR_RGB2GRAY)
    edges = cv2.Canny(gray, lo, hi)
    return Image.fromarray(np.stack([edges] * 3, axis=-1))


# ──────────────────────────────────────────────
# VLM (Qwen2-VL) — appearance extraction
# ──────────────────────────────────────────────

_vlm_model = None
_vlm_proc  = None


def load_vlm(model_name: str = "Qwen/Qwen2-VL-7B-Instruct"):
    global _vlm_model, _vlm_proc
    from transformers import Qwen2VLForConditionalGeneration, AutoProcessor, BitsAndBytesConfig
    bnb = BitsAndBytesConfig(load_in_8bit=True)
    _vlm_model = Qwen2VLForConditionalGeneration.from_pretrained(
        model_name, quantization_config=bnb, device_map="auto"
    )
    _vlm_proc = AutoProcessor.from_pretrained(model_name)


def unload_vlm():
    global _vlm_model, _vlm_proc
    del _vlm_model, _vlm_proc
    _vlm_model = _vlm_proc = None
    gc.collect(); torch.cuda.empty_cache()


def run_vlm(img: Image.Image) -> dict:
    system = (
        "You are a visual analysis assistant. "
        "Analyze the stuffed animal in the image and return ONLY a JSON object "
        "with these fields: animal_type, body_color, secondary_colors (list), "
        "body_shape, eye_style, accessories (list), distinctive_features (list), "
        "controlnet_scale (float 0.45-0.85). "
        "controlnet_scale: 0.45 if no face, 0.5 if pillow-shaped, "
        "0.75 for normal, 0.85 for limbless/round. "
        "No explanation, no markdown, only JSON."
    )
    messages = [{"role": "user", "content": [
        {"type": "image", "image": img},
        {"type": "text",  "text": "Analyze this stuffed animal and return JSON."},
    ]}]
    from qwen_vl_utils import process_vision_info
    text = _vlm_proc.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    image_inputs, _ = process_vision_info(messages)
    inputs = _vlm_proc(text=[text], images=image_inputs, return_tensors="pt")
    inputs = {k: v.to(_vlm_model.device) for k, v in inputs.items()}
    with torch.no_grad():
        out = _vlm_model.generate(**inputs, max_new_tokens=512, temperature=0.1)
    raw = _vlm_proc.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
    m = re.search(r"\{.*\}", raw, re.DOTALL)
    return json.loads(m.group()) if m else {}


def vlm_json_to_prompt(data: dict, extra_en: str = "") -> tuple[str, float]:
    animal    = data.get("animal_type", "plush toy")
    body_col  = data.get("body_color", "colorful")
    sec_cols  = ", ".join(data.get("secondary_colors", []))
    shape     = data.get("body_shape", "round")
    eyes      = data.get("eye_style", "round eyes")
    acc       = ", ".join(data.get("accessories", []))
    feat      = ", ".join(data.get("distinctive_features", []))
    cn_scale  = float(data.get("controlnet_scale", 0.75))

    parts = [
        f"monglestyle, {body_col} {animal} plush",
        shape, eyes,
    ]
    if sec_cols: parts.append(sec_cols)
    if acc:      parts.append(acc)
    if feat:     parts.append(feat)
    if extra_en: parts.append(extra_en)
    parts += [
        "single stuffed animal toy mascot character, full body, centered",
        "front view, cute chibi proportions, 32-bit pixel art sprite",
        "soft pixel shading, clean silhouette, soft brown outline",
        "pure white background",
    ]
    return ", ".join(p for p in parts if p), cn_scale


# ──────────────────────────────────────────────
# ControlNet generation
# ──────────────────────────────────────────────

_pipe = None


def load_pipeline(lcm: bool = True):
    global _pipe
    from diffusers import StableDiffusionXLControlNetPipeline, ControlNetModel
    from diffusers.schedulers import LCMScheduler

    cn = ControlNetModel.from_pretrained(
        "diffusers/controlnet-canny-sdxl-1.0", torch_dtype=torch.float16
    )
    _pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
        "stabilityai/stable-diffusion-xl-base-1.0",
        controlnet=cn, torch_dtype=torch.float16,
    ).to("cuda")

    if lcm:
        _pipe.load_lora_weights("latent-consistency/lcm-lora-sdxl", adapter_name="lcm")
        _pipe.load_lora_weights(str(LORA_DIR), adapter_name="style")
        _pipe.set_adapters(["lcm", "style"], adapter_weights=[1.0, 0.9])
        _pipe.scheduler = LCMScheduler.from_config(_pipe.scheduler.config)
    else:
        _pipe.load_lora_weights(str(LORA_DIR), adapter_name="style")
        _pipe.set_adapters(["style"], adapter_weights=[0.9])


def unload_pipeline():
    global _pipe
    del _pipe; _pipe = None
    gc.collect(); torch.cuda.empty_cache()


def generate_character(
    canny_img: Image.Image,
    prompt: str,
    cn_scale: float = 0.75,
    steps: int = 8,
    guidance: float = 1.5,
    seed: int = 42,
    lora_scale: float = 0.9,
) -> Image.Image:
    neg = "blurry, watermark, text, low quality, deformed, realistic photo, 3d render"
    gen = torch.Generator("cuda").manual_seed(seed)
    out = _pipe(
        prompt=prompt,
        negative_prompt=neg,
        image=canny_img,
        num_inference_steps=steps,
        guidance_scale=guidance,
        controlnet_conditioning_scale=cn_scale,
        cross_attention_kwargs={"scale": lora_scale},
        generator=gen,
    )
    return out.images[0]


# ──────────────────────────────────────────────
# Main API
# ──────────────────────────────────────────────

def run_pipeline(
    image_pil: Image.Image,
    char_desc_en: str = None,
    lcm: bool = True,
    lora_scale: float = 0.9,
    cn_scale_override: float = None,
    steps: int = 8,
    seed: int = 42,
    out_dir: str = None,
    sam_model: str = "facebook/sam-vit-base",
    vlm_model: str = "Qwen/Qwen2-VL-7B-Instruct",
) -> dict:
    """
    Full photo-to-pixel-art pipeline.

    Args:
        image_pil       : Input PIL image (stuffed animal photo)
        char_desc_en    : Optional English description to supplement VLM output
        lcm             : Use LCM LoRA for fast 8-step generation
        lora_scale      : Character LoRA weight (default 0.9)
        cn_scale_override: Override ControlNet scale (None = VLM recommendation)
        steps           : Inference steps (8 with LCM, 25-30 without)
        seed            : Random seed
        out_dir         : Save intermediate outputs here (optional)
        sam_model       : SAM model ID
        vlm_model       : Qwen2-VL model ID

    Returns dict with keys:
        result, result_nobg, canny, flat_color, appearance, prompt, cn_scale
    """
    if out_dir:
        Path(out_dir).mkdir(parents=True, exist_ok=True)

    # STEP 1: Preprocess
    print("[1/5] Preprocessing...")
    sq   = make_square(image_pil)
    nobg = remove_bg(sq)

    # STEP 2: SAM → flat color → Canny
    print("[2/5] SAM + Canny edge extraction...")
    mask  = run_sam(nobg, sam_model)
    flat  = build_flat_color(nobg, mask)
    canny = extract_canny(flat)

    # STEP 3: VLM appearance analysis
    print("[3/5] VLM appearance analysis...")
    load_vlm(vlm_model)
    appearance = run_vlm(nobg)
    unload_vlm()

    prompt, cn_scale = vlm_json_to_prompt(appearance, char_desc_en or "")
    if cn_scale_override is not None:
        cn_scale = cn_scale_override

    # STEP 4: Generate character
    print("[4/5] Generating pixel art character...")
    guidance = 1.5 if lcm else 7.5
    load_pipeline(lcm=lcm)
    result = generate_character(
        canny, prompt, cn_scale=cn_scale,
        steps=steps, guidance=guidance, seed=seed, lora_scale=lora_scale,
    )
    unload_pipeline()

    # STEP 5: Remove background from result
    print("[5/5] Final background removal...")
    result_nobg_rgba = remove_bg_rgba(result)

    # Save outputs
    if out_dir:
        d = Path(out_dir)
        nobg.save(d / "nobg.png")
        flat.save(d / "flat_color.png")
        canny.save(d / "canny.png")
        result.save(d / "result.png")
        result_nobg_rgba.save(d / "result_nobg.png")
        (d / "appearance.json").write_text(json.dumps(appearance, ensure_ascii=False, indent=2))
        (d / "prompt.txt").write_text(prompt)
        print(f"Saved to: {out_dir}")

    return {
        "result":      result,
        "result_nobg": result_nobg_rgba,
        "canny":       canny,
        "flat_color":  flat,
        "appearance":  appearance,
        "prompt":      prompt,
        "cn_scale":    cn_scale,
    }


# ──────────────────────────────────────────────
# CLI
# ──────────────────────────────────────────────

if __name__ == "__main__":
    import argparse
    p = argparse.ArgumentParser(description="Mongle character pipeline")
    p.add_argument("--image",    required=True, help="Input photo path")
    p.add_argument("--out-dir",  default="output", help="Output directory")
    p.add_argument("--desc",     default=None,  help="English character description (optional)")
    p.add_argument("--no-lcm",   dest="lcm", action="store_false", default=True)
    p.add_argument("--cn-scale", type=float, default=None)
    p.add_argument("--steps",    type=int,   default=8)
    p.add_argument("--seed",     type=int,   default=42)
    args = p.parse_args()

    result = run_pipeline(
        image_pil       = Image.open(args.image),
        char_desc_en    = args.desc,
        lcm             = args.lcm,
        cn_scale_override = args.cn_scale,
        steps           = args.steps,
        seed            = args.seed,
        out_dir         = args.out_dir,
    )
    print(f"\nPrompt: {result['prompt']}")
    print(f"cn_scale: {result['cn_scale']}")
    print(f"Done → {args.out_dir}/result_nobg.png")