#!/usr/bin/env python3
"""Local ComfyUI + Grounded-SAM + vtracer character asset pipeline.

This script is designed for a creator workstation, not for the WordPress
server. It can optionally queue a local ComfyUI workflow, then uses local
Grounding DINO + Segment Anything checkpoints to cut semantic character parts
such as "head", "left arm", and "torso". If vtracer is installed, each part is
converted to SVG and combined into one grouped SVG file.

No image is uploaded to haotianblog.com. The only network call this script makes
is to the local ComfyUI server when --comfyui-url is provided.

Examples:
  python comfyui_grounded_sam_vtracer_pipeline.py \
    --source exploded-character.png \
    --grounding-config GroundingDINO_SwinT_OGC.py \
    --grounding-weights groundingdino_swint_ogc.pth \
    --sam-checkpoint sam_vit_h_4b8939.pth \
    --part-prompts part_prompts.example.json \
    --out out/character-assets

  python comfyui_grounded_sam_vtracer_pipeline.py \
    --comfyui-url http://127.0.0.1:8188 \
    --workflow-template comfyui_exploded_character_workflow.template.json \
    --reference front.png --reference side.png --pose-image exploded_view_openpose_skeleton.png \
    --positive "same wolf character, clean cel shading, separated parts, flat background" \
    --grounding-config GroundingDINO_SwinT_OGC.py \
    --grounding-weights groundingdino_swint_ogc.pth \
    --sam-checkpoint sam_vit_h_4b8939.pth \
    --out out/wolf-rig-assets

Optional dependencies:
  pip install pillow numpy requests vtracer torch torchvision
  pip install git+https://github.com/IDEA-Research/GroundingDINO.git
  pip install git+https://github.com/facebookresearch/segment-anything.git
"""

from __future__ import annotations

import argparse
import json
import re
import sys
import time
import urllib.parse
import urllib.request
import uuid
import xml.etree.ElementTree as ET
from dataclasses import dataclass
from pathlib import Path
from typing import Any


DEFAULT_PARTS = [
    {"id": "head", "prompt": "head, face, hair"},
    {"id": "eyes", "prompt": "eyes"},
    {"id": "torso", "prompt": "torso, chest, clothes"},
    {"id": "left-arm", "prompt": "left arm"},
    {"id": "right-arm", "prompt": "right arm"},
    {"id": "left-hand", "prompt": "left hand"},
    {"id": "right-hand", "prompt": "right hand"},
    {"id": "left-leg", "prompt": "left leg"},
    {"id": "right-leg", "prompt": "right leg"},
    {"id": "tail", "prompt": "tail"},
    {"id": "accessory", "prompt": "headband, weapon, accessory"},
]


@dataclass
class PartSpec:
    part_id: str
    prompt: str


@dataclass
class LayerRecord:
    part_id: str
    prompt: str
    score: float
    box_xyxy: list[float]
    png_path: Path
    svg_path: Path | None


def slugify(value: str) -> str:
    slug = re.sub(r"[^a-zA-Z0-9_-]+", "-", value.strip().lower())
    return re.sub(r"-{2,}", "-", slug).strip("-") or "part"


def read_part_specs(path: Path | None) -> list[PartSpec]:
    raw: Any = DEFAULT_PARTS
    if path:
        raw = json.loads(path.read_text(encoding="utf-8"))
    specs: list[PartSpec] = []
    for item in raw:
        if isinstance(item, str):
            specs.append(PartSpec(slugify(item), item))
        elif isinstance(item, dict):
            prompt = str(item.get("prompt") or item.get("text") or item.get("name") or "").strip()
            part_id = slugify(str(item.get("id") or prompt))
            if prompt:
                specs.append(PartSpec(part_id, prompt))
    if not specs:
        raise ValueError("No part prompts were provided.")
    return specs


def replace_placeholders(value: Any, mapping: dict[str, str]) -> Any:
    if isinstance(value, str):
        for key, replacement in mapping.items():
            value = value.replace(key, replacement)
        return value
    if isinstance(value, list):
        return [replace_placeholders(item, mapping) for item in value]
    if isinstance(value, dict):
        return {key: replace_placeholders(item, mapping) for key, item in value.items()}
    return value


