const lex = @import("lex.zig");
const fail = @import("fail.zig");
const std = @import("std");

pub const Scalar = enum {
    U8,
    U16,
    U32,
    U64,

    pub fn width(self: Scalar) usize {
        return switch (self) {
            Scalar.U8 => 1,
            Scalar.U16 => 2,
            Scalar.U32 => 4,
            Scalar.U64 => 8,
        };
    }

    pub fn pba_name(self: Scalar) []const u8 {
        return switch (self) {
            Scalar.U8 => "u8",
            Scalar.U16 => "u16",
            Scalar.U32 => "u32",
            Scalar.U64 => "u64",
        };
    }

    pub fn c_type(self: Scalar) []const u8 {
        return switch (self) {
            Scalar.U8 => "uint8_t",
            Scalar.U16 => "uint16_t",
            Scalar.U32 => "uint32_t",
            Scalar.U64 => "uint64_t",
        };
    }
};

pub const ArrayLength = union(enum) {
    fixed: usize,
    ref: []const u8,
    open: void,
};

pub const FieldType = union(enum) {
    single: Scalar,
    array: struct {
        scalar_type: Scalar,
        length: ArrayLength,
    },
};

pub const Field = struct {
    name: []const u8,
    field_type: FieldType,
    position: lex.Position,

    pub fn parse(token_stream: []lex.Token) struct { Field, usize } {
        if (token_stream.len < 2) {
            fail.fail(.{}, "incomplete field declaration");
        }
        var res: Field = .{ .field_type = undefined, .name = undefined, .position = token_stream[0].pos };
        var scalar_type: Scalar = undefined;
        switch (token_stream[0].tok) {
            .identifier => |type_ident| {
                if (std.mem.eql(u8, type_ident, "u8")) {
                    scalar_type = Scalar.U8;
                } else if (std.mem.eql(u8, type_ident, "u16")) {
                    scalar_type = Scalar.U16;
                } else if (std.mem.eql(u8, type_ident, "u32")) {
                    scalar_type = Scalar.U32;
                } else if (std.mem.eql(u8, type_ident, "u64")) {
                    scalar_type = Scalar.U64;
                } else {
                    fail.fail(token_stream[0].pos, "expected a scalar type (u8, u16, u32, or u64)");
                }
            },
            else => fail.fail(token_stream[0].pos, "expected a scalar type (u8, u16, u32, or u64)"),
        }
        switch (token_stream[1].tok) {
            .identifier => |name| {
                res.name = name;
                res.field_type = .{ .single = scalar_type };
                return .{ res, 2 };
            },
            .lbracket => {},
            else => fail.fail(token_stream[1].pos, "invalid field declaration"),
        }
        if (token_stream.len < 4) {
            fail.fail(token_stream[0].pos, "incomplete field declaration (expecting an array declaration)");
        }
        var arr_len: ArrayLength = undefined;
        var cursor: usize = 2;
        switch (token_stream[cursor].tok) {
            .rbracket => {
                arr_len = .open;
                cursor += 1;
            },
            .scalar => |len| {
                if (token_stream.len < 5 or !std.meta.eql(token_stream[3].tok, .rbracket)) {
                    fail.fail(token_stream[cursor].pos, "invalid field declaration");
                }
                arr_len = .{ .fixed = len };
                cursor += 2;
            },
            .identifier => |name| {
                if (token_stream.len < 5 or !std.meta.eql(token_stream[3].tok, .rbracket)) {
                    fail.fail(token_stream[cursor].pos, "invalid field declaration");
                }
                arr_len = .{ .ref = name };
                cursor += 2;
            },
            else => fail.fail(token_stream[2].pos, "invalid field declaration"),
        }
        res.field_type = .{ .array = .{ .length = arr_len, .scalar_type = scalar_type } };
        switch (token_stream[cursor].tok) {
            .identifier => |name| {
                res.name = name;
                return .{ res, cursor + 1 };
            },
            else => fail.fail(token_stream[cursor].pos, "invalid field declaration"),
        }
        unreachable;
    }
};

pub const Builtin = enum {
    MSG_PRECEDING,
    FRAG_PRECEDING,
    MSG_REMAINING,
    FRAG_REMAINING,
    FRAGMENT_ID,
};

