#!/usr/bin/env python3
"""
golden_record.py — Extract and save winning ComfyUI generation parameters.

Reads PNG metadata embedded by ComfyUI, extracts key parameters (seed, prompt,
steps, cfg, resolution, models), and saves them as a Golden Record JSON file.
Can also regenerate a ComfyUI API workflow from a saved record.

Usage:
    # Extract parameters from a winning image
    python3 golden_record.py extract <image.png>

    # Extract and save as a Golden Record
    python3 golden_record.py save <image.png> [--name "descriptive name"] [--out <path>]

    # Show a saved Golden Record
    python3 golden_record.py show <record.json>

    # Generate a ComfyUI API workflow from a Golden Record
    python3 golden_record.py workflow <record.json> [--prompt "new prompt"] [--seed <n>]

    # List all saved Golden Records
    python3 golden_record.py list [--dir <records_dir>]
"""

import argparse
import struct
import json
import os
import sys
from datetime import datetime
from pathlib import Path

DEFAULT_RECORDS_DIR = os.path.expanduser("~/ComfyUI/golden_records")


def read_png_text_chunks(filepath):
    """Read tEXt chunks from a PNG file without external dependencies."""
    chunks = {}
    with open(filepath, "rb") as f:
        sig = f.read(8)
        if sig[:4] != b"\x89PNG":
            raise ValueError(f"Not a PNG file: {filepath}")
        while True:
            data = f.read(4)
            if len(data) < 4:
                break
            length = struct.unpack(">I", data)[0]
            chunk_type = f.read(4).decode("ascii", errors="replace")
            chunk_data = f.read(length)
            _ = f.read(4)  # CRC
            if chunk_type == "tEXt":
                null_idx = chunk_data.index(b"\x00")
                key = chunk_data[:null_idx].decode("latin-1")
                value = chunk_data[null_idx + 1 :].decode("latin-1")
                chunks[key] = value
            elif chunk_type == "IEND":
                break
    return chunks


def extract_parameters(workflow):
    """Extract key generation parameters from a ComfyUI workflow dict."""
    params = {
        "seed": None,
        "steps": None,
        "cfg": None,
        "sampler": None,
        "scheduler": None,
        "denoise": None,
        "prompt": None,
        "negative_prompt": None,
        "width": None,
        "height": None,
        "model": None,
        "clip_model": None,
        "vae_model": None,
        "filename_prefix": None,
    }

    for node_id, node in workflow.items():
        cls = node.get("class_type", "")
        inputs = node.get("inputs", {})

        if cls == "KSampler":
            params["seed"] = inputs.get("seed")
            params["steps"] = inputs.get("steps")
            params["cfg"] = inputs.get("cfg")
            params["sampler"] = inputs.get("sampler_name")
            params["scheduler"] = inputs.get("scheduler")
            params["denoise"] = inputs.get("denoise")

        elif cls == "CLIPTextEncode":
            text = inputs.get("text", "")
            if text:
                if params["prompt"] is None:
                    params["prompt"] = text
                else:
                    # First non-empty is positive, second is negative
                    params["negative_prompt"] = text
            elif params["prompt"] is not None:
                # Empty text after a filled one = negative prompt is empty
                pass

        elif cls == "EmptyLatentImage":
            params["width"] = inputs.get("width")
            params["height"] = inputs.get("height")

        elif cls == "UnetLoaderGGUF":
            params["model"] = inputs.get("unet_name")

        elif cls == "CLIPLoader":
            params["clip_model"] = inputs.get("clip_name")

        elif cls == "VAELoader":
            params["vae_model"] = inputs.get("vae_name")

        elif cls == "SaveImage":
            params["filename_prefix"] = inputs.get("filename_prefix")

    # Resolve prompt ordering: in ComfyUI, the positive prompt is the one
    # connected to KSampler's "positive" input. Walk the connections.
    positive_prompt, negative_prompt = _resolve_prompt_order(workflow)
    if positive_prompt is not None:
        params["prompt"] = positive_prompt
    if negative_prompt is not None:
        params["negative_prompt"] = negative_prompt

    return params