def clean_comfyui_prompt(value: Any) -> Any:
    if isinstance(value, dict):
        return {key: clean_comfyui_prompt(item) for key, item in value.items() if not str(key).startswith("_")}
    if isinstance(value, list):
        return [clean_comfyui_prompt(item) for item in value]
    if isinstance(value, str) and re.fullmatch(r"-?\d+", value):
        return int(value)
    if isinstance(value, str) and re.fullmatch(r"-?\d+\.\d+", value):
        return float(value)
    return value


def post_json(url: str, payload: dict[str, Any]) -> dict[str, Any]:
    data = json.dumps(payload).encode("utf-8")
    request = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"}, method="POST")
    with urllib.request.urlopen(request, timeout=60) as response:
        return json.loads(response.read().decode("utf-8"))


def get_json(url: str) -> dict[str, Any]:
    with urllib.request.urlopen(url, timeout=60) as response:
        return json.loads(response.read().decode("utf-8"))


def upload_comfyui_image(base_url: str, path: Path) -> str:
    boundary = "----haotianblog-" + uuid.uuid4().hex
    body = bytearray()
    body.extend(f"--{boundary}\r\n".encode())
    body.extend(
        (
            'Content-Disposition: form-data; name="image"; '
            f'filename="{path.name}"\r\nContent-Type: application/octet-stream\r\n\r\n'
        ).encode()
    )
    body.extend(path.read_bytes())
    body.extend(f"\r\n--{boundary}--\r\n".encode())
    request = urllib.request.Request(
        base_url.rstrip("/") + "/upload/image",
        data=bytes(body),
        headers={"Content-Type": f"multipart/form-data; boundary={boundary}"},
        method="POST",
    )
    with urllib.request.urlopen(request, timeout=120) as response:
        payload = json.loads(response.read().decode("utf-8"))
    return str(payload.get("name") or path.name)


def render_with_comfyui(args: argparse.Namespace, out_dir: Path) -> Path:
    if not args.workflow_template:
        raise ValueError("--workflow-template is required when --comfyui-url is used.")

    workflow = json.loads(args.workflow_template.read_text(encoding="utf-8"))
    reference_names = [upload_comfyui_image(args.comfyui_url, path) for path in args.reference]
    pose_name = upload_comfyui_image(args.comfyui_url, args.pose_image) if args.pose_image else ""

    mapping = {
        "__POSITIVE_PROMPT__": args.positive,
        "__NEGATIVE_PROMPT__": args.negative,
        "__OPENPOSE_IMAGE__": pose_name,
        "__SEED__": str(args.seed),
    }
    for index, name in enumerate(reference_names, start=1):
        mapping[f"__REFERENCE_IMAGE_{index}__"] = name
    if reference_names:
        mapping["__REFERENCE_IMAGE__"] = reference_names[0]

    prompt = clean_comfyui_prompt(replace_placeholders(workflow, mapping))
    client_id = uuid.uuid4().hex
    queued = post_json(args.comfyui_url.rstrip("/") + "/prompt", {"prompt": prompt, "client_id": client_id})
    prompt_id = str(queued["prompt_id"])
    deadline = time.monotonic() + args.comfyui_timeout

    history: dict[str, Any] | None = None
    while time.monotonic() < deadline:
        history_payload = get_json(args.comfyui_url.rstrip("/") + "/history/" + urllib.parse.quote(prompt_id))
        if prompt_id in history_payload:
            history = history_payload[prompt_id]
            break
        time.sleep(1.5)
    if history is None:
        raise TimeoutError(f"ComfyUI prompt {prompt_id} did not finish within {args.comfyui_timeout} seconds.")

    outputs = history.get("outputs", {})
    downloaded: list[Path] = []
    for node_output in outputs.values():
        for image in node_output.get("images", []):
            query = urllib.parse.urlencode(
                {
                    "filename": image.get("filename", ""),
                    "subfolder": image.get("subfolder", ""),
                    "type": image.get("type", "output"),
                }
            )
            image_url = args.comfyui_url.rstrip("/") + "/view?" + query
            target = out_dir / "comfyui-generated" / str(image.get("filename", f"{prompt_id}.png"))
            target.parent.mkdir(parents=True, exist_ok=True)
            with urllib.request.urlopen(image_url, timeout=120) as response:
                target.write_bytes(response.read())
            downloaded.append(target)
    if not downloaded:
        raise RuntimeError("ComfyUI finished but did not return an image output.")
    return downloaded[0]


