// SPDX-FileCopyrightText: 2024 Himbeer <himbeer@disroot.org>
//
// SPDX-License-Identifier: AGPL-3.0-or-later

const std = @import("std");

const slice = @import("slice.zig");

pub var default: Tree = undefined;

pub const ParseError = error{
    BadMagic,
    BadToken,
    ExpectedBeginNodeToken,
    ExpectedPropToken,
    DuplicateProperty,
};

pub const NodeError = error{
    NoParent,
    BadCellSize,
    PropertyNotFound,
};

pub const MemReservation = struct {
    const sentinel = .{
        .addr = 0,
        .size = 0,
    };

    addr: u64,
    size: u64,
};

pub const Reg = struct {
    start: usize,
    len: usize,
};

pub const Tree = struct {
    nodes: std.ArrayList(Node),
};

pub const Node = struct {
    name: []const u8,
    props: std.StringHashMap([]const u8),
    parent: ?*Node,
    subnodes: std.ArrayList(Node),

    pub fn unitAddr(self: Node) !?usize {
        var nameSplit = std.mem.splitScalar(u8, self.name, '@');
        _ = nameSplit.next();
        const unitAddrStr = nameSplit.next() orelse return null;

        return try std.fmt.parseInt(usize, unitAddrStr, 16);
    }

    pub fn preferredDriver(self: Node, drivers: []const []const u8) ?[]const u8 {
        const compatible_prop = self.props.get("compatible") orelse return null;

        var compatibles = std.mem.tokenizeScalar(u8, compatible_prop, '\x00');
        while (compatibles.next()) |compatible| {
            for (drivers) |driver| {
                if (std.mem.eql(u8, driver, compatible)) return compatible;
            }
        }

        return null;
    }

    pub fn isCompatible(self: Node, with: []const u8) bool {
        return self.preferredDriver(&[_][]const u8{with}) != null;
    }

    pub fn reg(self: Node, allocator: std.mem.Allocator) !std.ArrayList(Reg) {
        if (self.parent == null) return NodeError.NoParent;

        const address_cells_bytes = self.parent.?.props.get("#address-cells");
        const size_cells_bytes = self.parent.?.props.get("#size-cells");

        const address_cells = if (address_cells_bytes) |bytes| std.mem.readInt(u32, bytes[0..4], .Big) else 2;
        const size_cells = if (size_cells_bytes) |bytes| std.mem.readInt(u32, bytes[0..4], .Big) else 1;

        if (address_cells == 0 or size_cells == 0) return NodeError.BadCellSize;

        const reg_prop = self.props.get("reg") orelse return NodeError.PropertyNotFound;
        const reg_elem_len = 4 * address_cells + 4 * size_cells;

        const n = reg_prop.len / reg_elem_len;

        var regs = std.ArrayList(Reg).init(allocator);

        for (0..n) |i| {
            const start_offset = i * reg_elem_len;
            const len_offset = start_offset + 4 * address_cells;

            try regs.append(.{
                .start = std.mem.readVarInt(usize, reg_prop[start_offset .. start_offset + 4 * address_cells], .Big),
                .len = std.mem.readVarInt(usize, reg_prop[len_offset .. len_offset + 4 * size_cells], .Big),
            });
        }

        return regs;
    }
};

pub const Property = struct {
    name: []const u8,
    value: []const u8,
};

pub const PropertyDesc = struct {
    len: u32,
    name_offset: u32,

    pub fn fromPtr(ptr: *PropertyDesc) PropertyDesc {
        return .{
            .len = std.mem.bigToNative(u32, ptr.len),
            .name_offset = std.mem.bigToNative(u32, ptr.name_offset),
        };
    }
};

fn ParseResult(comptime T: type) type {
    return struct {
        addr: usize,
        value: T,
    };
}

pub const RawHeader = struct {
    magic: u32,
    total_size: u32,
    offset_dt_struct: u32,
    offset_dt_strings: u32,
    offset_mem_reserve_map: u32,
    version: u32,
    last_compat_version: u32,

    // v2
    boot_cpuid_phys: u32,

    // v3
    size_dt_strings: u32,

    // v17
    size_dt_struct: u32,
};