pub const Value = union(enum) {
    literal: usize,
    field_ref: []const u8,
    builtin: Builtin,

    pub fn parse(token_stream: []lex.Token) struct { Value, usize } {
        if (token_stream.len < 1) {
            fail.fail(.{}, "expected a value, got EOF");
        }
        switch (token_stream[0].tok) {
            .scalar => |literal| {
                return .{ .{ .literal = literal }, 1 };
            },
            .identifier => |name| {
                if (std.mem.eql(u8, name, "msg_preceding")) {
                    return .{ .{ .builtin = Builtin.MSG_PRECEDING }, 1 };
                } else if (std.mem.eql(u8, name, "msg_remaining")) {
                    return .{ .{ .builtin = Builtin.MSG_REMAINING }, 1 };
                } else if (std.mem.eql(u8, name, "frag_preceding")) {
                    return .{ .{ .builtin = Builtin.FRAG_PRECEDING }, 1 };
                } else if (std.mem.eql(u8, name, "frag_remaining")) {
                    return .{ .{ .builtin = Builtin.FRAG_REMAINING }, 1 };
                } else if (std.mem.eql(u8, name, "fragment_id")) {
                    return .{ .{ .builtin = Builtin.FRAGMENT_ID }, 1 };
                }
                return .{ .{ .field_ref = name }, 1 };
            },
            else => fail.fail(token_stream[0].pos, "expected a value"),
        }
        unreachable;
    }
};

pub const RegionAlgorithm = union(enum) {
    chacha20: struct {
        nonce: Value,
        key: Value,
    },
    poly1305: struct {
        nonce: Value,
        key: Value,
        authenticates: Value,
    },
    clear: void,

    pub fn tag_len(self: RegionAlgorithm) usize {
        return switch (self) {
            .poly1305 => 16,
            else => 0,
        };
    }

    pub fn parse(token_stream: []lex.Token) struct { RegionAlgorithm, usize } {
        if (token_stream.len < 1) {
            fail.fail(.{}, "incomplete region declaration");
        }
        var alg_name: []const u8 = undefined;
        switch (token_stream[0].tok) {
            .identifier => |name| {
                alg_name = name;
            },
            else => fail.fail(token_stream[0].pos, "not an encryption/MAC algorithm or 'clear;"),
        }
        if (std.mem.eql(u8, alg_name, "clear")) {
            return .{ .clear, 1 };
        } else if (std.mem.eql(u8, alg_name, "chacha20")) {
            if (token_stream.len < 6) {
                fail.fail(token_stream[0].pos, "incomplete ChaCha20 region declaration");
            }
            switch (token_stream[1].tok) {
                .lparen => {},
                else => fail.fail(token_stream[1].pos, "expected parameter list"),
            }
            var cursor: usize = 2;
            const nonce, var advance = Value.parse(token_stream[cursor..]);
            cursor += advance;
            const key, advance = Value.parse(token_stream[cursor..]);
            cursor += advance;
            switch (token_stream[cursor].tok) {
                .rparen => {},
                else => fail.fail(token_stream[cursor].pos, "expected closing parenthesis"),
            }
            return .{ .{ .chacha20 = .{ .nonce = nonce, .key = key } }, cursor + 1 };
        } else if (std.mem.eql(u8, alg_name, "poly1305")) {
            if (token_stream.len < 6) {
                fail.fail(token_stream[0].pos, "incomplete Poly1305 region declaration");
            }
            switch (token_stream[1].tok) {
                .lparen => {},
                else => fail.fail(token_stream[1].pos, "expected parameter list"),
            }
            var cursor: usize = 2;
            const nonce, var advance = Value.parse(token_stream[cursor..]);
            cursor += advance;
            const key, advance = Value.parse(token_stream[cursor..]);
            cursor += advance;
            const authenticates, advance = Value.parse(token_stream[cursor..]);
            cursor += advance;
            switch (token_stream[cursor].tok) {
                .rparen => {},
                else => fail.fail(token_stream[cursor].pos, "expected closing parenthesis"),
            }
            return .{ .{ .poly1305 = .{ .nonce = nonce, .key = key, .authenticates = authenticates } }, cursor + 1 };
        } else {
            fail.fail(token_stream[0].pos, "unknown AEAD algorithm");
        }
        unreachable;
    }
};

pub const DigestAlgorithm = enum {
    SHA256,
    SHA512,

    pub fn digest_len(self: DigestAlgorithm) usize {
        return switch (self) {
            DigestAlgorithm.SHA256 => 32,
            DigestAlgorithm.SHA512 => 64,
        };
    }

    pub fn parse(token_stream: []lex.Token) struct { DigestAlgorithm, usize } {
        if (token_stream.len < 1) {
            fail.fail(.{}, "expected a digest algorithm");
        }
        switch (token_stream[0].tok) {
            .identifier => |name| {
                if (std.mem.eql(u8, name, "sha256")) {
                    return .{ DigestAlgorithm.SHA256, 1 };
                } else if (std.mem.eql(u8, name, "sha512")) {
                    return .{ DigestAlgorithm.SHA512, 1 };
                }
                fail.fail(token_stream[0].pos, "unknown digest algorithm");
            },
            else => fail.fail(token_stream[0].pos, "not a digest algorithm"),
        }
        unreachable;
    }
};

