#!/usr/bin/env python
"""Generate matching court-card artwork with the local Draw Things server.

The Hearts family is already approved.  The other suits use the corresponding
Heart card as an img2img composition/style reference, while prompts vary the
character and garment motifs.  Output remains raster artwork; card indices,
suit marks and number-card geometry are added later by build_deck.py.

Usage:
    ./.venv/bin/python render_court.py QD
    ./.venv/bin/python render_court.py JD QD KD JC QC KC JS QS KS
    ./.venv/bin/python render_court.py --remaining
"""

import argparse
import datetime
import hashlib
import json
import os
import plistlib
import sys
import time

import numpy as np
from PIL import Image

import dt_client as dt
import imageService_pb2 as pb


HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
CONCEPTS = os.path.join(ROOT, "concepts")
OUTDIR = os.path.join(ROOT, "art", "courts", "raw")

MODEL = "qwen_image_2512_bf16_i8x.ckpt"
DT_APP = "/Applications/Draw Things.app/Contents/Info.plist"
WIDTH, HEIGHT = 640, 896

RANKS = ("J", "Q", "K")
SUITS = ("H", "D", "C", "S")
REMAINING = tuple(rank + suit for suit in "DCS" for rank in RANKS)

HEART_REFERENCE = {
    "J": os.path.join(CONCEPTS, "dt-court-jh-v3.png"),
    "Q": os.path.join(CONCEPTS, "dt-court-qh-v3.png"),
    "K": os.path.join(CONCEPTS, "dt-court-kh-v3.png"),
}

FIGURE = {
    "J": "a young knave with shoulder-length hair holding an upright halberd",
    "Q": "a regal queen wearing an ornate golden crown and holding a stylized flower",
    "K": "a mature bearded king wearing an ornate golden crown and holding an upright sword",
}

SUIT_DIRECTION = {
    "D": ("the diamond court, with angular diamond-shaped garment ornament; "
          "red and gold are dominant with muted-blue accents"),
    "C": ("the club court, with rounded trefoil garment ornament; black and "
          "muted blue are dominant with red and gold accents"),
    "S": ("the spade court, with elegant pointed garment ornament; black and "
          "muted blue are dominant with red and gold accents"),
}

STYLE = (
    "original Anglo-American playing-card court illustration, crisp flat "
    "vector-like drawing, precise black keylines, restrained red, black, "
    "golden-yellow, muted-blue and warm-skin palette on white card stock, "
    "mirrored two-way composition, richly patterned robes, traditional but "
    "entirely original character design, clean white corners"
)

NEGATIVE = (
    "rank letter, corner index, corner marking, suit symbol, heart symbol, "
    "diamond symbol, club symbol, spade symbol, pip, text, letters, numbers, "
    "typography, watermark, signature, logo, branded artwork, photorealistic, "
    "photograph, 3d render, soft gradient shading, blur, noise, drop shadow"
)

# Stable, visibly different seeds per card.
DEFAULT_SEEDS = {
    "JD": 181, "QD": 182, "KD": 183,
    "JC": 281, "QC": 282, "KC": 283,
    "JS": 381, "QS": 382, "KS": 383,
}


def sha256_file(path):
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def drawthings_version():
    try:
        with open(DT_APP, "rb") as f:
            return plistlib.load(f).get("CFBundleShortVersionString", "unknown")
    except Exception:
        return "unknown"


def checkpoint_hashes(stub, files):
    response = stub.FilesExist(
        pb.FileListRequest(files=files, filesWithHash=files), timeout=900
    )
    hashes = list(response.hashes)
    result = {}
    for i, (filename, exists) in enumerate(zip(response.files, response.existences)):
        digest = hashes[i].hex() if i < len(hashes) and hashes[i] else None
        result[filename] = {"exists": bool(exists), "sha256": digest}
    return result


def client_revision():
    return sha256_file(os.path.join(HERE, "dt_client.py"))[:16]


def build_reference(code):
    """Load the approved matching-rank Heart and paint out only its pip ink.

    Earlier versions blanked a rectangular region.  That could survive img2img
    as a visible white block when a generated sleeve or face drifted into the
    area.  A heart-shaped mask removes the reference symbol without reserving a
    rectangle in the composition.
    """
    source = HEART_REFERENCE[code[0]]
    image = Image.open(source).convert("RGB")
    pixels = np.asarray(image).copy()
    r, g, b = pixels[..., 0], pixels[..., 1], pixels[..., 2]
    red_ink = (r > 90) & (r > g * 1.06) & (r > b * 1.06)

    # Flood only the red component under each known pip seed. This avoids
    # touching nearby red garment/crown details and leaves no geometric patch.
    for seed_x, seed_y in ((177, 155), (480, 745)):
        stack = [(seed_x, seed_y)]
        seen = set()
        while stack:
            x, y = stack.pop()
            if (x, y) in seen or not (0 <= x < WIDTH and 0 <= y < HEIGHT):
                continue
            seen.add((x, y))
            if not red_ink[y, x]:
                continue
            pixels[y, x] = (252, 252, 250)
            for dy in (-1, 0, 1):
                for dx in (-1, 0, 1):
                    if dx or dy:
                        stack.append((x + dx, y + dy))
    return Image.fromarray(pixels), source