pub const Header = struct {
    const magic = 0xd00dfeed;

    const token_begin_node = 0x00000001;
    const token_end_node = 0x00000002;
    const token_prop = 0x00000003;
    const token_nop = 0x00000004;
    const token_end = 0x00000009;

    raw_hdr: *RawHeader,

    magic: u32,
    total_size: u32,
    offset_dt_struct: u32,
    offset_dt_strings: u32,
    offset_mem_reserve_map: u32,
    version: u32,
    last_compat_version: u32,

    // v2
    boot_cpuid_phys: u32,

    // v3
    size_dt_strings: u32,

    // v17
    size_dt_struct: u32,

    pub fn parse(raw_hdr: *RawHeader) !Header {
        const hdr = .{
            .raw_hdr = raw_hdr,

            .magic = std.mem.bigToNative(u32, raw_hdr.magic),
            .total_size = std.mem.bigToNative(u32, raw_hdr.total_size),
            .offset_dt_struct = std.mem.bigToNative(u32, raw_hdr.offset_dt_struct),
            .offset_dt_strings = std.mem.bigToNative(u32, raw_hdr.offset_dt_strings),
            .offset_mem_reserve_map = std.mem.bigToNative(u32, raw_hdr.offset_mem_reserve_map),
            .version = std.mem.bigToNative(u32, raw_hdr.version),
            .last_compat_version = std.mem.bigToNative(u32, raw_hdr.last_compat_version),

            // v2
            .boot_cpuid_phys = std.mem.bigToNative(u32, raw_hdr.boot_cpuid_phys),

            // v3
            .size_dt_strings = std.mem.bigToNative(u32, raw_hdr.size_dt_strings),

            // v17
            .size_dt_struct = std.mem.bigToNative(u32, raw_hdr.size_dt_struct),
        };

        if (hdr.magic != magic) {
            return ParseError.BadMagic;
        }

        return hdr;
    }

    pub fn memoryReservations(self: *const Header) []MemReservation {
        const mem_reserve_map_addr = @intFromPtr(self.raw_hdr) + @as(usize, self.offset_mem_reserve_map);
        const mem_reservations: [*:MemReservation.sentinel]MemReservation = @ptrFromInt(mem_reserve_map_addr);

        var i: usize = 0;
        while (!std.meta.eql(mem_reservations[i], MemReservation.sentinel)) {
            i += 1;
        }

        return @ptrCast(mem_reservations[0..i]);
    }

    pub fn parseTree(self: *const Header, allocator: std.mem.Allocator) !Tree {
        var nodes = std.ArrayList(Node).init(allocator);

        var dt_struct_addr = @intFromPtr(self.raw_hdr) + @as(usize, self.offset_dt_struct);

        const dt_strings_addr = @intFromPtr(self.raw_hdr) + @as(usize, self.offset_dt_strings);
        const dt_strings: [*:0]const u8 = @ptrFromInt(dt_strings_addr);

        while (std.mem.bigToNative(u32, @as(*u32, @ptrFromInt(dt_struct_addr)).*) != token_end) {
            const parsed = try parseNode(allocator, dt_struct_addr, dt_strings);
            dt_struct_addr = parsed.addr;
            try nodes.insert(0, parsed.value);
        }

        var tree = Tree{ .nodes = nodes };
        initParentPtrs(&tree);
        return tree;
    }
};

pub fn findNode(nodes: []const Node, name: []const u8) slice.Filter(Node, []const u8) {
    return slice.Filter(Node, []const u8).new(nodes, nodeNameFilter, name);
}

pub fn findNodeExact(nodes: []const Node, name: []const u8) ?struct { usize, Node } {
    for (nodes, 0..) |node, i| {
        if (std.mem.eql(u8, node.name, name)) {
            return .{ i, node };
        }
    }

    return null;
}

pub fn findPath(dt: *const Tree, path: []const u8) ?Node {
    if (dt.nodes.items.len < 1) return null;

    const trimmed_path = std.mem.trim(u8, path, "/");

    var node = dt.nodes.items[0];

    var segments = std.mem.tokenizeScalar(u8, trimmed_path, '/');
    while (segments.next()) |segment| {
        var nodes = findNode(node.subnodes.items, segment);
        if (nodes.next()) |result| {
            node = result;
        } else return null;
    }

    return node;
}

