#!/usr/bin/env python

import enum
import re


# Parse net9p_defs.txt #########################################################

class Atom(enum.Enum):
    u8 = 1
    u16 = 2
    u32 = 3
    u64 = 4


class Struct:
    name: str
    members: list["Member"]


class List:
    cnt: str
    typ: Atom | Struct

    def __init__(self, /, *, cnt: str, typ: Atom | Struct) -> None:
        self.cnt = cnt
        self.typ = typ


class Member:
    name: str
    typ: Atom | Struct | List

    def __init__(self, /, *, name: str, typ: Atom | Struct | List) -> None:
        self.name = name
        self.typ = typ


def parse_members(
    env: dict[str, Atom | Struct], existing: list[Member], specs: str
) -> list[Member]:
    ret = existing
    for spec in specs.split():
        m = re.fullmatch(r"(.+)\[([^*]+)(?:\*([^*]+))?\]", spec)
        if not m:
            raise SyntaxError(f"invalid member spec {repr(spec)}")
        if m.group(2) not in env:
            raise NameError(f"Unknown type {repr(m.group(2))}")
        name = m.group(1)
        typ = env[m.group(2)]
        if any(x.name == name for x in ret):
            raise ValueError(f"duplicate member name {repr(name)}")
        if cnt := m.group(3):
            if len(ret) == 0 or ret[-1].name != cnt:
                raise ValueError(f"list count must be previous item: {repr(cnt)}")
            if not isinstance(ret[-1].typ, Atom):
                raise ValueError(f"list count must be an integer type: {repr(cnt)}")
            ret += [Member(name=name, typ=List(cnt=cnt, typ=typ))]
        else:
            ret += [Member(name=name, typ=typ)]
    return ret


class Message:
    id: int
    name: str
    members: list[Member]


def parse_file(filename: str) -> tuple[list[Struct], list[Message]]:
    msgs: list[Message] = []
    env: dict[str, Atom | Struct] = {
        "1": Atom.u8,
        "2": Atom.u16,
        "4": Atom.u32,
        "8": Atom.u64,
    }
    with open(filename, "r") as fh:
        prev: Struct | Message | None = None
        for line in fh:
            line = line.split("#", 1)[0].strip()
            if not line:
                continue
            if m := re.fullmatch(r'([0-9]+)\s*=\s*(\S+)\s*"([^"]*)"', line):
                msg = Message()
                msg.id = int(m.group(1))
                msg.name = m.group(2)
                msg.members = parse_members(env, [], m.group(3))
                msgs += [msg]
                prev = msg
            elif m := re.fullmatch(r'(\S+)\s*=\s*"([^"]*)"', line):
                struct = Struct()
                struct.name = m.group(1)
                struct.members = parse_members(env, [], m.group(2))
                env[struct.name] = struct
                prev = struct
            elif m := re.fullmatch(r'"([^"]*)"', line):
                if not prev:
                    raise SyntaxError(
                        "a continuation line must come after a struct line"
                    )
                prev.members = parse_members(env, prev.members, line.strip('"'))
            else:
                raise SyntaxError(f"invalid line {repr(line)}")
    structs = [x for x in env.values() if isinstance(x, Struct)]
    return structs, msgs

# Generate C ###################################################################

def shortname(typ: Atom | Struct | Message) -> str:
    match typ:
        case Atom.u8:
            return "1"
        case Atom.u16:
            return "2"
        case Atom.u32:
            return "4"
        case Atom.u64:
            return "8"
        case Struct():
            return typ.name
        case Message():
            return 'msg_'+typ.name
        case _:
            raise ValueError(f"not a type: {typ.__class__.__name__}")

def c_typename(typ: Atom | Struct | List | Message) -> str:
    match typ:
        case Atom.u8:
            return "uint8_t"
        case Atom.u16:
            return "uint16_t"
        case Atom.u32:
            return "uint32_t"
        case Atom.u64:
            return "uint64_t"
        case Struct():
            return "struct v9fs_" + typ.name
        case Message():
            return "struct v9fs_msg_" + typ.name
        case List():
            return c_typename(typ.typ) + "*"
        case _:
            raise ValueError(f"not a type: {typ.__class__.__name__}")


