#!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "websocket-client>=1.6",
#     "requests>=2.31",
#     "pillow>=10.0",
# ]
# ///
"""
age_shift_cli.py — detect current age and generate target-age portraits
via ComfyUI + ComfyUI-AgeShift + ComfyUI-PuLID-Flux.

Usage:
    # Detect only
    uv run python age_shift_cli.py face.jpg --detect-only --ip 10.8.0.12

    # Single target age (e.g. age this person up to 65)
    uv run python age_shift_cli.py face.jpg --target-age 65 --ip 10.8.0.12

    # Age sequence (5/15/25/45/65)
    uv run python age_shift_cli.py face.jpg \\
        --target-ages 5,15,25,45,65 \\
        --gender male --out ./out --ip 10.8.0.12

    # Dump workflow JSONs for debugging
    uv run python age_shift_cli.py face.jpg --target-age 65 \\
        --dump-workflow /tmp/wf.json --ip 10.8.0.12
"""
import argparse
import json
import sys
import time
import uuid
from pathlib import Path
from typing import Any, Dict, List

import requests
from websocket import create_connection


# =============================================================================
# ComfyClient
# =============================================================================
class ComfyClient:
    def __init__(self, server: str = "127.0.0.1:8188"):
        self.server = server
        self.client_id = str(uuid.uuid4())
        self.http = f"http://{server}"
        self.ws = f"ws://{server}/ws?clientId={self.client_id}"

    def upload_image(self, file_path: Path, overwrite: bool = True) -> str:
        with open(file_path, "rb") as f:
            files = {"image": (file_path.name, f, "application/octet-stream")}
            data = {"overwrite": "true" if overwrite else "false", "type": "input"}
            r = requests.post(f"{self.http}/upload/image", files=files, data=data, timeout=60)
        r.raise_for_status()
        return r.json().get("name", file_path.name)

    def submit(self, workflow: dict) -> str:
        payload = {"prompt": workflow, "client_id": self.client_id}
        r = requests.post(f"{self.http}/prompt", json=payload, timeout=30)
        if r.status_code != 200:
            raise RuntimeError(f"/prompt {r.status_code}: {r.text[:1000]}")
        return r.json()["prompt_id"]

    def history(self, prompt_id: str) -> dict:
        r = requests.get(f"{self.http}/history/{prompt_id}", timeout=30)
        r.raise_for_status()
        data = r.json()
        return data.get(prompt_id, {})

    def view(self, filename: str, subfolder: str = "", type_: str = "output") -> bytes:
        params = {"filename": filename, "subfolder": subfolder, "type": type_}
        r = requests.get(f"{self.http}/view", params=params, timeout=120)
        r.raise_for_status()
        return r.content

    def run(self, workflow: dict, timeout_s: float = 600) -> dict:
        ws = create_connection(self.ws, timeout=timeout_s)
        try:
            pid = self.submit(workflow)
            print(f"      prompt_id={pid}")
            t0 = time.time()
            last_node = None
            while True:
                if time.time() - t0 > timeout_s:
                    raise TimeoutError(f"Workflow timed out after {timeout_s}s")
                msg = ws.recv()
                if isinstance(msg, bytes):
                    continue
                evt = json.loads(msg)
                t = evt.get("type")
                d = evt.get("data", {})
                if d.get("prompt_id") and d["prompt_id"] != pid:
                    continue
                if t == "progress":
                    cur, tot = d.get("value", 0), d.get("max", 1)
                    print(f"\r      ▸ {cur:>4}/{tot}                          ", end="", flush=True)
                elif t == "executing":
                    n = d.get("node")
                    if n is None:
                        print()  # finalize line
                        break
                    if n != last_node:
                        print(f"\r      ▸ node {str(n):<32}", end="", flush=True)
                        last_node = n
                elif t == "execution_error":
                    raise RuntimeError(
                        f"Error in node {d.get('node_id')}: {d.get('exception_message')}"
                    )
            return self.history(pid)
        finally:
            ws.close()