def _resolve_prompt_order(workflow):
    """Determine which CLIPTextEncode is positive vs negative by tracing KSampler connections."""
    positive = None
    negative = None

    for node_id, node in workflow.items():
        if node.get("class_type") == "KSampler":
            inputs = node.get("inputs", {})
            pos_ref = inputs.get("positive")
            neg_ref = inputs.get("negative")

            if isinstance(pos_ref, list) and len(pos_ref) >= 1:
                pos_node_id = str(pos_ref[0])
                pos_node = workflow.get(pos_node_id, {})
                if pos_node.get("class_type") == "CLIPTextEncode":
                    positive = pos_node.get("inputs", {}).get("text", "")

            if isinstance(neg_ref, list) and len(neg_ref) >= 1:
                neg_node_id = str(neg_ref[0])
                neg_node = workflow.get(neg_node_id, {})
                if neg_node.get("class_type") == "CLIPTextEncode":
                    negative = neg_node.get("inputs", {}).get("text", "")

            break

    return positive, negative


def build_golden_record(image_path, name=None):
    """Build a complete Golden Record from an image file."""
    chunks = read_png_text_chunks(image_path)
    if "prompt" not in chunks:
        raise ValueError(f"No ComfyUI workflow found in PNG metadata: {image_path}")

    workflow = json.loads(chunks["prompt"])
    params = extract_parameters(workflow)

    record = {
        "golden_record": {
            "name": name or Path(image_path).stem,
            "source_image": str(Path(image_path).resolve()),
            "created": datetime.now().isoformat(),
        },
        "parameters": params,
        "raw_workflow": workflow,
    }

    return record


def record_to_workflow(record, prompt_override=None, seed_override=None):
    """Reconstruct a ComfyUI API workflow from a Golden Record.

    Starts from the raw_workflow and applies any overrides.
    """
    workflow = json.loads(json.dumps(record["raw_workflow"]))  # deep copy

    for node_id, node in workflow.items():
        cls = node.get("class_type", "")
        inputs = node.get("inputs", {})

        if cls == "KSampler" and seed_override is not None:
            inputs["seed"] = seed_override

        if cls == "CLIPTextEncode" and prompt_override is not None:
            # Only override the positive prompt (non-empty one)
            if inputs.get("text", ""):
                inputs["text"] = prompt_override
                prompt_override = None  # Only replace the first non-empty

    return workflow


def save_record(record, output_path=None):
    """Save a Golden Record to disk."""
    if output_path is None:
        os.makedirs(DEFAULT_RECORDS_DIR, exist_ok=True)
        safe_name = record["golden_record"]["name"].replace(" ", "_").lower()
        output_path = os.path.join(DEFAULT_RECORDS_DIR, f"{safe_name}.json")

    os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
    with open(output_path, "w") as f:
        json.dump(record, f, indent=2)

    return output_path


def format_params(params):
    """Format parameters for human-readable display."""
    lines = []
    lines.append(f"  Seed:       {params.get('seed', '?')}")
    lines.append(f"  Steps:      {params.get('steps', '?')}")
    lines.append(f"  CFG:        {params.get('cfg', '?')}")
    lines.append(f"  Sampler:    {params.get('sampler', '?')}")
    lines.append(f"  Scheduler:  {params.get('scheduler', '?')}")
    lines.append(f"  Resolution: {params.get('width', '?')}x{params.get('height', '?')}")
    lines.append(f"  Model:      {params.get('model', '?')}")
    lines.append(f"  CLIP:       {params.get('clip_model', '?')}")
    lines.append(f"  VAE:        {params.get('vae_model', '?')}")
    lines.append(f"  Prompt:     {params.get('prompt', '?')[:120]}...")
    neg = params.get("negative_prompt", "")
    if neg:
        lines.append(f"  Negative:   {neg[:80]}")
    return "\n".join(lines)


def cmd_extract(args):
    """Extract and display parameters from a PNG."""
    chunks = read_png_text_chunks(args.image)
    if "prompt" not in chunks:
        print(f"ERROR: No ComfyUI workflow found in {args.image}", file=sys.stderr)
        sys.exit(1)

    workflow = json.loads(chunks["prompt"])
    params = extract_parameters(workflow)

    print(f"Parameters from: {args.image}")
    print(format_params(params))

    if args.json:
        print("\n--- JSON ---")
        print(json.dumps(params, indent=2))


def cmd_save(args):
    """Extract and save a Golden Record."""
    record = build_golden_record(args.image, name=args.name)
    output_path = save_record(record, args.out)

    print(f"Golden Record saved: {output_path}")
    print(f"Name: {record['golden_record']['name']}")
    print(format_params(record["parameters"]))


