#!/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 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)


def describe(opt: dict) -> str:
    text = ' '.join(str(opt.get('description', '')).split())
    hint = ' '.join(str(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 prop_for(name: str, opt: dict) -> dict:
    kind = opt.get('type')
    setting = camel(name)
    prop: dict = {'title': opt.get('label') or setting, 'description': describe(opt)}

    if kind == 'select':
        values = opt.get('values') or {}
        # 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(str(values[k]).split())[:70] or k for k in keys
        ]
        prop['default'] = ''
    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'}
    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 = reference['options']
    categories = reference['categories']

    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))

    groups = []
    for _, cat in categories.items():
        fields = [camel(n) for n in cat['options'] 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:]

    # 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


def main():
    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())

    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()