#!/usr/bin/env python3
"""
Rewrite the settings block of the SD.cpp generation nodes from the server's own
option reference.

The server publishes every field the generation endpoints accept at
/options/generation, with a label, a type, a default, a description and which
endpoints it applies to. Typing that list into five node files by hand is how
they drifted in the first place - three of them exposed ten fields out of
forty-six and hid the rest behind an "Extra Options" JSON box.

So the list is generated. Run this against a server, check the diff, commit the
result. The nodes stay self-contained JavaScript afterwards - nothing here runs
at execution time.

  python3 scripts/gen-sdcpp-generation-options.py http://mulan:8077
"""

import json
import re
import sys
import urllib.request
from pathlib import Path

NODES = {
    'sdcpp-txt2img': 'txt2img',
    'sdcpp-img2img': 'img2img',
    'sdcpp-edit': 'img2img',
    'sdcpp-txt2vid': 'txt2vid',
    'sdcpp-upscale': 'upscale',
}

# Settings that already exist under a name of their own, kept so workflows that
# were configured before this generator do not lose their values.
KEEP_NAME = {
    'negative_prompt': 'negativePrompt',
    'cfg_scale': 'cfgScale',
    'batch_count': 'batchCount',
    'clip_skip': 'clipSkip',
    'init_image_base64': 'initImageBase64',
    'mask_image_base64': 'maskImageBase64',
    'image_base64': 'imageBase64',
    'upscale_factor': 'upscaleFactor',
    'tile_size': 'tileSize',
    'video_frames': 'videoFrames',
    'ref_images': 'refImages',
    'ref_image_args': 'refImageArgs',
}

# Fields the node handles itself rather than passing through, or that make no
# sense to type into a form.
SKIP = {'title'}

# Fields the endpoints accept that /options/generation does not describe yet.
# They are in the OpenAPI request schemas, so they are real - the per-field
# reference simply has not caught up. Declared here in the same shape the
# reference uses, so that when it does catch up these can be deleted and
# nothing else changes.
# The values map for cache_mode lists two modes; the field's own description in
# the schema names six. Neither source is machine-readable and complete, so the
# list is written here from the description, with a note of where it came from.
CACHE_MODES = {
    'easycache': 'EasyCache - single threshold, simple',
    'ucache': 'UCache',
    'dbcache': 'DBCache',
    'taylorseer': 'TaylorSeer',
    'cache_dit': 'Cache-DiT',
    'spectrum': 'Spectrum - frequency-domain analysis (best quality/speed tradeoff)',
}

UNDOCUMENTED = {
    'ip_adapter_image_base64': {
        'applies_to': ['txt2img', 'img2img', 'txt2vid'],
        'label': 'IP-Adapter Image (base64)', 'type': 'string', 'default': '',
        'description': 'A reference image whose style and subject guide the result, as base64. '
                       'Needs an IP-Adapter loaded alongside the model - see the Load Model node.',
        'recommended': 'Take it from a Download or Fetch Output node rather than pasting one in.',
        'category': 'image_input',
    },
    'ip_adapter_strength': {
        'applies_to': ['txt2img', 'img2img', 'txt2vid'],
        'label': 'IP-Adapter Strength', 'type': 'number', 'default': 1.0,
        'description': 'How strongly the reference image guides the result.',
        'recommended': '1.0 is the upstream default. Lower it when the reference is overwhelming the prompt.',
        'category': 'image_input',
    },
    'ref_audios': {
        'applies_to': ['txt2vid'],
        'label': 'Reference Audios', 'type': 'array<string>', 'default': [],
        'description': 'Reference audio as base64-encoded WAV, mono or stereo PCM (i16/i24/i32/f32).',
        'recommended': 'For models that take audio guidance.',
        'category': 'video',
    },
    'ref_videos': {
        'applies_to': ['txt2vid'],
        'label': 'Reference Videos', 'type': 'array<object>', 'default': [],
        'description': 'Reference videos. Each entry is an object with frames (base64 images), '
                       'fps (default 24) and an optional audio_wav_base64.',
        'recommended': 'Built by an earlier node rather than typed.',
        'category': 'video',
    },
}

