#!/usr/bin/env python
"""
csv_to_styles.py - Convert a styles CSV into the JSON format the RM Styles nodes read.

Run with no arguments to get file-picker dialogs for the input CSV and the output
JSON location:

    python tools/csv_to_styles.py

Or pass paths directly:

    python tools/csv_to_styles.py "C:\\...\\styles\\StylesSFW.csv" "C:\\...\\my_styles.json"

The converter is built around the real defects found in the ComfyUI styles CSVs on
this machine, not a hypothetical clean file. It repairs:

  * Encoding      - files saved as cp1252/latin-1 that fail to decode as UTF-8
                    (Styles.csv has a 0xa0 non-breaking space at byte 45146)
  * BOM           - leading UTF-8 byte-order mark
  * Headers       - Number/Name/Positive/Negative/Lora, the A1111
                    name/prompt/negative_prompt shape, and a trailing comma in the
                    header row producing a nameless column
  * Lora syntax   - "<lora:Raven.safetensors:1)" (closing paren instead of >),
                    space-separated instead of comma-separated entries, leading-dot
                    weights (".8"), stray trailing commas
  * Text          - non-breaking spaces, control characters, embedded newlines,
                    doubled commas, leading/trailing whitespace and commas
  * Numbering     - blank, non-numeric, zero, negative and duplicate style numbers
                    (the node silently drops number <= 0 and lets a duplicate
                    overwrite the earlier style)
  * Junk rows     - fully empty rows and placeholder rows such as "999,,,,,"

Every style is emitted with the full field set the node expects, so no output is
ever missing a key.
"""

import argparse
import json
import os
import re
import sys
import unicodedata
from pathlib import Path

# Encodings tried in order. utf-8-sig first so a BOM is consumed rather than
# becoming part of the first header name.
ENCODINGS = ("utf-8-sig", "utf-8", "cp1252", "latin-1")

# Header aliases -> canonical field. Compared lowercased with punctuation stripped.
HEADER_ALIASES = {
    "number": "number", "num": "number", "#": "number", "id": "number", "index": "number",
    "name": "name", "title": "name", "style": "name", "stylename": "name",
    "positive": "positive", "prompt": "positive", "positiveprompt": "positive",
    "pos": "positive", "text": "positive",
    "negative": "negative", "negativeprompt": "negative", "neg": "negative",
    "motion": "motion", "motionprompt": "motion", "video": "motion",
    "lora": "imageLoras", "loras": "imageLoras", "imagelora": "imageLoras",
    "imageloras": "imageLoras", "pdxllora": "imageLoras", "pdxlloras": "imageLoras",
    "motionlora": "motionLoras", "motionloras": "motionLoras",
    "motionlorahigh": "motionLorasHigh", "motionlorashigh": "motionLorasHigh",
    "motionhigh": "motionLorasHigh",
    "motionloralow": "motionLorasLow", "motionloraslow": "motionLorasLow",
    "motionlow": "motionLorasLow",
}

TEXT_FIELDS = ("name", "positive", "negative", "motion")
LORA_FIELDS = ("imageLoras", "motionLoras", "motionLorasHigh", "motionLorasLow")

# Tolerant lora matcher. Accepts >, ), ] or end-of-string as the terminator, any
# amount of internal whitespace, and weights written as "1", "0.8" or ".8".
LORA_RE = re.compile(
    r"<\s*lora\s*:\s*([^:<>)\]]+?)\s*:\s*([0-9]*\.?[0-9]+)\s*(?:[>)\]]|$)",
    re.IGNORECASE,
)


class Report:
    """Collects what was repaired so the run can be summarised honestly."""

    def __init__(self):
        self.counts = {}
        self.notes = []

    def fix(self, label, n=1):
        if n:
            self.counts[label] = self.counts.get(label, 0) + n

    def note(self, msg):
        self.notes.append(msg)


def read_csv_text(path, forced_encoding=None, report=None):
    """Read the CSV as text, finding an encoding that actually decodes it."""
    raw = Path(path).read_bytes()

    if forced_encoding:
        return raw.decode(forced_encoding), forced_encoding

    for enc in ENCODINGS:
        try:
            text = raw.decode(enc)
        except UnicodeDecodeError:
            continue
        if report and enc not in ("utf-8", "utf-8-sig"):
            report.fix(f"decoded as {enc} (not valid UTF-8)")
        if report and enc == "utf-8-sig" and raw[:3] == b"\xef\xbb\xbf":
            report.fix("stripped UTF-8 BOM")
        return text, enc

    # latin-1 maps every byte, so this is effectively unreachable
    return raw.decode("utf-8", errors="replace"), "utf-8 (with replacements)"