# =============================================================================
# Workflow builders (API format — node ids are strings)
# =============================================================================
def build_detect_workflow(
    image_name: str,
    *,
    model_id: str = "nateraw/vit-age-classifier",
    device: str = "auto",
    dtype: str = "fp32",
) -> Dict[str, Any]:
    return {
        "load": {
            "class_type": "LoadImage",
            "inputs": {"image": image_name, "upload": "image"},
        },
        "loader": {
            "class_type": "AgeDetectorLoader",
            "inputs": {"model_id": model_id, "device": device, "dtype": dtype},
        },
        "detect": {
            "class_type": "AgeDetect",
            "inputs": {
                "detector": ["loader", 0],
                "image": ["load", 0],
                "method": "weighted_mean",
            },
        },
        "show_age": {
            "class_type": "PreviewAny",
            "inputs": {"source": ["detect", 0]},
        },
        "show_band": {
            "class_type": "PreviewAny",
            "inputs": {"source": ["detect", 1]},
        },
        "show_conf": {
            "class_type": "PreviewAny",
            "inputs": {"source": ["detect", 2]},
        },
        "show_summary": {
            "class_type": "PreviewAny",
            "inputs": {"source": ["detect", 3]},
        },
    }


def build_shift_workflow(
    image_name: str,
    target_age: int,
    *,
    # flux model file names (must be present in ComfyUI/models/*)
    flux_unet: str,
    flux_clip_l: str,
    flux_t5: str,
    flux_vae: str,
    pulid_weights: str,
    insightface_provider: str = "CUDA",
    # generation knobs
    width: int = 1024,
    height: int = 1024,
    steps: int = 20,
    cfg: float = 1.0,
    guidance: float = 3.5,
    sampler: str = "euler",
    scheduler: str = "simple",
    seed: int = 0,
    # pulid knobs
    pulid_weight: float = 0.9,
    pulid_start: float = 0.0,
    pulid_end: float = 1.0,
    # prompt knobs
    subject: str = "person",
    gender: str = "unspecified",
    style: str = "photorealistic portrait, sharp focus, natural studio lighting",
    extra_positive: str = "same identity, same facial structure, same eye shape",
    negative_base: str = (
        "cartoon, illustration, painting, anime, 3d render, "
        "deformed, distorted, blurry, low quality, plastic skin, watermark, text"
    ),
) -> Dict[str, Any]:
    """Build a full PuLID-Flux + AgeShift workflow (API format).

    Requires ComfyUI-PuLID-Flux installed alongside this package.
    """
    return {
        # Inputs
        "load_face": {
            "class_type": "LoadImage",
            "inputs": {"image": image_name, "upload": "image"},
        },
        # Age-conditioned prompt
        "prompt": {
            "class_type": "AgePromptBuilder",
            "inputs": {
                "target_age": int(target_age),
                "subject": subject,
                "gender": gender,
                "style": style,
                "extra_positive": extra_positive,
                "negative_base": negative_base,
                "include_negative_age_drift": True,
            },
        },
        # Flux backbone
        "unet": {
            "class_type": "UNETLoader",
            "inputs": {"unet_name": flux_unet, "weight_dtype": "default"},
        },
        "dual_clip": {
            "class_type": "DualCLIPLoader",
            "inputs": {
                "clip_name1": flux_clip_l,
                "clip_name2": flux_t5,
                "type": "flux",
            },
        },
        "vae": {"class_type": "VAELoader", "inputs": {"vae_name": flux_vae}},
        # PuLID-Flux identity injection
        "pulid_model": {
            "class_type": "PulidFluxModelLoader",
            "inputs": {"pulid_file": pulid_weights},
        },
        "eva_clip": {"class_type": "PulidFluxEvaClipLoader", "inputs": {}},
        "insightface": {
            "class_type": "PulidFluxInsightFaceLoader",
            "inputs": {"provider": insightface_provider},
        },
        "apply_pulid": {
            "class_type": "ApplyPulidFlux",
            "inputs": {
                "model": ["unet", 0],
                "pulid_flux": ["pulid_model", 0],
                "eva_clip": ["eva_clip", 0],
                "face_analysis": ["insightface", 0],
                "image": ["load_face", 0],
                "weight": pulid_weight,
                "start_at": pulid_start,
                "end_at": pulid_end,
                "fusion": "mean",
                "fusion_weight_max": 1.0,
                "fusion_weight_min": 0.0,
                "train_step": 1000,
                "use_gray": True,
            },
        },
        # Text → conditioning
        "pos_encode": {
            "class_type": "CLIPTextEncode",
            "inputs": {"clip": ["dual_clip", 0], "text": ["prompt", 0]},
        },
        "neg_encode": {
            "class_type": "CLIPTextEncode",
            "inputs": {"clip": ["dual_clip", 0], "text": ["prompt", 1]},
        },
        "guidance_node": {
            "class_type": "FluxGuidance",
            "inputs": {"conditioning": ["pos_encode", 0], "guidance": guidance},
        },
        # Latent + sample
        "latent": {
            "class_type": "EmptySD3LatentImage",
            "inputs": {"width": width, "height": height, "batch_size": 1},
        },
        "sampler_node": {
            "class_type": "KSampler",
            "inputs": {
                "model": ["apply_pulid", 0],
                "positive": ["guidance_node", 0],
                "negative": ["neg_encode", 0],
                "latent_image": ["latent", 0],
                "seed": int(seed),
                "steps": int(steps),
                "cfg": float(cfg),
                "sampler_name": sampler,
                "scheduler": scheduler,
                "denoise": 1.0,
            },
        },
        "decode": {
            "class_type": "VAEDecode",
            "inputs": {"samples": ["sampler_node", 0], "vae": ["vae", 0]},
        },
        "save": {
            "class_type": "SaveImage",
            "inputs": {
                "images": ["decode", 0],
                "filename_prefix": f"ageshift_age{int(target_age):03d}",
            },
        },
        "show_prompt_summary": {
            "class_type": "PreviewAny",
            "inputs": {"source": ["prompt", 2]},
        },
    }