def static_size(typ: Atom | Struct | List | Message) -> int | None:
    match typ:
        case Atom.u8:
            return 1
        case Atom.u16:
            return 2
        case Atom.u32:
            return 4
        case Atom.u64:
            return 8
        case Struct() | Message():
            size = 0
            for member in typ.members:
                msize = static_size(member.typ)
                if msize is None:
                    return None
                size += msize
            return size
        case List():
            return None
        case _:
            raise ValueError(f"not a type: {typ.__class__.__name__}")

        
def gen_h(structs: list[Struct], msgs: list[Message]) -> str:
    ret = ""
    ret += "/* Generated by ./net9p_defs.gen.  DO NOT EDIT!  */\n"
    ret += "\n"
    ret += "#ifndef _NET9P_DEFS_H_\n"
    ret += "#define _NET9P_DEFS_H_\n"
    ret += "\n"
    ret += "#define V9FS_EWRONGSIZE 7\n"
    ret += "#define V9FS_EBADTYPE   4\n"

    ret += "\n"
    for struct in structs:
        ret += c_typename(struct) + " {\n"
        typewidth = max(len(c_typename(member.typ)) for member in struct.members)
        for member in struct.members:
            ret += f"\t{c_typename(member.typ).ljust(typewidth)} {member.name};\n"
        ret += "};\n"

    ret += "\n"
    ret += "enum v9fs_msg_type {\n"
    namewidth = max(len(msg.name) for msg in msgs)
    for msg in msgs:
        ret += f"\tV9FS_TYP_{msg.name.ljust(namewidth)} = {msg.id},\n"
    ret += "};\n"

    ret += "\n"
    for msg in msgs:
        if not msg.members:
            ret += c_typename(msg) + " {};\n"
        else:
            ret += c_typename(msg) + " {\n"
            typewidth = max(len(c_typename(member.typ)) for member in msg.members)
            for member in msg.members:
                ret += f"\t{c_typename(member.typ).ljust(typewidth)} {member.name};\n"
            ret += "};\n"

    ret += """
/**
 * @param net_bytes : the complete request, starting with the "size[4]"
 * @param out_tag : the message-ID tag
 * @param out_body : a pointer that get set to the parsed body, whose
 *                   type is known by the return value, will need to
 *                   be free()d
 * @return -V9FS_E{error} on errors, V9FS_TYP_{type} on success
 */
int v9fs_unmarshal_msg(uint8_t *net_bytes, uint16_t *out_tag, void **out_body);
"""

    ret += """
/**
 * @param uint16_t in_msgid : the message-ID tag
 * @param struct v9fs_msg_{type} in_msg : the message to encode
 * @param uint8_t *out_buf : the buffer to encode to
 * @return uint32_t : the encoded length
 */
#define v9fs_marshal_msg(in_msgid, in_msg, out_buf) _Generic((in_msg)"""
    for msg in msgs:
        ret += f", \\\n\t\t{c_typename(msg)}: _v9fs_marshal_{shortname(msg)}(in_msgid, in_msg, out_buf)"
    ret += "\\\n\t)(in_msg)\n"
    for msg in msgs:
        ret += f"uint32_t _v9fs_marshal_{shortname(msg)}(uint16_t in_msgid, {c_typename(msg)} in_msg, uint8_t *out_buf);\n"

    ret += "\n"
    ret += "#endif /* _NET9P_DEFS_H_ */\n"
    return ret