def normalise_key(key):
    """Reduce a header cell to a comparable token."""
    if key is None:
        return ""
    return re.sub(r"[^a-z0-9#]", "", key.strip().lower())


def map_headers(fieldnames, report):
    """Map the CSV's headers onto canonical field names."""
    mapping = {}
    unknown = []
    for raw in fieldnames or []:
        key = normalise_key(raw)
        if not key:
            # trailing comma in the header row yields a nameless column
            if raw is not None:
                report.fix("ignored nameless trailing column")
            continue
        canonical = HEADER_ALIASES.get(key)
        if canonical:
            # first header wins if a file repeats a column
            mapping.setdefault(raw, canonical)
        else:
            unknown.append(raw)

    if unknown:
        report.note(f"Ignored unrecognised column(s): {', '.join(repr(u) for u in unknown)}")
    return mapping


def clean_text(value, report=None):
    """Normalise a free-text prompt cell into something safe for the node."""
    if value is None:
        return ""

    original = value
    text = unicodedata.normalize("NFC", str(value))

    if "\xa0" in text:
        if report:
            report.fix("replaced non-breaking space")
        text = text.replace("\xa0", " ")

    # Embedded newlines and control characters survive CSV quoting but make the
    # prompt unreadable in the editor and mangle the node's console preview.
    if any(ch == "\n" or ch == "\r" for ch in text):
        if report:
            report.fix("flattened embedded newline")
    text = re.sub(r"[\r\n\t]+", " ", text)

    stripped = "".join(ch for ch in text if ch >= " " or ch == " ")
    if stripped != text and report:
        report.fix("removed control character")
    text = stripped

    text = re.sub(r"[ ]{2,}", " ", text)

    # ",,"  and stray edge commas are harmless to JSON but produce empty prompt
    # tokens that some samplers weight oddly.
    collapsed = re.sub(r"\s*,(\s*,)+", ",", text)
    if collapsed != text and report:
        report.fix("collapsed doubled comma")
    text = collapsed

    text = text.strip().strip(",").strip()

    if text != original and report is not None and text == original.strip():
        report.fix("trimmed whitespace")

    return text


def parse_loras(value, report=None):
    """Parse a lora cell into [{'path': ..., 'weight': ...}], repairing bad syntax."""
    if not value:
        return []

    text = clean_text(value)
    if not text:
        return []

    loras = []
    matches = list(LORA_RE.finditer(text))

    if matches:
        # Flag the malformed spellings we actually repaired
        for m in matches:
            terminator = m.group(0)[-1]
            if terminator in ")]":
                if report:
                    report.fix("repaired lora closing bracket")
            elif terminator not in ">":
                if report:
                    report.fix("repaired unterminated lora tag")

        if len(matches) > 1:
            between = text[matches[0].end():matches[1].start()]
            if between.strip() == "" and report:
                report.fix("split space-separated lora tags")

        for m in matches:
            path = m.group(1).strip().strip(",").strip()
            if not path:
                continue
            try:
                weight = float(m.group(2))
            except ValueError:
                weight = 1.0
                if report:
                    report.fix("defaulted unparseable lora weight to 1.0")
            loras.append({"path": path, "weight": round(weight, 4)})
        return loras

    # No <lora:...> syntax at all - treat as bare comma-separated filenames
    for part in text.split(","):
        part = part.strip()
        if not part:
            continue
        loras.append({"path": part, "weight": 1.0})
    if loras and report:
        report.fix("converted bare lora name(s) to tagged form")
    return loras


def row_is_empty(style):
    """A row with no prompt, no name and no loras carries no information."""
    if any(style[f] for f in TEXT_FIELDS):
        return False
    if any(style[f] for f in LORA_FIELDS):
        return False
    return True