pub const Appears = enum {
    ONCE_PER_FRAGMENT,
    ONCE_PER_MESSAGE,
};

pub const Block = struct {
    algorithm: RegionAlgorithm,
    fields: ?[]Field,
    appears: Appears,
    pos: lex.Position,

    pub fn parse(token_stream: []lex.Token, allocator: std.mem.Allocator) struct { Block, usize } {
        var fields = std.ArrayList(Field).empty;
        if (token_stream.len < 2) {
            fail.fail(.{}, "incomplete or missing block declaration");
        }
        var cursor: usize = 0;
        var appears = Appears.ONCE_PER_MESSAGE;
        switch (token_stream[cursor].tok) {
            .bang => {
                appears = Appears.ONCE_PER_FRAGMENT;
                cursor += 1;
            },
            else => {},
        }
        const reg_alg, var advance = RegionAlgorithm.parse(token_stream[cursor..]);
        cursor += advance;
        if (cursor >= token_stream.len) {
            fail.fail(.{}, "expected field list, got EOF");
        }
        switch (token_stream[cursor].tok) {
            .lbrace => {},
            else => {
                return .{ .{ .algorithm = reg_alg, .fields = null, .appears = appears, .pos = token_stream[0].pos }, cursor };
            },
        }
        cursor += 1;
        while (cursor < token_stream.len and !std.meta.eql(token_stream[cursor].tok, .rbrace)) {
            const next_field, advance = Field.parse(token_stream[cursor..]);
            cursor += advance;
            fields.append(allocator, next_field) catch fail.fail(token_stream[cursor].pos, "out of memory");
        }
        if (cursor >= token_stream.len) {
            fail.fail(.{}, "expected closing brace, got EOF");
        }
        switch (token_stream[cursor].tok) {
            .rbrace => {},
            else => fail.fail(token_stream[cursor].pos, "expected closing brace"),
        }
        cursor += 1;
        return .{ .{ .algorithm = reg_alg, .fields = fields.items, .appears = appears, .pos = token_stream[0].pos }, cursor };
    }
};

pub const Message = struct {
    name: []const u8,
    params: []Field,
    blocks: []Block,
    pos: lex.Position,
    digest: ?DigestAlgorithm,

    pub fn parse(token_stream: []lex.Token, allocator: std.mem.Allocator) struct { Message, usize } {
        var blocks = std.ArrayList(Block).empty;
        if (token_stream.len < 3) {
            fail.fail(.{}, "incomplete or missing message declaration");
        }
        var cursor: usize = 0;
        var name: []const u8 = undefined;
        switch (token_stream[cursor].tok) {
            .identifier => |ident| {
                name = ident;
                cursor += 1;
            },
            else => fail.fail(token_stream[cursor].pos, "expected message name"),
        }
        var params = std.ArrayList(Field).empty;
        switch (token_stream[cursor].tok) {
            .lparen => {
                cursor += 1;
                while (cursor < token_stream.len and !std.meta.eql(token_stream[cursor].tok, .rparen)) {
                    const next_param, const advance = Field.parse(token_stream[cursor..]);
                    cursor += advance;
                    params.append(allocator, next_param) catch fail.fail(.{}, "out of memory");
                }
                if (cursor >= token_stream.len) {
                    fail.fail(.{}, "expected closing parenthesis, got EOF");
                }
                cursor += 1;
            },
            .identifier => {},
            else => fail.fail(token_stream[cursor].pos, "expected message name"),
        }
        var digest: ?DigestAlgorithm = null;
        switch (token_stream[cursor].tok) {
            .identifier => {
                digest, const advance = DigestAlgorithm.parse(token_stream[cursor..]);
                cursor += advance;
            },
            .lbrace => {},
            else => fail.fail(token_stream[cursor].pos, "expected digest algorithm or block list"),
        }
        switch (token_stream[cursor].tok) {
            .lbrace => {
                cursor += 1;
            },
            else => fail.fail(token_stream[cursor].pos, "expected digest algorithm or block list"),
        }
        while (cursor < token_stream.len and !std.meta.eql(token_stream[cursor].tok, .rbrace)) {
            const next_block, const advance = Block.parse(token_stream[cursor..], allocator);
            cursor += advance;
            blocks.append(allocator, next_block) catch fail.fail(token_stream[cursor].pos, "out of memory");
        }
        if (cursor >= token_stream.len) {
            fail.fail(.{}, "expected closing brace, got EOF");
        }
        cursor += 1;
        return .{ .{ .name = name, .params = params.items, .blocks = blocks.items, .pos = token_stream[0].pos, .digest = digest }, cursor };
    }
};

