#!/usr/bin/env python3
"""Generate immutable JVM values and lossless adapters for the adopted user-state FBS."""

import argparse
import re
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]
SCHEMAS = ROOT / "indexer/src/main/fbs/user_state"
OUT = ROOT / "indexer/src/main/generated/kotlin/link/roadio/riff/userstate/model"


def camel(value):
    return re.sub(r"_([a-z0-9])", lambda m: m[1].upper(), value)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--check", action="store_true")
    args = parser.parse_args()

    def output(path, content):
        if args.check:
            if not path.is_file() or path.read_text() != content:
                raise SystemExit(f"Generated user-state source is stale: {path}")
        else:
            path.parent.mkdir(parents=True, exist_ok=True)
            path.write_text(content)

    paths = sorted(SCHEMAS.glob("*.fbs")) + sorted(
        (SCHEMAS.parent / "user_state_ipc").glob("*.fbs")
    )
    source = "\n".join(p.read_text() for p in paths)
    enums = {}
    for name, values in re.findall(r"enum (\w+):\w+\s*\{([^}]+)\}", source):
        counter = -1
        enum = {}
        for value in values.split(","):
            parts = value.strip().split("=")
            if not parts[0]:
                continue
            counter = int(parts[1]) if len(parts) == 2 else counter + 1
            enum[parts[0].strip()] = counter
        enums[name] = enum
    definitions = {}
    for kind, name, body in re.findall(r"(table|struct) (\w+)\s*\{([^}]+)\}", source):
        fields = []
        for field, typ, default, attributes in re.findall(
            r"(\w+):([\w\[\]]+)(?:\s*=\s*([\w.-]+))?\s*(\([^;]+\))?;", body
        ):
            fields.append((field, typ, default, "required" in attributes))
        definitions[name] = (kind, fields)
    primitive = {
        "long": "Long",
        "ulong": "Long",
        "uint": "Int",
        "ubyte": "Int",
        "bool": "Boolean",
        "float": "Float",
        "double": "Double",
    }

    def kotlin_type(typ):
        if typ.startswith("["):
            return "List<" + kotlin_type(typ[1:-1]) + ">"
        return primitive.get(
            typ,
            "Int" if typ in enums else "String" if typ == "string" else typ + "Model",
        )

    def nullable(typ, default, required, kind):
        return (
            default == "null"
            or (kind == "table" and typ in definitions and not required)
            or (typ == "string" and not required)
        )

    def default_value(typ, default, required, kind):
        if nullable(typ, default, required, kind):
            return "null"
        if typ.startswith("["):
            return "emptyList()"
        if typ in enums:
            return str(enums[typ].get(default, 0))
        if typ in definitions:
            return typ + "Model()"
        if typ == "string":
            return '""'
        value = default or ("false" if typ == "bool" else "0")
        return value + (
            "L"
            if typ in ("long", "ulong")
            else (
                "f"
                if typ == "float"
                else ".0" if typ == "double" and "." not in value else ""
            )
        )

    models = [
        "// Generated by tools/flatbuffers/generate_user_state.py; do not edit.",
        "package link.roadio.riff.userstate.model",
        "",
    ]
    for name, (kind, fields) in definitions.items():
        models.append("data class " + name + "Model(")
        for field, typ, default, required in fields:
            models.append(
                "    val "
                + camel(field)
                + ": "
                + kotlin_type(typ)
                + ("?" if nullable(typ, default, required, kind) else "")
                + " = "
                + default_value(typ, default, required, kind)
                + ","
            )
        models.append(")\n")
    code = [
        "// Generated by tools/flatbuffers/generate_user_state.py; do not edit.",
        "@file:OptIn(kotlin.ExperimentalUnsignedTypes::class)",
        "package link.roadio.riff.userstate.model",
        "",
        "import com.google.flatbuffers.kotlin.*",
        "import link.roadio.riff.state.fbs.*",
        "import link.roadio.riff.userstate.BoundedReadList",
        "import link.roadio.riff.indexer.storage.flatbuffers.ReadOnlyByteBuffer",
        "import java.nio.ByteBuffer",
        "",
        "object UserStateCodec {",
        "    private fun count(n: Int): Int { require(n in 0..1_000_000); return n }",
    ]

    def wire_value(typ, value):
        if typ in enums:
            return typ + "(" + value + ".toUByte())"
        return value + (
            {"ulong": ".toULong()", "uint": ".toUInt()", "ubyte": ".toUByte()"}.get(
                typ, ""
            )
        )

    def flatten_struct(name, expr):
        result = []
        for field, typ, _, _ in definitions[name][1]:
            val = expr + "." + camel(field)
            result.extend(
                flatten_struct(typ, val)
                if typ in definitions
                else [wire_value(typ, val)]
            )
        return result

    for name, (kind, fields) in definitions.items():
        code.append(
            f"    private fun put{name}(b: FlatBufferBuilder, v: {name}Model): Offset<{name}> {{"
        )
        if kind == "struct":
            code.append(
                f"        return {name}.create{name}(b, "
                + ", ".join(flatten_struct(name, "v"))
                + ")"
            )
        else:
            for field, typ, default, required in fields:
                f = camel(field)
                val = "v." + f
                if typ == "string":
                    code.append(
                        f"        val {f}Offset = "
                        + (
                            f"{val}?.let {{ b.createString(it) }}"
                            if not required
                            else f"b.createString({val})"
                        )
                    )
                elif typ in definitions and definitions[typ][0] == "table":
                    code.append(
                        f"        val {f}Offset = "
                        + (
                            f"{val}?.let {{ put{typ}(b, it) }}"
                            if nullable(typ, default, required, kind)
                            else f"put{typ}(b, {val})"
                        )
                    )
                elif typ.startswith("["):
                    elem = typ[1:-1]
                    if elem in definitions and definitions[elem][0] == "table":
                        code.append(
                            f"        val {f}Offset = b.createVectorOfTables({val}.map {{ put{elem}(b, it) }}.toTypedArray())"
                        )
                    else:
                        size = 32 if elem == "HistoryEntry" else 8
                        code.append(f"        b.startVector({size}, {val}.size, 8)")
                        code.append(
                            f"        {val}.asReversed().forEach {{ "
                            + (
                                f"put{elem}(b, it)"
                                if elem in definitions
                                else "b.add(it.toULong())"
                            )
                            + " }"
                        )
                        code.append(
                            f'        val {f}Offset = b.endVector<{elem if elem in definitions else "ULong"}>()'
                        )
            code.append(f"        {name}.start{name}(b)")
            for field, typ, default, required in reversed(fields):
                f = camel(field)
                method = f[0].upper() + f[1:]
                val = "v." + f
                if (
                    typ.startswith("[")
                    or typ == "string"
                    or (typ in definitions and definitions[typ][0] == "table")
                ):
                    ref = f + "Offset"
                    code.append(
                        f"        "
                        + (
                            f"{ref}?.let {{ {name}.add{method}(b, it) }}"
                            if nullable(typ, default, required, kind)
                            else f"{name}.add{method}(b, {ref})"
                        )
                    )
                elif typ in definitions:
                    code.append(
                        f"        "
                        + (
                            f"{val}?.let {{ {name}.add{method}(b, put{typ}(b, it)) }}"
                            if nullable(typ, default, required, kind)
                            else f"{name}.add{method}(b, put{typ}(b, {val}))"
                        )
                    )
                else:
                    expr = wire_value(typ, "it" if default == "null" else val)
                    code.append(
                        f"        "
                        + (
                            f"{val}?.let {{ {name}.add{method}(b, {expr}) }}"
                            if default == "null"
                            else f"{name}.add{method}(b, {expr})"
                        )
                    )
            code.append(f"        return {name}.end{name}(b)")
        code.append("    }\n")
        code.append(
            f"    private fun read{name}(v: {name}): {name}Model = {name}Model("
        )
        for field, typ, default, required in fields:
            f = camel(field)
            val = "v." + f
            if typ.startswith("["):
                elem = typ[1:-1]
                exp = f"v.{f}(i)"
                exp = (
                    f"read{elem}(checkNotNull({exp}))"
                    if elem in definitions
                    else exp + ".toLong()"
                )
                value = f"BoundedReadList(count(v.{f}Length)) {{ i -> {exp} }}"
            elif typ in definitions:
                value = f"{val}?.let(::read{typ})"
                if required or kind == "struct":
                    value = f"read{typ}(checkNotNull({val}))"
            elif typ in enums:
                value = val + ".value.toInt()"
            elif typ == "string":
                value = f"checkNotNull({val})" if required else val
            else:
                value = val + (
                    {
                        "ulong": "?.toLong()" if default == "null" else ".toLong()",
                        "uint": ".toInt()",
                        "ubyte": ".toInt()",
                    }.get(typ, "")
                )
            code.append(f"        {f} = {value},")
        code.append("    )\n")
    roots = {
        "ProfileState": "PFST",
        "ListeningHistory": "PHST",
        "InstallationState": "INST",
        "StateManifest": "SMNF",
        "StatePointer": "SPTR",
        "UserStateCommand": "UCMD",
    }
    for name, identifier in roots.items():
        code.extend(
            [
                f"    fun encode(v: {name}Model): ByteArray {{",
                "        val b = FlatBufferBuilder(1024)",
                f'        b.finish(put{name}(b, v), "{identifier}")',
                "        return b.sizedByteArray()",
                "    }",
                f"    fun decode{name}(bytes: ByteArray): {name}Model = read{name}(ByteBuffer.wrap(bytes.copyOf()).asReadOnlyBuffer())",
                f"    fun read{name}(bytes: ByteBuffer): {name}Model {{",
                "        require(bytes.isReadOnly)",
                "        val view = bytes.slice().asReadOnlyBuffer()",
                f'        link.roadio.riff.userstate.UserStateWireVerifier.verify(view, "{name}")',
                "        val buffer = ReadOnlyByteBuffer(view)",
                f"        require({name}.{name}BufferHasIdentifier(buffer))",
                f"        return read{name}({name}.asRoot(buffer))",
                "    }",
            ]
        )
    code.append("}")
    sizes = {
        "long": (8, 8),
        "ulong": (8, 8),
        "double": (8, 8),
        "float": (4, 4),
        "uint": (4, 4),
        "ubyte": (1, 1),
        "bool": (1, 1),
    }
    layouts = {}

    def layout(name):
        if name in sizes:
            return sizes[name]
        if name in enums:
            return (1, 1)
        if name in layouts:
            return layouts[name][0]
        kind, fields = definitions[name]
        if kind == "table":
            return (4, 4)
        offset, alignment, placed = 0, 1, []
        for field, typ, default, required in fields:
            size, align = layout(typ)
            offset = ((offset + align - 1) // align) * align
            placed.append(offset)
            offset += size
            alignment = max(alignment, align)
        result = (((offset + alignment - 1) // alignment) * alignment, alignment)
        layouts[name] = result, placed
        return result

    def wire_type(typ):
        return (
            "enum:" + ",".join(str(n) for n in enums[typ].values())
            if typ in enums
            else typ
        )

    shapes = [
        "// Generated by tools/flatbuffers/generate_user_state.py; do not edit.",
        "package link.roadio.riff.userstate",
        "",
        "internal object UserStateWireTypes {",
        "    val shapes = mapOf(",
    ]
    for name, (kind, fields) in definitions.items():
        size = layout(name)[0] if kind == "struct" else 0
        positions = layouts[name][1] if kind == "struct" else list(range(len(fields)))
        values = [
            f'WireField({pos}, "{wire_type(typ)}", {str(required).lower()})'
            for pos, (_, typ, _, required) in zip(positions, fields)
        ]
        shapes.append(
            f'        "{name}" to WireShape({size}, listOf(' + ", ".join(values) + ")),"
        )
    shapes.extend(["    )", "}"])
    wire_output = (
        ROOT
        / "indexer/src/main/generated/kotlin/link/roadio/riff/userstate/UserStateWireTypes.kt"
    )
    output(wire_output, "\n".join(shapes) + "\n")
    output(OUT / "UserStateModels.kt", "\n".join(models).rstrip() + "\n")
    output(OUT / "UserStateCodec.kt", "\n".join(code) + "\n")


if __name__ == "__main__":
    main()