def convert(csv_path, forced_encoding=None, start_number=1):
    """Convert a CSV file into the node's styles structure. Returns (data, report)."""
    import csv as csv_module

    report = Report()
    text, encoding = read_csv_text(csv_path, forced_encoding, report)
    report.note(f"Read {csv_path} as {encoding}")

    reader = csv_module.DictReader(text.splitlines(True))
    header_map = map_headers(reader.fieldnames, report)

    if "positive" not in header_map.values():
        raise SystemExit(
            f"Could not find a prompt column in {csv_path}.\n"
            f"  Headers found: {reader.fieldnames}\n"
            f"  Expected one of: Positive / Prompt / Positive Prompt / Text"
        )

    parsed = []
    dropped_empty = 0

    for raw_row in reader:
        style = {
            "number": None, "name": "", "positive": "", "negative": "", "motion": "",
            "imageLoras": [], "motionLoras": [], "motionLorasHigh": [], "motionLorasLow": [],
        }
        raw_number = None

        for column, canonical in header_map.items():
            cell = raw_row.get(column)
            if canonical == "number":
                raw_number = cell
            elif canonical in LORA_FIELDS:
                style[canonical] = parse_loras(cell, report)
            else:
                style[canonical] = clean_text(cell, report)

        if row_is_empty(style):
            dropped_empty += 1
            continue

        style["number"] = raw_number
        parsed.append(style)

    if dropped_empty:
        report.fix("dropped empty/placeholder row", dropped_empty)

    # --- numbering -------------------------------------------------------
    # Keep valid, unique, positive numbers; assign fresh ones to everything else.
    taken = set()
    needs_number = []

    for style in parsed:
        raw = (style["number"] or "").strip() if isinstance(style["number"], str) else style["number"]
        value = None
        if raw not in (None, ""):
            try:
                value = int(float(raw))
            except (TypeError, ValueError):
                report.fix("replaced non-numeric style number")
                value = None

        if value is None:
            needs_number.append(style)
            continue
        if value < 1:
            # the node skips any style with number <= 0
            report.fix("renumbered style with number <= 0")
            needs_number.append(style)
            continue
        if value in taken:
            # the node keys styles by number, so a duplicate silently wins
            report.fix("renumbered duplicate style number")
            needs_number.append(style)
            continue

        style["number"] = value
        taken.add(value)

    next_free = max(start_number, 1)
    for style in needs_number:
        while next_free in taken:
            next_free += 1
        style["number"] = next_free
        taken.add(next_free)

    parsed.sort(key=lambda s: s["number"])

    data = {"styles": [
        {
            "number": s["number"],
            "name": s["name"],
            "positive": s["positive"],
            "negative": s["negative"],
            "motion": s["motion"],
            "imageLoras": s["imageLoras"],
            "motionLoras": s["motionLoras"],
            "motionLorasHigh": s["motionLorasHigh"],
            "motionLorasLow": s["motionLorasLow"],
        }
        for s in parsed
    ]}

    return data, report


def verify(json_path):
    """Re-read the written file exactly the way the node does, and report the count."""
    with open(json_path, "r", encoding="utf-8") as f:
        data = json.load(f)

    loadable = {}
    for style in data.get("styles", []):
        num = style.get("number")
        if num and num > 0:
            loadable[num] = style
    return loadable


# --- file pickers --------------------------------------------------------

# Set when a dialog or console prompt was used, so the launcher knows to hold the
# window open afterwards (a dropped CSV still needs a Save As dialog).
USED_DIALOG = False


def pick_files(csv_path, json_path):
    """Fill in whichever paths weren't given, using dialogs when available."""
    global USED_DIALOG

    if csv_path and json_path:
        return csv_path, json_path

    USED_DIALOG = True

    try:
        import tkinter as tk
        from tkinter import filedialog
    except Exception:
        # No GUI available - fall back to prompting on the console
        if not csv_path:
            csv_path = input("Path to source CSV: ").strip().strip('"').strip("'")
        if not json_path:
            json_path = input("Path for output JSON: ").strip().strip('"').strip("'")
        return csv_path, json_path

    root = tk.Tk()
    root.withdraw()
    root.attributes("-topmost", True)

    # tools/ -> ComfyUI-RMAutomation -> custom_nodes -> ComfyUI root
    try:
        default_dir = str(Path(__file__).resolve().parents[3] / "styles")
    except IndexError:
        default_dir = ""
    if not default_dir or not Path(default_dir).is_dir():
        default_dir = os.getcwd()

    if not csv_path:
        csv_path = filedialog.askopenfilename(
            title="Select the styles CSV to convert",
            initialdir=default_dir,
            filetypes=[("CSV files", "*.csv"), ("All files", "*.*")],
        )
        if not csv_path:
            root.destroy()
            raise SystemExit("Cancelled - no input file selected.")

    if not json_path:
        suggested = Path(csv_path).stem + ".json"
        json_path = filedialog.asksaveasfilename(
            title="Save the converted styles JSON as",
            initialdir=str(Path(csv_path).parent),
            initialfile=suggested,
            defaultextension=".json",
            filetypes=[("JSON files", "*.json"), ("All files", "*.*")],
        )
        if not json_path:
            root.destroy()
            raise SystemExit("Cancelled - no output file selected.")

    root.destroy()
    return csv_path, json_path