pub const ProtocolInstance = struct {
    version: usize,
    header: []Field,
    messages: []Message,
    pos: lex.Position,

    pub fn parse(token_stream: []lex.Token, allocator: std.mem.Allocator) struct { ProtocolInstance, usize } {
        var res: ProtocolInstance = .{ .header = undefined, .version = undefined, .messages = undefined, .pos = undefined };
        var header_fields = std.ArrayList(Field).empty;
        if (token_stream.len < 7) {
            fail.fail(.{}, "incomplete version declaration");
        }
        res.pos = token_stream[0].pos;
        switch (token_stream[0].tok) {
            .identifier => |version_ident| {
                if (!std.mem.eql(u8, version_ident, "version")) {
                    fail.fail(token_stream[0].pos, "expected a version declaration");
                }
            },
            else => fail.fail(token_stream[0].pos, "expected a version declaration"),
        }
        switch (token_stream[1].tok) {
            .scalar => |version| {
                if (version <= 0) {
                    fail.fail(token_stream[1].pos, "version number must be positive");
                }
                res.version = version;
            },
            else => fail.fail(token_stream[1].pos, "version must be a scalar"),
        }
        switch (token_stream[2].tok) {
            .lbrace => {},
            else => fail.fail(token_stream[2].pos, "expected opening brace"),
        }
        switch (token_stream[3].tok) {
            .identifier => |header_ident| {
                if (!std.mem.eql(u8, header_ident, "header")) {
                    fail.fail(token_stream[3].pos, "expected a header declaration");
                }
            },
            else => fail.fail(token_stream[3].pos, "expected a header declaration"),
        }
        switch (token_stream[4].tok) {
            .lbrace => {},
            else => fail.fail(token_stream[4].pos, "expected opening brace"),
        }
        var current_idx: usize = 5;
        while (current_idx < token_stream.len and !std.meta.eql(token_stream[current_idx].tok, .rbrace)) {
            const next_header_field, const advance = Field.parse(token_stream[current_idx..]);
            header_fields.append(allocator, next_header_field) catch fail.fail(.{}, "out of memory");
            current_idx += advance;
        }
        if (current_idx >= token_stream.len) {
            fail.fail(.{}, "expected closing brace, got EOF");
        }
        current_idx += 1;
        res.header = header_fields.items;
        var messages = std.ArrayList(Message).empty;
        while (current_idx < token_stream.len and !std.meta.eql(token_stream[current_idx].tok, .rbrace)) {
            const next_message, const advance = Message.parse(token_stream[current_idx..], allocator);
            messages.append(allocator, next_message) catch fail.fail(.{}, "out of memory");
            current_idx += advance;
        }
        if (current_idx >= token_stream.len) {
            fail.fail(.{}, "expected closing brace, got EOF");
        }
        current_idx += 1;
        return .{ res, current_idx };
    }
};

pub const Protocol = struct {
    name: []const u8,
    mtu: usize,
    instances: []ProtocolInstance,

    pub fn parse(tokens: []lex.Token, allocator: std.mem.Allocator) Protocol {
        var res: Protocol = .{ .name = undefined, .mtu = undefined, .instances = undefined };
        var instances = std.ArrayList(ProtocolInstance).empty;
        if (tokens.len < 3) {
            fail.fail(.{}, "expected at least a protocol name and MTU declaration");
        }
        switch (tokens[0].tok) {
            .identifier => |protocol_name| res.name = protocol_name,
            else => fail.fail(tokens[0].pos, "expected the protocol name"),
        }
        switch (tokens[1].tok) {
            .identifier => |mtu_ident| {
                if (!std.mem.eql(u8, mtu_ident, "mtu")) {
                    fail.fail(tokens[1].pos, "expected an MTU declaration");
                }
            },
            else => fail.fail(tokens[1].pos, "expected an MTU declaration"),
        }
        switch (tokens[2].tok) {
            .scalar => |mtu| {
                if (mtu <= 0) {
                    fail.fail(tokens[2].pos, "MTU must be positive");
                }
                res.mtu = mtu;
            },
            else => fail.fail(tokens[2].pos, "MTU must be a scalar"),
        }
        var current_idx: usize = 3;
        while (current_idx < tokens.len and !std.meta.eql(tokens[current_idx].tok, .eof)) {
            const next_instance, const advance = ProtocolInstance.parse(tokens[current_idx..], allocator);
            current_idx += advance;
            instances.append(allocator, next_instance) catch fail.fail(.{}, "out of memory");
        }
        res.instances = instances.items;
        return res;
    }
};

const auto_header_len = 6;