def cmd_show(args):
    """Display a saved Golden Record."""
    try:
        with open(args.record) as f:
            record = json.load(f)
    except FileNotFoundError:
        print(f"ERROR: Record not found: {args.record}", file=sys.stderr)
        sys.exit(1)
    except json.JSONDecodeError as e:
        print(f"ERROR: Invalid JSON in {args.record}: {e}", file=sys.stderr)
        sys.exit(1)

    meta = record.get("golden_record", {})
    print(f"Golden Record: {meta.get('name', '?')}")
    print(f"Source: {meta.get('source_image', '?')}")
    print(f"Created: {meta.get('created', '?')}")
    print(format_params(record.get("parameters", {})))


def cmd_workflow(args):
    """Generate a ComfyUI API workflow from a Golden Record."""
    try:
        with open(args.record) as f:
            record = json.load(f)
    except FileNotFoundError:
        print(f"ERROR: Record not found: {args.record}", file=sys.stderr)
        sys.exit(1)
    except json.JSONDecodeError as e:
        print(f"ERROR: Invalid JSON in {args.record}: {e}", file=sys.stderr)
        sys.exit(1)

    workflow = record_to_workflow(
        record,
        prompt_override=args.prompt,
        seed_override=args.seed,
    )

    if args.out:
        with open(args.out, "w") as f:
            json.dump(workflow, f, indent=2)
        print(f"Workflow saved: {args.out}")
    else:
        print(json.dumps(workflow, indent=2))


def cmd_list(args):
    """List all saved Golden Records."""
    records_dir = args.dir or DEFAULT_RECORDS_DIR
    if not os.path.isdir(records_dir):
        print(f"No records directory found at {records_dir}")
        sys.exit(0)

    files = sorted(Path(records_dir).glob("*.json"))
    if not files:
        print("No Golden Records saved yet.")
        sys.exit(0)

    print(f"Golden Records in {records_dir}:\n")
    for f in files:
        try:
            with open(f) as fh:
                record = json.load(fh)
            meta = record.get("golden_record", {})
            params = record.get("parameters", {})
            print(f"  {f.name}")
            print(f"    Name: {meta.get('name', '?')}")
            print(f"    Seed: {params.get('seed', '?')} | {params.get('width', '?')}x{params.get('height', '?')} | {params.get('steps', '?')} steps")
            print(f"    Prompt: {(params.get('prompt') or '?')[:80]}...")
            print()
        except (json.JSONDecodeError, KeyError):
            print(f"  {f.name} (invalid)")


def main():
    parser = argparse.ArgumentParser(
        description="Golden Record — Save and reproduce winning ComfyUI parameters"
    )
    subparsers = parser.add_subparsers(dest="command")

    # extract
    p_extract = subparsers.add_parser("extract", help="Extract parameters from a PNG")
    p_extract.add_argument("image", help="Path to ComfyUI output PNG")
    p_extract.add_argument("--json", action="store_true", help="Also output raw JSON")

    # save
    p_save = subparsers.add_parser("save", help="Save a Golden Record from a PNG")
    p_save.add_argument("image", help="Path to ComfyUI output PNG")
    p_save.add_argument("--name", help="Descriptive name for this record")
    p_save.add_argument("--out", help="Output path (default: ~/ComfyUI/golden_records/)")

    # show
    p_show = subparsers.add_parser("show", help="Display a saved Golden Record")
    p_show.add_argument("record", help="Path to Golden Record JSON")

    # workflow
    p_workflow = subparsers.add_parser("workflow", help="Generate workflow from record")
    p_workflow.add_argument("record", help="Path to Golden Record JSON")
    p_workflow.add_argument("--prompt", help="Override the prompt text")
    p_workflow.add_argument("--seed", type=int, help="Override the seed")
    p_workflow.add_argument("--out", help="Save workflow to file (default: stdout)")

    # list
    p_list = subparsers.add_parser("list", help="List saved Golden Records")
    p_list.add_argument("--dir", help="Records directory")

    args = parser.parse_args()

    if args.command == "extract":
        cmd_extract(args)
    elif args.command == "save":
        cmd_save(args)
    elif args.command == "show":
        cmd_show(args)
    elif args.command == "workflow":
        cmd_workflow(args)
    elif args.command == "list":
        cmd_list(args)
    else:
        parser.print_help()


if __name__ == "__main__":
    main()