# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

"""Generate SpeechRecognitionModels.h and SpeechRecognitionModelDisplayInfo.sys.mjs
from models.yaml."""

import json

import yaml


def locale_array(name, locales):
    """A null-terminated static array of locale strings, 8 per line."""
    for locale in locales:
        if not isinstance(locale, str):
            raise ValueError(
                f"{name}: locale {locale!r} is not a string. YAML reads a bare "
                "no, yes, on, off, y or n as a boolean, so a locale tag that "
                "collides with one must be quoted in models.yaml."
            )
    items = [f'"{locale}"' for locale in locales] + ["nullptr"]
    lines = ["    " + ", ".join(items[i : i + 8]) for i in range(0, len(items), 8)]
    return f"static const char* const {name}[] = {{\n" + ",\n".join(lines) + "};\n"


def gen_models_header(output, input):
    with open(input) as f:
        data = yaml.safe_load(f)

    models = data["models"]

    output.write(
        "/* This file is generated by gen_speech_models.py. Do not edit. */\n"
        "#pragma once\n\n"
        "#include <stdint.h>\n\n"
        "namespace mozilla::dom {\n\n"
        "struct SpeechRecognitionModelInfo {\n"
        "  const char* id;\n"
        "  const char* const* supported_locales; // null-terminated\n"
        "  const char* repo;\n"
        "  const char* filename;\n"
        "  const char* revision;\n"
        "  uint32_t size_mb;\n"
        "  const char* quant;\n"
        "  uint32_t latency_ms;\n"
        "  bool is_streaming; // true = streaming RNN-T, false = offline CTC batch\n"
        "};\n\n"
    )

    # Emit per-model locale arrays.
    for m in models:
        output.write(
            locale_array(
                f"kSpeechModelSupportedLocales_{m['id']}", m["supported_locales"]
            )
        )

    output.write("\n")

    # Emit the table, sentinel-terminated, in models.yaml (preference) order.
    output.write(
        "static const SpeechRecognitionModelInfo kSpeechRecognitionModels[] = {\n"
    )
    for m in models:
        is_streaming = "true" if m.get("streaming", True) else "false"
        output.write(
            f'  {{"{m["id"]}", kSpeechModelSupportedLocales_{m["id"]},\n'
            f'   "{m["repo"]}", "{m["filename"]}", "{m["revision"]}",\n'
            f'   {m["size_mb"]}, "{m["quant"]}", {m["latency_ms"]}, {is_streaming}}},\n'
        )
    output.write("  {nullptr, nullptr, nullptr, nullptr, nullptr, 0, nullptr, 0}\n")
    output.write("};\n\n")

    output.write("}  // namespace mozilla::dom\n")


def gen_models_display_info(output, input):
    """Generate a JS module mapping each model's downloaded file name to the
    human-readable name and upstream Hugging Face model card shown in
    about:addons' "Manage On-Device AI Models" panel."""
    with open(input) as f:
        data = yaml.safe_load(f)

    info = {
        m["filename"]: {"name": m["display_name"], "hfUrl": m["hf_url"]}
        for m in data["models"]
    }

    output.write(
        "/* This file is generated by gen_speech_models.py. Do not edit. */\n\n"
        "export default " + json.dumps(info, indent=2) + ";\n"
    )