def ensure_grounded_sam_args(args: argparse.Namespace) -> None:
    missing = [
        name
        for name, value in (
            ("--grounding-config", args.grounding_config),
            ("--grounding-weights", args.grounding_weights),
            ("--sam-checkpoint", args.sam_checkpoint),
        )
        if not value
    ]
    if missing:
        raise ValueError("Grounded-SAM segmentation requires " + ", ".join(missing))


def grounded_sam_layers(source: Path, specs: list[PartSpec], out_dir: Path, args: argparse.Namespace) -> list[LayerRecord]:
    ensure_grounded_sam_args(args)
    try:
        import numpy as np
        import torch
        from PIL import Image
        from groundingdino.util.inference import load_image, load_model, predict
        from segment_anything import SamPredictor, sam_model_registry
        from torchvision.ops import box_convert
    except Exception as exc:  # pragma: no cover - optional dependency path
        raise RuntimeError(
            "Grounded-SAM dependencies are not installed. Install GroundingDINO, segment-anything, torch, "
            "torchvision, pillow, and numpy on your local workstation."
        ) from exc

    image_source, image_tensor = load_image(str(source))
    height, width = image_source.shape[:2]
    pil = Image.open(source).convert("RGBA")
    if pil.size != (width, height):
        pil = pil.resize((width, height), Image.Resampling.LANCZOS)

    grounding_model = load_model(str(args.grounding_config), str(args.grounding_weights))
    device = "cuda" if torch.cuda.is_available() and not args.cpu else "cpu"
    sam = sam_model_registry[args.sam_model_type](checkpoint=str(args.sam_checkpoint))
    sam.to(device=device)
    predictor = SamPredictor(sam)
    predictor.set_image(image_source)

    layers_dir = out_dir / "layers"
    layers_dir.mkdir(parents=True, exist_ok=True)
    records: list[LayerRecord] = []

    for spec in specs:
        boxes, logits, _phrases = predict(
            model=grounding_model,
            image=image_tensor,
            caption=spec.prompt,
            box_threshold=args.box_threshold,
            text_threshold=args.text_threshold,
        )
        if len(boxes) == 0:
            print(f"warning: no box found for {spec.part_id}: {spec.prompt}", file=sys.stderr)
            continue
        best_index = int(torch.argmax(logits).item())
        box = boxes[best_index : best_index + 1] * torch.tensor([width, height, width, height])
        xyxy = box_convert(boxes=box, in_fmt="cxcywh", out_fmt="xyxy").to(device)
        transformed = predictor.transform.apply_boxes_torch(xyxy, image_source.shape[:2])
        masks, _scores, _logits = predictor.predict_torch(
            point_coords=None,
            point_labels=None,
            boxes=transformed,
            multimask_output=False,
        )
        mask = masks[0][0].detach().cpu().numpy().astype(bool)
        rgba = np.array(pil)
        layer = np.zeros_like(rgba)
        layer[mask] = rgba[mask]
        png_path = layers_dir / f"{spec.part_id}.png"
        Image.fromarray(layer, "RGBA").save(png_path)
        records.append(
            LayerRecord(
                part_id=spec.part_id,
                prompt=spec.prompt,
                score=float(logits[best_index].detach().cpu().item()),
                box_xyxy=[round(float(v), 2) for v in xyxy[0].detach().cpu().tolist()],
                png_path=png_path,
                svg_path=None,
            )
        )
    return records


def vectorize_layers(records: list[LayerRecord]) -> None:
    try:
        import vtracer
    except Exception:
        print("warning: vtracer is not installed; PNG layers were generated without SVG vectorization.", file=sys.stderr)
        return

    for record in records:
        svg_path = record.png_path.with_suffix(".svg")
        vtracer.convert_image_to_svg_py(
            str(record.png_path),
            str(svg_path),
            colormode="color",
            hierarchical="stacked",
            mode="spline",
            filter_speckle=6,
            color_precision=7,
            layer_difference=12,
            corner_threshold=65,
            length_threshold=3.5,
            max_iterations=10,
            splice_threshold=45,
            path_precision=3,
        )
        record.svg_path = svg_path


def svg_inner(svg_path: Path) -> str:
    root = ET.parse(svg_path).getroot()
    return "\n".join(ET.tostring(child, encoding="unicode") for child in list(root))