def prompt_for(code):
    rank, suit = code
    return "%s. Subject: %s, %s. Constraints: no rank, no text, no suit pips." % (
        STYLE, FIGURE[rank], SUIT_DIRECTION[suit]
    )


def render(stub, code, seed, strength, steps, guidance, sampler, meta_common):
    reference, reference_path = build_reference(code)
    prompt = prompt_for(code)
    config = dt.build_config(
        model=MODEL,
        width=WIDTH,
        height=HEIGHT,
        steps=steps,
        guidance_scale=guidance,
        seed=seed,
        sampler=sampler,
        strength=strength,
    )

    started = time.time()
    images = dt.generate(
        stub,
        prompt,
        NEGATIVE,
        config,
        image=np.asarray(reference),
        timeout=3600,
    )
    elapsed = time.time() - started
    if not images:
        raise RuntimeError("Draw Things returned no image for %s" % code)

    os.makedirs(OUTDIR, exist_ok=True)
    stem = code.lower()
    png_path = os.path.join(OUTDIR, stem + ".png")
    with open(png_path, "wb") as f:
        f.write(dt.decode_png(images[0]))

    metadata = dict(meta_common)
    metadata.update({
        "card": code,
        "output": os.path.relpath(png_path, ROOT),
        "output_sha256": sha256_file(png_path),
        "prompt": prompt,
        "negative_prompt": NEGATIVE,
        "reference": {
            "source": os.path.relpath(reference_path, ROOT),
            "source_sha256": sha256_file(reference_path),
            "reference_heart_pips_painted_out": True,
        },
        "settings": {
            "width": WIDTH,
            "height": HEIGHT,
            "steps": steps,
            "guidance_scale": guidance,
            "seed": seed,
            "strength": strength,
            "sampler": sampler,
            "seed_mode": "ScaleAlike",
        },
        "elapsed_seconds": round(elapsed, 1),
    })
    with open(os.path.join(OUTDIR, stem + ".json"), "w") as f:
        json.dump(metadata, f, indent=2, sort_keys=True)

    print("%s  %5.0fs  -> %s" % (code, elapsed, os.path.relpath(png_path, ROOT)), flush=True)
    return png_path


def valid_code(value):
    code = value.upper()
    if len(code) != 2 or code[0] not in RANKS or code[1] not in SUITS:
        raise argparse.ArgumentTypeError("expected one of J/Q/K plus H/D/C/S")
    if code[1] == "H":
        raise argparse.ArgumentTypeError("Hearts are already approved v3 assets")
    return code


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("cards", nargs="*", type=valid_code)
    parser.add_argument("--remaining", action="store_true", help="render all nine non-Heart courts")
    parser.add_argument("--force", action="store_true", help="overwrite existing renders")
    parser.add_argument("--seed", type=int, help="override the per-card seed")
    parser.add_argument("--strength", type=float, default=0.68)
    parser.add_argument("--steps", type=int, default=30)
    parser.add_argument("--guidance", type=float, default=4.0)
    parser.add_argument("--sampler", default="DPMPP2MTrailing")
    args = parser.parse_args()

    cards = list(REMAINING if args.remaining else args.cards)
    if not cards:
        parser.error("provide card codes or --remaining")

    stub = dt.connect()
    meta_common = {
        "generator": "Draw Things gRPC (local)",
        "drawthings_version": drawthings_version(),
        "model_file": MODEL,
        "model_hashes": checkpoint_hashes(
            stub, [MODEL, "qwen_image_vae_f16.ckpt", "qwen_2.5_vl_7b_q8p.ckpt"]
        ),
        "client_revision": client_revision(),
        "generated_at": datetime.datetime.now().astimezone().isoformat(),
        "status": "game court artwork",
    }

    for code in cards:
        existing = os.path.join(OUTDIR, code.lower() + ".png")
        if os.path.exists(existing) and not args.force:
            print("%s  exists; skipping (use --force to replace)" % code, flush=True)
            continue
        seed = args.seed if args.seed is not None else DEFAULT_SEEDS[code]
        render(stub, code, seed, args.strength, args.steps, args.guidance, args.sampler, meta_common)


if __name__ == "__main__":
    main()