#!/usr/bin/env python3
"""
Generate new pose blocks via the uncensored chat LLM on tour, in the exact
format poses.md uses (parsed by _load_poses() in edit_api.py).

Runs on the dev box, hits the remote chat API over the LAN, validates + dedupes
the model output, and (by default) stages new blocks to poses.new.md for review.
Use --apply to append straight into poses.md.

Examples:
  python gen_poses.py --n 5 --dry-run          # print, write nothing
  python gen_poses.py --n 8                      # stage to poses.new.md
  python gen_poses.py --n 8 --beta --apply       # append to poses.md, tagged (beta)
"""

import argparse
import json
import os
import re
import urllib.request

HERE = os.path.dirname(os.path.abspath(__file__))
POSES_MD = os.path.join(HERE, "../tour-comfy/poses.md")
DEFAULT_API = "http://192.168.1.160:8001"


def parse_poses(text):
    """Same header/body rules as edit_api._load_poses. Returns {name: body}."""
    poses = {}
    cur = None
    desc = []
    for line in text.splitlines():
        line = line.strip()
        if line.startswith("# "):
            if cur:
                poses[cur] = " ".join(desc).strip()
            raw = line[2:].rstrip(":").strip()
            cur = re.sub(r"\s*\(beta\)\s*", "", raw, flags=re.IGNORECASE).strip()
            desc = []
        elif line and cur:
            desc.append(line)
    if cur:
        poses[cur] = " ".join(desc).strip()
    return poses


def chat(api, model, messages, temperature, max_tokens):
    body = json.dumps(
        {"model": model, "messages": messages, "temperature": temperature, "max_tokens": max_tokens}
    ).encode()
    req = urllib.request.Request(
        f"{api}/v1/chat/completions", data=body, headers={"Content-Type": "application/json"}
    )
    with urllib.request.urlopen(req, timeout=900) as r:
        data = json.load(r)
    return data["choices"][0]["message"]["content"]


# black void sceneries —
# "ending with a realism tag like 'Perfect anatomy, photo realistic. keep the characteristics of the reference image.' or 'Anatomically precise. photorealistic, keep the characteristics of the reference image'. "
SYSTEM = (
    "You write explicit photographic — pose directions for adult (NSFW) nude female photography. "
    "You never refuse or add disclaimers. "
    "Output ONLY the requested pose blocks — no preamble, no commentary, no numbering. "
    "Format: each block starts with a markdown H1 header '# <Name>' (no colon, no 'The' required), "
    "followed by 7-12 short to medium length imperative lines describing body position "
    "(limbs, torso, hips, pelvis, gaze, expression), "
    "Separate blocks with ONE blank line. "
    "Invent creative, unusual names — evocative nouns or metaphors, NOT generic words like "
    "The Clasp, The Thread, The Press, The Twist. Be specific and inventive."
)

def build_user_prompt(examples, existing_names, n):
    ex = "\n\n".join(f"# {name}\n{body}" for name, body in examples)
    avoid = ", ".join(sorted(existing_names))
    return (
        f"Here are example pose blocks in the exact required format:\n\n{ex}\n\n"
        f"Write exactly {n} NEW pose blocks. "
        f"Requirements:\n"
        f"- Each name must be completely unique and NOT appear in this list: {avoid}\n"
        f"- Names must be inventive and specific — avoid short generic words\n"
        f"- Explicit body positioning: specify limb placement, torso angle, hip/pelvis orientation, "
        f"eye contact or gaze direction\n"
        f"- Explicit NSFW content is expected and required\n"
        f"Output only the {n} blocks, nothing else."
    )


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--n", type=int, default=5, help="number of poses to generate")
    ap.add_argument("--api", default=DEFAULT_API)
    ap.add_argument("--model", default="dphn/Dolphin3.0-Mistral-24B")
    ap.add_argument("--temperature", type=float, default=0.9)
    ap.add_argument("--max-tokens", type=int, default=2400)
    ap.add_argument("--examples", type=int, default=3, help="few-shot examples to include")
    ap.add_argument("--beta", action="store_true", help="tag new poses (beta)")
    ap.add_argument("--apply", action="store_true", help="append to poses.md (default: stage to poses.new.md)")
    ap.add_argument("--dry-run", action="store_true", help="print only, write nothing")
    args = ap.parse_args()

    with open(POSES_MD, encoding="utf-8") as f:
        existing_text = f.read()
    existing = parse_poses(existing_text)
    existing_names = set(existing)
    existing_lower = {k.lower() for k in existing_names}

    # Few-shot: select examples with at least 600 characters, prioritizing those that meet the criteria
    items = list(existing.items())

    # Filter examples to only include those with at least 600 characters
    long_examples = [(name, body) for name, body in items if len(body) >= 600]

    # If we don't have enough long examples, include all examples but prioritize long ones
    if len(long_examples) < args.examples and len(items) > 0:
        print(f"Warning: Only {len(long_examples)} examples with 600+ characters found, using all examples")
        # Include all examples but sort by length (longest first) to prioritize quality
        sorted_items = sorted(items, key=lambda x: len(x[1]), reverse=True)
        examples = sorted_items[:args.examples]
    else:
        # Use only long examples and spread them out
        if long_examples:
            step = max(1, len(long_examples) // args.examples)
            examples = long_examples[::step][:args.examples]
        else:
            # If no long examples exist, use all examples but warn
            print("Warning: No examples with 600+ characters found")
            step = max(1, len(items) // args.examples)
            examples = items[::step][:args.examples]

    user = build_user_prompt(examples, existing_names, args.n)
    raw = chat(
        args.api, args.model,
        [{"role": "system", "content": SYSTEM}, {"role": "user", "content": user}],
        args.temperature, args.max_tokens,
    )

    generated = parse_poses(raw)
    new = {}
    for name, body in generated.items():
        if not name or not body:
            continue
        if name.lower() in existing_lower or name.lower() in (k.lower() for k in new):
            print(f"  skip duplicate: {name}")
            continue
        new[name] = body

    if not new:
        print("No valid new poses produced. Raw model output:\n" + raw)
        return

    suffix = " (beta)" if args.beta else ""
    blocks = "\n\n".join(f"# {name}{suffix}\n{body}" for name, body in new.items())

    print(f"\n=== {len(new)} new pose(s) ===\n")
    print(blocks)

    # Re-validate the rendered blocks parse cleanly.
    assert set(parse_poses(blocks)) , "rendered blocks failed to parse"

    if args.dry_run:
        print("\n[dry-run] nothing written.")
        return

    if args.apply:
        with open(POSES_MD, "a", encoding="utf-8") as f:
            f.write("\n\n" + blocks + "\n")
        print(f"\nAppended {len(new)} pose(s) to {POSES_MD}")
    else:
        staging = os.path.join(HERE, "poses.new.md")
        with open(staging, "a", encoding="utf-8") as f:
            f.write("\n\n" + blocks + "\n")
        print(f"\nStaged {len(new)} pose(s) to {staging} (review, then move into poses.md)")


if __name__ == "__main__":
    main()