def write_combined_svg(source: Path, records: list[LayerRecord], out_dir: Path) -> Path | None:
    try:
        from PIL import Image
    except Exception:
        return None
    svg_records = [record for record in records if record.svg_path]
    if not svg_records:
        return None
    width, height = Image.open(source).size
    lines = [
        f'<svg xmlns="http://www.w3.org/2000/svg" width="{width}" height="{height}" viewBox="0 0 {width} {height}">',
        '  <metadata>Generated locally by haotianblog Animation Asset Lab pipeline.</metadata>',
    ]
    for record in svg_records:
        lines.append(f'  <g id="{record.part_id}" data-prompt="{record.prompt}" data-score="{record.score:.4f}">')
        for inner_line in svg_inner(record.svg_path).splitlines():
            lines.append("    " + inner_line)
        lines.append("  </g>")
    lines.append("</svg>")
    combined = out_dir / "character-rig-layered.svg"
    combined.write_text("\n".join(lines) + "\n", encoding="utf-8")
    return combined


def write_manifest(source: Path, records: list[LayerRecord], combined_svg: Path | None, out_dir: Path) -> None:
    manifest = {
        "source": str(source),
        "privacy": "All segmentation, vectorization, and packaging was local. No website upload was made.",
        "combined_svg": combined_svg.name if combined_svg else None,
        "layers": [
            {
                "id": record.part_id,
                "prompt": record.prompt,
                "score": round(record.score, 4),
                "box_xyxy": record.box_xyxy,
                "png": str(record.png_path.relative_to(out_dir)),
                "svg": str(record.svg_path.relative_to(out_dir)) if record.svg_path else None,
            }
            for record in records
        ],
    }
    (out_dir / "character-rig-manifest.json").write_text(json.dumps(manifest, indent=2, ensure_ascii=False), encoding="utf-8")


def main() -> int:
    parser = argparse.ArgumentParser(description="Generate 2D character rig assets with local ComfyUI, Grounded-SAM, and vtracer.")
    parser.add_argument("--source", type=Path, help="Existing exploded-view character image. If omitted, --comfyui-url must render one.")
    parser.add_argument("--out", type=Path, default=Path("character-rig-assets"))
    parser.add_argument("--part-prompts", type=Path, help="JSON list of semantic parts to cut out.")
    parser.add_argument("--comfyui-url", help="Local ComfyUI URL, for example http://127.0.0.1:8188")
    parser.add_argument("--workflow-template", type=Path, help="ComfyUI API workflow JSON with placeholders.")
    parser.add_argument("--reference", type=Path, action="append", default=[], help="Reference character image for IP-Adapter. Repeatable.")
    parser.add_argument("--pose-image", type=Path, help="Exploded-view OpenPose layout image.")
    parser.add_argument("--positive", default="same character, clean cel shaded exploded view, separated body parts on a flat background")
    parser.add_argument("--negative", default="overlapping limbs, cropped body, realistic photo, noisy background, extra fingers")
    parser.add_argument("--seed", type=int, default=123456)
    parser.add_argument("--comfyui-timeout", type=int, default=600)
    parser.add_argument("--grounding-config", type=Path)
    parser.add_argument("--grounding-weights", type=Path)
    parser.add_argument("--sam-checkpoint", type=Path)
    parser.add_argument("--sam-model-type", default="vit_h", choices=("vit_h", "vit_l", "vit_b"))
    parser.add_argument("--box-threshold", type=float, default=0.28)
    parser.add_argument("--text-threshold", type=float, default=0.25)
    parser.add_argument("--cpu", action="store_true", help="Force SAM to CPU even when CUDA is available.")
    args = parser.parse_args()

    args.out.mkdir(parents=True, exist_ok=True)
    source = args.source
    if args.comfyui_url:
        source = render_with_comfyui(args, args.out)
    if source is None or not source.exists():
        print("Provide --source or a working --comfyui-url + --workflow-template render path.", file=sys.stderr)
        return 2

    specs = read_part_specs(args.part_prompts)
    records = grounded_sam_layers(source, specs, args.out, args)
    vectorize_layers(records)
    combined_svg = write_combined_svg(source, records, args.out)
    write_manifest(source, records, combined_svg, args.out)

    print(f"source: {source}")
    print(f"layers: {len(records)}")
    if combined_svg:
        print(f"combined svg: {combined_svg}")
    print(f"manifest: {args.out / 'character-rig-manifest.json'}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())