# =============================================================================
# History helpers
# =============================================================================
def extract_text(history: dict, node_id: str) -> str:
    outs = history.get("outputs", {})
    if node_id not in outs:
        return ""
    no = outs[node_id]
    if "text" in no:
        v = no["text"]
        return "\n".join(str(x) for x in v) if isinstance(v, list) else str(v)
    parts = []
    for v in no.values():
        if isinstance(v, list):
            parts.extend(str(x) for x in v)
    return "\n".join(parts)


def extract_images(history: dict, node_id: str) -> List[Dict[str, str]]:
    outs = history.get("outputs", {})
    if node_id not in outs:
        return []
    return outs[node_id].get("images", [])


# =============================================================================
# Main
# =============================================================================
def main():
    ap = argparse.ArgumentParser(
        description="AgeShift: detect current age and generate target-age portraits "
                    "via ComfyUI + PuLID-Flux",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    ap.add_argument("image", type=Path, help="Source face image (jpg/png)")

    # server
    ap.add_argument("--ip", default="127.0.0.1")
    ap.add_argument("--port", type=int, default=8188)
    ap.add_argument("--server", default=None, help="full host:port (overrides --ip --port)")

    # mode
    ap.add_argument("--detect-only", action="store_true", help="only run age detection")
    ap.add_argument("--target-age", type=int, default=None, help="single target age to generate")
    ap.add_argument(
        "--target-ages",
        default=None,
        help="csv list of target ages, e.g. 5,15,25,45,65",
    )

    # generation knobs
    ap.add_argument("--subject", default="person")
    ap.add_argument(
        "--gender",
        default="unspecified",
        choices=["unspecified", "male", "female", "non-binary"],
    )
    ap.add_argument(
        "--style",
        default="photorealistic portrait, sharp focus, natural studio lighting, neutral background",
    )
    ap.add_argument(
        "--extra-positive",
        default="same identity, same facial structure, same eye shape",
    )
    ap.add_argument("--width", type=int, default=1024)
    ap.add_argument("--height", type=int, default=1024)
    ap.add_argument("--steps", type=int, default=20)
    ap.add_argument("--cfg", type=float, default=1.0)
    ap.add_argument("--guidance", type=float, default=3.5)
    ap.add_argument("--sampler", default="euler")
    ap.add_argument("--scheduler", default="simple")
    ap.add_argument("--seed", type=int, default=0, help="0 = random per age")
    ap.add_argument("--pulid-weight", type=float, default=0.9)
    ap.add_argument("--pulid-start", type=float, default=0.0)
    ap.add_argument("--pulid-end", type=float, default=1.0)

    # model files (must already exist on the ComfyUI server)
    ap.add_argument("--flux-unet", default="flux1-dev.safetensors")
    ap.add_argument("--flux-clip-l", default="clip_l.safetensors")
    ap.add_argument("--flux-t5", default="t5xxl_fp8_e4m3fn.safetensors")
    ap.add_argument("--flux-vae", default="ae.safetensors")
    ap.add_argument("--pulid-weights", default="pulid_flux_v0.9.1.safetensors")
    ap.add_argument(
        "--insightface-provider",
        default="CUDA",
        choices=["CUDA", "CPU", "ROCM"],
    )

    # detector
    ap.add_argument("--detector-model", default="nateraw/vit-age-classifier")
    ap.add_argument(
        "--detector-device",
        default="auto",
        choices=["auto", "cuda", "cuda:0", "cuda:1", "cpu"],
    )
    ap.add_argument("--detector-dtype", default="fp32", choices=["fp32", "fp16", "bf16"])

    # output / debug
    ap.add_argument("--out", type=Path, default=Path("./out"))
    ap.add_argument(
        "--dump-workflow",
        type=Path,
        default=None,
        help="dump generated workflow JSON(s) for debugging",
    )
    ap.add_argument("--timeout", type=float, default=600.0)

    args = ap.parse_args()
    server = args.server or f"{args.ip}:{args.port}"
    args.out.mkdir(parents=True, exist_ok=True)

    if not args.image.exists():
        print(f"ERROR: image not found: {args.image}", file=sys.stderr)
        sys.exit(2)

    client = ComfyClient(server)

    # ---- upload ----
    print(f"[1/3] Uploading {args.image.name} → {server} ...")
    server_name = client.upload_image(args.image)
    print(f"      uploaded as: {server_name}")

    # ---- detect ----
    print(f"[2/3] Detecting age ...")
    detect_wf = build_detect_workflow(
        server_name,
        model_id=args.detector_model,
        device=args.detector_device,
        dtype=args.detector_dtype,
    )
    if args.dump_workflow:
        dump = args.dump_workflow.with_suffix(".detect.json")
        dump.parent.mkdir(parents=True, exist_ok=True)
        dump.write_text(json.dumps(detect_wf, indent=2, ensure_ascii=False))
        print(f"      wrote {dump}")
    hist = client.run(detect_wf, timeout_s=args.timeout)
    summary = extract_text(hist, "show_summary") or "(no summary)"
    age_str = extract_text(hist, "show_age").strip()
    try:
        detected_age = int(float(age_str))
    except (ValueError, TypeError):
        detected_age = None
    print(f"      → {summary}")

    if args.detect_only:
        print("Done (detect only).")
        return

    # ---- shift ----
    targets: List[int] = []
    if args.target_age is not None:
        targets.append(args.target_age)
    if args.target_ages:
        for s in args.target_ages.split(","):
            s = s.strip()
            if s:
                try:
                    targets.append(int(s))
                except ValueError:
                    print(f"      ! skip invalid age: {s!r}", file=sys.stderr)
    if not targets:
        print(
            "No --target-age / --target-ages given; nothing to generate.",
            file=sys.stderr,
        )
        sys.exit(0)
    # dedup, keep order
    seen = set()
    targets = [a for a in targets if not (a in seen or seen.add(a))]

    print(f"[3/3] Generating {len(targets)} age-shifted image(s) ...")
    for i, age in enumerate(targets, 1):
        print(f"  ({i}/{len(targets)}) target age = {age}")
        seed = args.seed if args.seed else int(time.time() * 1000) % (2**31 - 1) + i
        wf = build_shift_workflow(
            server_name,
            age,
            flux_unet=args.flux_unet,
            flux_clip_l=args.flux_clip_l,
            flux_t5=args.flux_t5,
            flux_vae=args.flux_vae,
            pulid_weights=args.pulid_weights,
            insightface_provider=args.insightface_provider,
            width=args.width,
            height=args.height,
            steps=args.steps,
            cfg=args.cfg,
            guidance=args.guidance,
            sampler=args.sampler,
            scheduler=args.scheduler,
            seed=seed,
            pulid_weight=args.pulid_weight,
            pulid_start=args.pulid_start,
            pulid_end=args.pulid_end,
            subject=args.subject,
            gender=args.gender,
            style=args.style,
            extra_positive=args.extra_positive,
        )
        if args.dump_workflow:
            dump = args.dump_workflow.with_suffix(f".shift_age{age}.json")
            dump.write_text(json.dumps(wf, indent=2, ensure_ascii=False))
            print(f"      wrote {dump}")
        try:
            hist = client.run(wf, timeout_s=args.timeout)
        except Exception as e:
            print(f"      ! failed: {e}", file=sys.stderr)
            continue
        prompt_summary = extract_text(hist, "show_prompt_summary").strip()
        if prompt_summary:
            print(f"      prompt: {prompt_summary}")
        imgs = extract_images(hist, "save")
        if not imgs:
            print(f"      ! no output images for age {age}", file=sys.stderr)
            continue
        for img in imgs:
            data = client.view(
                img["filename"],
                img.get("subfolder", ""),
                img.get("type", "output"),
            )
            out_path = args.out / f"age_{age:03d}_{img['filename']}"
            out_path.write_bytes(data)
            print(f"      → {out_path}  ({len(data)//1024} KB)")

    print("All done.")


if __name__ == "__main__":
    main()