pub fn findPathExact(dt: *const Tree, path: []const u8) ?Node {
    if (dt.nodes.items.len < 1) return null;

    const trimmed_path = std.mem.trim(u8, path, "/");

    var node = dt.nodes.items[0];

    var segments = std.mem.tokenizeScalar(u8, trimmed_path, '/');
    while (segments.next()) |segment| {
        if (findNodeExact(node.subnodes.items, segment)) |result| {
            node = result[1];
        } else return null;
    }

    return node;
}

fn nodeNameFilter(node: Node, name: []const u8) bool {
    var it = std.mem.splitScalar(u8, node.name, '@');
    const trueName = it.first();
    return std.mem.eql(u8, trueName, name);
}

fn parseNode(allocator: std.mem.Allocator, dt_struct_addr: usize, dt_strings: [*:0]const u8) !ParseResult(Node) {
    var props = std.StringHashMap([]const u8).init(allocator);
    var subnodes = std.ArrayList(Node).init(allocator);

    var addr = dt_struct_addr;

    // Skip Nop tokens
    while (std.mem.bigToNative(u32, @as(*u32, @ptrFromInt(addr)).*) == Header.token_nop) {
        addr += @sizeOf(u32);
    }

    if (std.mem.bigToNative(u32, @as(*u32, @ptrFromInt(addr)).*) != Header.token_begin_node) {
        return ParseError.ExpectedBeginNodeToken;
    }

    addr += @sizeOf(u32);

    const name: [*:0]u8 = @ptrFromInt(addr);

    addr += std.mem.len(name) + 1;

    // Skip zeroed alignment padding
    addr += (4 - ((std.mem.len(name) + 1) % 4)) % 4;

    while (std.mem.bigToNative(u32, @as(*u32, @ptrFromInt(addr)).*) != Header.token_end_node) {
        switch (std.mem.bigToNative(u32, @as(*u32, @ptrFromInt(addr)).*)) {
            Header.token_prop => {
                const parsed = try parseProperty(addr, dt_strings);
                addr = parsed.addr;

                const result = try props.getOrPut(parsed.value.name);
                if (result.found_existing) {
                    return ParseError.DuplicateProperty;
                }

                result.value_ptr.* = parsed.value.value;
            },
            // Skip Nop tokens
            Header.token_nop => addr += @sizeOf(u32),
            Header.token_begin_node => {
                const parsed = try parseNode(allocator, addr, dt_strings);
                addr = parsed.addr;
                try subnodes.insert(0, parsed.value);
            },
            else => return ParseError.BadToken,
        }
    }

    addr += @sizeOf(u32);

    return .{
        .addr = addr,
        .value = .{
            .name = std.mem.span(name),
            .props = props,
            .parent = null,
            .subnodes = subnodes,
        },
    };
}

fn parseProperty(dt_struct_addr: usize, dt_strings: [*:0]const u8) !ParseResult(Property) {
    var addr = dt_struct_addr;

    // Skip Nop tokens
    while (std.mem.bigToNative(u32, @as(*u32, @ptrFromInt(addr)).*) == Header.token_nop) {
        addr += @sizeOf(u32);
    }

    if (std.mem.bigToNative(u32, @as(*u32, @ptrFromInt(addr)).*) != Header.token_prop) {
        return ParseError.ExpectedPropToken;
    }

    addr += @sizeOf(u32);

    const desc = PropertyDesc.fromPtr(@ptrFromInt(addr));

    addr += @sizeOf(PropertyDesc);

    const name = dt_strings + desc.name_offset;
    const value_bytes: [*]u8 = @ptrFromInt(addr);
    const value: []const u8 = @ptrCast(value_bytes[0..desc.len]);

    addr += desc.len;

    // Skip zeroed alignment padding
    addr += (4 - (desc.len % 4)) % 4;

    return .{
        .addr = addr,
        .value = .{
            .name = std.mem.span(name),
            .value = value,
        },
    };
}

fn initParentPtrs(dt: *Tree) void {
    for (dt.nodes.items, 0..) |_, i| {
        for (dt.nodes.items[i].subnodes.items, 0..) |_, j| {
            initParentPtrsRecursive(&dt.nodes.items[i].subnodes.items[j], &dt.nodes.items[i]);
        }
    }
}

fn initParentPtrsRecursive(node: *Node, parent: *Node) void {
    node.parent = parent;
    for (node.subnodes.items, 0..) |_, i| {
        initParentPtrsRecursive(&node.subnodes.items[i], node);
    }
}