def main():
    parser = argparse.ArgumentParser(
        description="Convert a styles CSV into RM Styles JSON.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="Run with no arguments to pick both files with a dialog.",
    )
    parser.add_argument("csv", nargs="?", help="source CSV (prompts if omitted)")
    parser.add_argument("json", nargs="?", help="destination JSON (prompts if omitted)")
    parser.add_argument("--encoding", help="force a source encoding instead of detecting")
    parser.add_argument("--start-number", type=int, default=1,
                        help="first number to hand out when renumbering (default 1)")
    parser.add_argument("--dry-run", action="store_true",
                        help="report what would change without writing the file")
    args = parser.parse_args()

    csv_path, json_path = pick_files(args.csv, args.json)

    csv_path = os.path.expandvars(os.path.expanduser(str(csv_path).strip().strip('"').strip("'")))
    json_path = os.path.expandvars(os.path.expanduser(str(json_path).strip().strip('"').strip("'")))

    if not Path(csv_path).is_file():
        raise SystemExit(f"Input file not found: {csv_path}")

    data, report = convert(csv_path, args.encoding, args.start_number)

    print()
    print("=" * 68)
    print(f"  {Path(csv_path).name}  ->  {Path(json_path).name}")
    print("=" * 68)
    for note in report.notes:
        print(f"  {note}")
    print()

    if report.counts:
        print("  Repairs applied:")
        for label, count in sorted(report.counts.items(), key=lambda kv: -kv[1]):
            print(f"    {count:>5}  {label}")
    else:
        print("  No repairs needed.")
    print()

    styles = data["styles"]
    with_loras = sum(1 for s in styles if s["imageLoras"])
    with_names = sum(1 for s in styles if s["name"])
    numbers = [s["number"] for s in styles]

    print(f"  Styles converted : {len(styles)}")
    if numbers:
        print(f"  Number range     : {min(numbers)}-{max(numbers)}")
    print(f"  With a name      : {with_names}")
    print(f"  With image loras : {with_loras}")
    print()

    if args.dry_run:
        print("  --dry-run: nothing written.")
        return

    out = Path(json_path)
    os.makedirs(out.parent, exist_ok=True)
    with open(out, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2, ensure_ascii=False)

    loadable = verify(out)
    print(f"  Written to {out}")
    print(f"  Verified: the node will load {len(loadable)} of {len(styles)} styles.")

    if len(loadable) != len(styles):
        missing = [s["number"] for s in styles if s["number"] not in loadable]
        print(f"  WARNING: unreadable style numbers: {missing[:20]}")
        sys.exit(1)
    print()
    print("  Point a RM Styles node's custom_json_path at that file.")
    print()


def run():
    """Entry point that survives being double-clicked.

    Launched from Explorer the console closes the moment the process exits, which
    would take the repair summary - or a traceback - with it. Hold the window open
    whenever the run was interactive: no arguments at all, or any run that had to
    ask for a path (dragging a CSV onto the launcher still needs a Save As dialog).
    """
    interactive = len(sys.argv) == 1
    code = 0

    try:
        main()
    except SystemExit as exc:
        # argparse and this script's own error paths both exit this way; print
        # the message here so it isn't swallowed along with the window.
        if isinstance(exc.code, str):
            print(f"\n{exc.code}")
            code = 1
        elif exc.code:
            code = exc.code
    except KeyboardInterrupt:
        print("\nCancelled.")
        code = 1
    except Exception:
        import traceback
        traceback.print_exc()
        code = 1

    if interactive or USED_DIALOG:
        try:
            input("Press Enter to close this window...")
        except EOFError:
            pass

    sys.exit(code)


if __name__ == "__main__":
    run()