# Fields a node takes that the reference files under a different endpoint. The
# server documents ref_images and ref_image_args as txt2img fields, but the
# image-edit node posts them to /img2img and the server accepts them - so they
# stay, rather than a node losing a working feature to a documentation table.
EXTRA_FOR = {
    'sdcpp-edit': ['ref_images', 'ref_image_args'],
}

# Which group the connection settings go in, and its position.
SERVER_GROUP = {'title': 'Server', 'fields': ['serverUrl', 'credentialId']}

# Pressing the button asks the server which architecture the loaded model is and
# writes that architecture's generation defaults into the form. Every preset
# carries width, height, steps, cfg_scale, sampler and scheduler; a few carry
# more. A field a preset does not mention is left alone rather than blanked.
PREFILL_FIELDS = [
    'width', 'height', 'steps', 'cfgScale', 'sampler', 'scheduler',
    'cacheMode', 'distilledGuidance', 'flowShift', 'negativePrompt',
    'videoFrames', 'fps',
]


def camel(name: str) -> str:
    if name in KEEP_NAME:
        return KEEP_NAME[name]
    head, *rest = name.split('_')
    return head + ''.join(w[:1].upper() + w[1:] for w in rest)


def js(value) -> str:
    """A JS literal. json.dumps is valid JS for everything used here."""
    return json.dumps(value, ensure_ascii=False)


# The server writes its descriptions with em-dashes. They are shown to people
# in our editor, where the house style is a plain hyphen, and the meaning is
# identical - so they are normalised on the way in rather than every reader
# meeting two conventions in one form.
def plain(text: str) -> str:
    return str(text).replace('\u2014', '-').replace('\u2013', '-')


def describe(opt: dict) -> str:
    text = ' '.join(plain(opt.get('description', '')).split())
    hint = ' '.join(plain(opt.get('recommended', '')).split())
    if hint:
        text = f'{text} Recommended: {hint}' if text else f'Recommended: {hint}'
    # Every one of these is optional: the server fills an absent field from the
    # loaded model's architecture preset, which is almost always the right
    # answer and is not something the node can know.
    return f'{text} Leave empty for the architecture default.'


def openapi_enums(spec: dict) -> dict:
    """The values each generation field accepts, from the request schema.

    /options/generation carries a values map with nice labels, and it has
    drifted: it lists 15 samplers spelled dpmpp2m where the server accepts 21
    spelled dpm++2m. The OpenAPI schema is generated from the running build and
    had them all, so it decides what the values are; the values map is only
    consulted for wording.

    This matters more than a missing entry in a dropdown. The server answers 202
    to any sampler name at all - including one that is simply wrong - and
    silently falls back to a default, so a misspelt value produces a different
    image with nothing to say so.
    """
    enums = {}
    schemas = spec.get('components', {}).get('schemas', {})
    for schema in schemas.values():
        for field, described in (schema.get('properties') or {}).items():
            values = described.get('enum')
            if values:
                enums.setdefault(field, list(values))
    return enums


def prop_for(name: str, opt: dict) -> dict:
    kind = opt.get('type')
    setting = camel(name)
    prop: dict = {'title': plain(opt.get('label') or setting), 'description': describe(opt)}

    if kind == 'select':
        values = opt.get('values') or {}
        authoritative = OPENAPI_ENUMS.get(name)
        if authoritative:
            # Keep the wording from the values map where there is any, but the
            # list itself comes from the schema.
            values = {v: values.get(v, v) for v in authoritative}
        elif name == 'cache_mode':
            values = dict(CACHE_MODES)
        # The empty entry is what "leave it to the architecture" looks like in a
        # dropdown; without it a select cannot express "unset". Some of the
        # server's own value maps already carry one, so it is not added twice.
        keys = [k for k in values.keys() if k != '']
        prop['type'] = 'string'
        prop['enum'] = [''] + keys
        prop['enumLabels'] = ['(architecture default)'] + [
            ' '.join(plain(values[k]).split())[:70] or k for k in keys
        ]
        prop['default'] = ''

        # For most of these, empty means "say nothing and let the loaded
        # architecture decide". For cache_mode the server also uses empty to
        # mean "off" - so leaving it empty is indistinguishable from not
        # choosing, and the preset wins. Z-Image's preset turns caching on,
        # which is why it could not be switched off. A separate entry says it
        # outright, and the body builder sends an explicit empty for it.
        if name == 'cache_mode':
            prop['enum'] = ['', 'off'] + keys
            prop['enumLabels'] = [
                '(architecture default, which may switch it on)',
                'Off - no caching, whatever the architecture prefers',
            ] + [' '.join(plain(values[k]).split())[:70] or k for k in keys]
    elif kind == 'boolean':
        prop['type'] = 'boolean'
    elif kind == 'number':
        prop['type'] = 'number'
    elif kind == 'array<string>':
        prop['type'] = 'array'
        prop['items'] = {'type': 'string'}
    elif kind == 'array<number>':
        prop['type'] = 'array'
        prop['items'] = {'type': 'number'}
    elif kind == 'array<object>':
        prop['type'] = 'array'
        prop['items'] = {'type': 'object'}
    else:
        prop['type'] = 'string'
        if name in ('prompt', 'negative_prompt'):
            prop['format'] = 'textarea'

    return prop