def gen_c(structs: list[Struct], msgs: list[Message]) -> str:
    ret = """
/* Generated by ./net9p_defs.gen.  DO NOT EDIT!  */

#include <stdbool.h>
#include <stdint.h>
#include <stdlib.h> /* for malloc() */

#include "net9p_defs.h"
"""

    # basic utilities ##########################################################
    ret += """
/* basic utilities ************************************************************/

#define UNUSED __attribute__ ((unused))

static inline uint8_t decode_u8le(uint8_t *in) {
	return in[0];
}
static inline uint16_t decode_u16le(uint8_t *in) {
	return (((uint16_t)(in[0])) <<	0)
		 | (((uint16_t)(in[1])) <<	8)
		 ;
}
static inline uint32_t decode_u32le(uint8_t *in) {
	return (((uint32_t)(in[0])) <<	0)
		 | (((uint32_t)(in[1])) <<	8)
		 | (((uint32_t)(in[2])) << 16)
		 | (((uint32_t)(in[3])) << 24)
		 ;
}
static inline uint64_t decode_u64le(uint8_t *in) {
	return (((uint64_t)(in[0])) <<	0)
		 | (((uint64_t)(in[1])) <<	8)
		 | (((uint64_t)(in[2])) << 16)
		 | (((uint64_t)(in[3])) << 24)
		 | (((uint64_t)(in[4])) << 32)
		 | (((uint64_t)(in[5])) << 40)
		 | (((uint64_t)(in[6])) << 48)
		 | (((uint64_t)(in[7])) << 56)
		 ;
}

static inline void encode_u8le(uint8_t in, uint8_t *out) {
	out[0] = in;
}
static inline void encode_u16le(uint16_t in, uint8_t *out) {
	out[0] = (uint8_t)((in >>  0) & 0xFF);
	out[1] = (uint8_t)((in >>  8) & 0xFF);
}
static inline void encode_u32le(uint32_t in, uint8_t *out) {
	out[0] = (uint8_t)((in >>  0) & 0xFF);
	out[1] = (uint8_t)((in >>  8) & 0xFF);
	out[2] = (uint8_t)((in >> 16) & 0xFF);
	out[3] = (uint8_t)((in >> 24) & 0xFF);
}
static inline void encode_u64le(uint64_t in, uint8_t *out) {
	out[0] = (uint8_t)((in >>  0) & 0xFF);
	out[1] = (uint8_t)((in >>  8) & 0xFF);
	out[2] = (uint8_t)((in >> 16) & 0xFF);
	out[3] = (uint8_t)((in >> 24) & 0xFF);
	out[4] = (uint8_t)((in >> 32) & 0xFF);
	out[5] = (uint8_t)((in >> 40) & 0xFF);
	out[6] = (uint8_t)((in >> 48) & 0xFF);
	out[7] = (uint8_t)((in >> 56) & 0xFF);
}
"""

    # checksize_* ##############################################################
    ret += """
/* checksize_* ****************************************************************/

typedef bool (*_checksize_fn_t)(uint32_t net_len, uint8_t *net_bytes, uint32_t *mut_net_offset, size_t *mut_host_extra);
static inline bool _checksize_list(size_t cnt, _checksize_fn_t fn, size_t host_size,
                                   uint32_t net_len, uint8_t *net_bytes, uint32_t *mut_net_offset, size_t *mut_host_extra) {
	for (size_t i = 0; i < cnt; i++)
		if (__builtin_add_overflow(*mut_host_extra, host_size, mut_host_extra)
		 || fn(net_len, net_bytes, mut_net_offset, mut_host_extra))
			return true;
	return false;
}
static inline bool checksize_1(uint32_t net_len, uint8_t *net_bytes UNUSED, uint32_t *mut_net_offset, size_t *mut_host_extra UNUSED) {
	return __builtin_add_overflow(*mut_net_offset, 1, mut_net_offset) || net_len < *mut_net_offset;
}
static inline bool checksize_2(uint32_t net_len, uint8_t *net_bytes UNUSED, uint32_t *mut_net_offset, size_t *mut_host_extra UNUSED) {
	return __builtin_add_overflow(*mut_net_offset, 2, mut_net_offset) || net_len < *mut_net_offset;
}
static inline bool checksize_4(uint32_t net_len, uint8_t *net_bytes UNUSED, uint32_t *mut_net_offset, size_t *mut_host_extra UNUSED) {
	return __builtin_add_overflow(*mut_net_offset, 4, mut_net_offset) || net_len < *mut_net_offset;
}
static inline bool checksize_8(uint32_t net_len, uint8_t *net_bytes UNUSED, uint32_t *mut_net_offset, size_t *mut_host_extra UNUSED) {
	return __builtin_add_overflow(*mut_net_offset, 8, mut_net_offset) || net_len < *mut_net_offset;
}
"""
    for struct in structs + msgs:
        argattr = ' UNUSED' if len(struct.members) == 0 else ''
        ret += f"static inline bool checksize_{shortname(struct)}(uint32_t net_len{argattr}, uint8_t *net_bytes{argattr}, uint32_t *mut_net_offset{argattr}, size_t *mut_host_extra{argattr}) {{\n"
        if len(struct.members) == 0:
            ret += "\treturn false;\n"
            ret += "}\n"
            continue
        prefix0 = "\treturn "
        prefix1 = "\t    || "
        prefix2 = "\t       "
        prefix = prefix0
        prev_size: int|None = None
        for member in struct.members:
            if isinstance(member.typ, List):
                ret += f"\n{prefix }_checksize_list(decode_u{prev_size*8}le(&net_bytes[(*mut_net_offset)-{prev_size}]), checksize_{shortname(member.typ.typ)}, sizeof({c_typename(member.typ.typ)}),"
                ret += f"\n{prefix2}                net_len, net_bytes, mut_net_offset, mut_host_extra)"
            else:
                ret += f"\n{prefix}checksize_{shortname(member.typ)}(net_len, net_bytes, mut_net_offset, mut_host_extra)"
            prefix = prefix1
            prev_size = static_size(member.typ)
        if struct.name == 's':
            ret += f"\n{prefix}__builtin_add_overflow(*mut_host_extra, 1, mut_host_extra)"
        ret += ";\n}\n"

    # unmarshal_* ##############################################################
    ret += """
/* unmarshal_* ****************************************************************/
/* checksize_XXX() should be called before unmarshal_XXX().  */

static inline void unmarshal_1(uint8_t *net_bytes, uint32_t *mut_net_offset, void **mut_host_extra UNUSED, uint8_t *out) {
	*out = decode_u8le(&net_bytes[*mut_net_offset]);
	*mut_net_offset += 1;
}
static inline void unmarshal_2(uint8_t *net_bytes, uint32_t *mut_net_offset, void **mut_host_extra UNUSED, uint16_t *out) {
	*out = decode_u16le(&net_bytes[*mut_net_offset]);
	*mut_net_offset += 2;
}
static inline void unmarshal_4(uint8_t *net_bytes, uint32_t *mut_net_offset, void **mut_host_extra UNUSED, uint32_t *out) {
	*out = decode_u32le(&net_bytes[*mut_net_offset]);
	*mut_net_offset += 4;
}
static inline void unmarshal_8(uint8_t *net_bytes, uint32_t *mut_net_offset, void **mut_host_extra UNUSED, uint64_t *out) {
	*out = decode_u64le(&net_bytes[*mut_net_offset]);
	*mut_net_offset += 8;
}
"""
    for struct in structs + msgs:
        argattr = ' UNUSED' if len(struct.members) == 0 else ''
        ret += f"static inline void unmarshal_{shortname(struct)}(uint8_t *net_bytes{argattr}, uint32_t *mut_net_offset{argattr}, void **mut_host_extra{argattr}, {c_typename(struct)} *out{argattr}) {{"
        if len(struct.members) == 0:
            ret += "}\n"
            continue
        ret += "\n"
        for member in struct.members:
            if isinstance(member.typ, List):
                ret += f"\tout->{member.name} = *mut_host_extra;\n"
                ret += f"\t*mut_host_extra += sizeof(out->{member.name}) * out->{member.typ.cnt};\n"
                ret += f"\tfor (typeof(out->{member.typ.cnt}) i = 0; i < out->{member.typ.cnt}; i++)\n"
                ret += f"\t\tunmarshal_{shortname(member.typ.typ)}(net_bytes, mut_net_offset, mut_host_extra, &(out->{member.name}[i]));\n"
            else:
                ret += f"\tunmarshal_{shortname(member.typ)}(net_bytes, mut_net_offset, mut_host_extra, &(out->{member.name}));\n"
        ret += "}\n"

    # v9fs_unmarshal_msg #######################################################
    ret += """
/* v9fs_unmarshal_msg *********************************************************/

int v9fs_unmarshal_msg(uint8_t *net_bytes, uint16_t *out_tag, void **out_body) {
	uint32_t net_len = decode_u32le(net_bytes);
	if (net_len < 7)
		return -V9FS_EWRONGSIZE;
	uint8_t typ = net_bytes[4];
	*out_tag = decode_u16le(&net_bytes[5]);

	uint32_t net_offset = 7;
	size_t host_size;
	void *host_extra;
	switch (typ) {
"""
    for msg in msgs:
        ret += f"\tcase V9FS_TYP_{msg.name}:\n"
        ret += f"\t\thost_size = sizeof({c_typename(msg)});\n"
        ret += f"\t\tif (checksize_{shortname(msg)}(net_len, net_bytes, &net_offset, &host_size))\n"
        ret += "\t\t\treturn -V9FS_EWRONGSIZE;\n"
        ret += "\n"
        ret += "\t\t*out_body = malloc(host_size);"
        ret += "\n"
        ret += "\t\tnet_offset = 7;\n"
        ret += f"\t\thost_extra = *out_body + sizeof({c_typename(msg)});\n"
        ret += f"\t\tunmarshal_{shortname(msg)}(net_bytes, &net_offset, &host_extra, *out_body);\n"
        ret += "\n"
        ret += "\t\tbreak;\n"
    ret += """
	default:
		return -V9FS_EBADTYPE;
	}
	return typ;
}
"""

    # marshal_* ################################################################
    ret += """
/* marshal_* ******************************************************************/

static inline void marshal_1(uint8_t val, uint8_t *out_net_bytes, uint32_t *mut_net_offset) {
    out_net_bytes[*mut_net_offset] = val;
	*mut_net_offset += 1;
}
static inline void marshal_2(uint16_t val, uint8_t *out_net_bytes, uint32_t *mut_net_offset) {
	encode_u16le(val, &out_net_bytes[*mut_net_offset]);
	*mut_net_offset += 2;
}
static inline void marshal_4(uint32_t val, uint8_t *out_net_bytes, uint32_t *mut_net_offset) {
	encode_u32le(val, &out_net_bytes[*mut_net_offset]);
	*mut_net_offset += 4;
}
static inline void marshal_8(uint64_t val, uint8_t *out_net_bytes, uint32_t *mut_net_offset) {
	encode_u64le(val, &out_net_bytes[*mut_net_offset]);
	*mut_net_offset += 8;
}
"""
    for struct in structs + msgs:
        argattr = ' UNUSED' if len(struct.members) == 0 else ''
        ret += f"static inline void marshal_{shortname(struct)}({c_typename(struct)} val{argattr}, uint8_t *out_net_bytes{argattr}, uint32_t *mut_net_offset{argattr}) {{"
        if len(struct.members) == 0:
            ret += "}\n"
            continue
        ret += "\n"
        for member in struct.members:
            if isinstance(member.typ, List):
                ret += f"\tfor (typeof(val.{member.typ.cnt}) i = 0; i < val.{member.typ.cnt}; i++)\n"
                ret += f"\t\tmarshal_{shortname(member.typ.typ)}(val.{member.name}[i], out_net_bytes, mut_net_offset);\n"
            else:
                ret += f"\tmarshal_{shortname(member.typ)}(val.{member.name}, out_net_bytes, mut_net_offset);\n"
        ret += "}\n"

    # _v9fs_marshal_msg_* ######################################################
    ret += """
/* _v9fs_marshal_msg_* ********************************************************/

"""
    for msg in msgs:
        ret += f"uint32_t _v9fs_marshal_{shortname(msg)}(uint16_t in_msgid, {c_typename(msg)} in_msg, uint8_t *out_buf) {{\n"
        ret += "\tuint32_t offset = 4;\n"
        ret += f"\tmarshal_1(V9FS_TYP_{msg.name}, out_buf, &offset);\n"
        ret += "\tmarshal_2(in_msgid, out_buf, &offset);\n"
        ret += f"\tmarshal_{shortname(msg)}(in_msg, out_buf, &offset);\n"
        ret += "\tencode_u32le(offset, out_buf);\n"
        ret += "\treturn offset;\n"
        ret += "}\n"
    ret += "\n"

    ############################################################################
    return ret


################################################################################

if __name__ == "__main__":
    structs, msgs = parse_file("net9p_defs.txt")
    with open('net9p_defs.h', 'w') as fh:
        fh.write(gen_h(structs, msgs))
    with open('net9p_defs.c', 'w') as fh:
        fh.write(gen_c(structs, msgs))