def render_properties(props: dict, indent: str = '        ') -> str:
    out = []
    for key, prop in props.items():
        inner = ', '.join(f'{k}: {js(v)}' for k, v in prop.items())
        out.append(f'{indent}{key}: {{ {inner} }}')
    return ',\n'.join(out)


def build(reference: dict, endpoint: str, extra: list) -> tuple:
    options = dict(reference['options'])
    categories = reference['categories']

    for name, described in UNDOCUMENTED.items():
        options[name] = {k: v for k, v in described.items() if k != 'category'}

    applicable = [
        (name, opt) for name, opt in options.items()
        if (endpoint in opt.get('applies_to', []) or name in extra) and name not in SKIP
    ]
    by_name = dict(applicable)

    props = {}
    table = []
    for name, opt in applicable:
        setting = camel(name)
        props[setting] = prop_for(name, opt)
        table.append((setting, name))

    # The undocumented fields belong in a group too, next to the ones they are
    # related to rather than dumped in "Other".
    extra_by_category: dict = {}
    for name, described in UNDOCUMENTED.items():
        if name in by_name:
            extra_by_category.setdefault(described['category'], []).append(name)

    groups = []
    for cat_key, cat in categories.items():
        fields = [camel(n) for n in list(cat['options']) + extra_by_category.get(cat_key, [])
                  if n in by_name]
        if fields:
            groups.append({'title': cat['label'], 'fields': fields})
    # Core first, then the rest as the server lists them.
    groups.sort(key=lambda g: 0 if g['title'] == 'Core' else 1)

    return props, table, groups


def rewrite(path: Path, endpoint: str, reference: dict) -> str:
    source = path.read_text()

    start = source.index('const configSchema')
    end = source.index('const inputSchema')

    props, table, groups = build(reference, endpoint, EXTRA_FOR.get(path.stem, []))

    # Node-owned settings that are not server generation options.
    tail_props = {
        'title': {'type': 'string', 'title': 'Job Title',
                  'description': 'Optional label stored with the job, useful for finding it again in the queue'},
        'extraOptions': {'type': 'object', 'title': 'Extra Options',
                         'description': 'Any other generation field passed straight through. Everything the server documents already has a setting above, so this is only needed for a field a newer server has gained'},
        'timeout': {'type': 'number', 'title': 'Timeout (ms)',
                    'description': 'Applies to queueing the job, not to the render. The call returns as soon as the job is accepted',
                    'default': 30000},
    }

    ui_groups = [SERVER_GROUP] + groups + [{'title': 'Job', 'fields': ['title', 'extraOptions', 'timeout']}]

    prefill = {
        'label': 'Take the architecture defaults',
        'description': "Fill these in from the preset for whichever model the server has loaded - the same values it would use if these were left empty",
        'node': 'sdcpp-architecture',
        'needs': ['serverUrl', 'credentialId'],
        'map': {f'defaults.{f}': f for f in PREFILL_FIELDS if f in props},
    }

    head = """const configSchema = {
    type: 'object',
    // Generated from the server's own reference at /options/generation - see
    // scripts/gen-sdcpp-generation-options.py. Every field the endpoint accepts
    // has a setting here, grouped the way the server groups them.
    uiGroups: %s,
    prefill: %s,
    properties: {
        serverUrl: {
            type: 'string', title: 'Server URL',
            description: 'Base address of the sdcpp-restapi server',
            default: 'http://localhost:8077'
        },
        credentialId: {
            type: 'string', title: 'Credential',
            description: 'A basic credential holding the sdcpp-restapi username and password',
            dynamicOptions: { source: 'credentials', filter: { type: ['sdcpp', 'basic'] } }
        },
""" % (json.dumps(ui_groups, ensure_ascii=False, indent=8).replace('\n', '\n    '),
        json.dumps(prefill, ensure_ascii=False, indent=8).replace('\n', '\n    '))

    body = render_properties(props) + ',\n' + render_properties(tail_props)
    new_schema = head + body + "\n    },\n    required: ['credentialId']\n};\n\n"

    source = source[:start] + new_schema + source[end:]

    # What the server says each choice-field accepts. Emitted because the
    # server does not check: it answers 202 to any sampler name at all,
    # including one that is simply a typo, and silently falls back to a default.
    # A generation that quietly used a different sampler than the one asked for
    # is not something anyone would notice from the result.
    choice_rows = []
    for name, opt in sorted(reference['options'].items()):
        if opt.get('type') != 'select' or (endpoint not in opt.get('applies_to', [])
                                           and name not in EXTRA_FOR.get(path.stem, [])):
            continue
        allowed = [k for k in (opt.get('values') or {}).keys() if k != '']
        choice_rows.append(f"    {js(camel(name))}: {js(allowed)}")
    choices_js = ("// What the server accepts for each choice, as it described them when this\n"
                  "// file was generated. See scripts/gen-sdcpp-generation-options.py.\n"
                  "const KNOWN_CHOICES = {\n" + ",\n".join(choice_rows) + "\n};\n")

    if 'const KNOWN_CHOICES' in source:
        source = re.sub(r'// What the server accepts for each choice.*?\n\};\n',
                        choices_js, source, flags=re.S)
    else:
        source = source.replace('async function execute(', choices_js + '\n' + 'async function execute(', 1)

    # The setting-to-field table the request is built from.
    rows = ',\n'.join(f"    {{ setting: {js(s)}, server: {js(n)} }}" for s, n in table)
    table_js = (
        "// Every generation field the server documents, and the setting it comes\n"
        "// from. Generated alongside the schema above so the two cannot drift.\n"
        "const GENERATION_OPTIONS = [\n" + rows + "\n];\n"
    )

    marker = 'async function execute('
    if 'const GENERATION_OPTIONS' in source:
        source = re.sub(r'// Every generation field the server documents.*?\n\];\n',
                        table_js, source, flags=re.S)
    else:
        source = source.replace(marker, table_js + '\n' + marker, 1)

    return source


OPENAPI_ENUMS: dict = {}


def main():
    global OPENAPI_ENUMS
    server = sys.argv[1] if len(sys.argv) > 1 else 'http://localhost:8077'
    with urllib.request.urlopen(server.rstrip('/') + '/options/generation', timeout=30) as f:
        reference = json.loads(f.read().decode())
    # Both are unauthenticated. The schema decides what a field accepts; the
    # reference above decides what it is called, what it does and where it
    # belongs.
    with urllib.request.urlopen(server.rstrip('/') + '/openapi.json', timeout=30) as f:
        OPENAPI_ENUMS = openapi_enums(json.loads(f.read().decode()))
    print(f'{len(OPENAPI_ENUMS)} field(s) have an authoritative list in the schema')

    root = Path(__file__).resolve().parent.parent / 'nodes' / 'sdcpp'
    for node, endpoint in NODES.items():
        path = root / f'{node}.js'
        path.write_text(rewrite(path, endpoint, reference))
        count = sum(1 for n, o in reference['options'].items()
                    if endpoint in o.get('applies_to', []) or n in EXTRA_FOR.get(node, []))
        print(f'{node:18} {endpoint:8} {count} options')


if __name__ == '__main__':
    main()