const std = @import("std");
const Fdt = @This();

pub const MAGIC = 0xd00dfeed;

pub const Token = enum(u32) {
    begin_node = 0x01,
    end_node = 0x02,
    prop = 0x03,
    nop = 0x04,
    end = 0x09,
};

pub const RawHeader = extern struct {
    magic: u32,
    totalsize: u32,
    off_dt_struct: u32,
    off_dt_strings: u32,
    off_mem_rsvmap: u32,
    version: u32,
    last_comp_version: u32,
    boot_cpuid_phys: u32,
    size_dt_strings: u32,
    size_dt_struct: u32,
};

pub const ReserveEntry = extern struct {
    address: u64,
    size: u64,

    pub fn toNative(self: *const ReserveEntry) struct { address: u64, size: u64 } {
        return .{
            .address = std.mem.bigToNative(u64, self.address),
            .size = std.mem.bigToNative(u64, self.size),
        };
    }
};

pub const Property = struct {
    name: []const u8,
    data: []const u8,

    pub fn asU32(self: Property) ?u32 {
        if (self.data.len != 4) return null;
        return std.mem.bigToNative(u32, @as(*const u32, @alignCast(@ptrCast(self.data.ptr))).*);
    }

    pub fn asU64(self: Property) ?u64 {
        if (self.data.len != 8) return null;
        return std.mem.bigToNative(u64, @as(*const u64, @alignCast(@ptrCast(self.data.ptr))).*);
    }

    pub fn asString(self: Property) ?[]const u8 {
        if (self.data.len == 0) return null;
        for (self.data, 0..) |c, i| {
            if (c == 0) return self.data[0..i];
        }
        return self.data;
    }

    pub fn asStringList(self: Property) StringListIterator {
        return .{ .data = self.data };
    }

    pub fn asU32Array(self: Property) U32ArrayIterator {
        return .{ .data = self.data };
    }

    pub fn asU64Array(self: Property) U64ArrayIterator {
        return .{ .data = self.data };
    }

    pub fn isEmpty(self: Property) bool {
        return self.data.len == 0;
    }
};

pub const StringListIterator = struct {
    data: []const u8,
    pos: usize = 0,

    pub fn next(self: *StringListIterator) ?[]const u8 {
        if (self.pos >= self.data.len) return null;

        const start = self.pos;
        while (self.pos < self.data.len and self.data[self.pos] != 0) {
            self.pos += 1;
        }
        const end = self.pos;
        self.pos += 1;

        if (start == end) return null;
        return self.data[start..end];
    }
};

pub const U32ArrayIterator = struct {
    data: []const u8,
    pos: usize = 0,

    pub fn next(self: *U32ArrayIterator) ?u32 {
        if (self.pos + 4 > self.data.len) return null;
        const value = std.mem.bigToNative(u32, @as(*const u32, @alignCast(@ptrCast(self.data.ptr + self.pos))).*);
        self.pos += 4;
        return value;
    }

    pub fn count(self: U32ArrayIterator) usize {
        return self.data.len / 4;
    }
};

pub const U64ArrayIterator = struct {
    data: []const u8,
    pos: usize = 0,

    pub fn next(self: *U64ArrayIterator) ?u64 {
        if (self.pos + 8 > self.data.len) return null;
        const value = std.mem.bigToNative(u64, @as(*const u64, @alignCast(@ptrCast(self.data.ptr + self.pos))).*);
        self.pos += 8;
        return value;
    }

    pub fn count(self: U64ArrayIterator) usize {
        return self.data.len / 8;
    }
};

pub const Node = struct {
    name: []const u8,
    fdt: *const Fdt,
    struct_offset: usize,

    pub fn unitAddress(self: Node) ?[]const u8 {
        for (self.name, 0..) |c, i| {
            if (c == '@') return self.name[i + 1 ..];
        }
        return null;
    }

    pub fn baseName(self: Node) []const u8 {
        for (self.name, 0..) |c, i| {
            if (c == '@') return self.name[0..i];
        }
        return self.name;
    }

    pub fn properties(self: Node) PropertyIterator {
        return .{
            .fdt = self.fdt,
            .pos = self.struct_offset,
        };
    }

    pub fn getProperty(self: Node, name: []const u8) ?Property {
        var iter = self.properties();
        while (iter.next()) |prop| {
            if (std.mem.eql(u8, prop.name, name)) {
                return prop;
            }
        }
        return null;
    }

    pub fn children(self: Node) ChildIterator {
        // Skip to the end of properties to find children
        var pos = self.struct_offset;
        while (pos < self.fdt.struct_block.len) {
            const token = self.fdt.readToken(pos) orelse break;
            pos += 4;

            switch (token) {
                .prop => {
                    const len = self.fdt.readU32(pos) orelse break;
                    pos += 8; // skip len and nameoff
                    pos += alignUp(len, 4);
                },
                .nop => {},
                .begin_node, .end_node, .end => {
                    pos -= 4; // back up to re-read this token
                    break;
                },
            }
        }

        return .{
            .fdt = self.fdt,
            .pos = pos,
            .depth = 0,
        };
    }

    pub fn getChild(self: Node, name: []const u8) ?Node {
        var iter = self.children();
        while (iter.next()) |child| {
            if (std.mem.eql(u8, child.name, name)) {
                return child;
            }
        }
        return null;
    }

    pub fn compatible(self: Node) ?StringListIterator {
        const prop = self.getProperty("compatible") orelse return null;
        return prop.asStringList();
    }

    pub fn isCompatible(self: Node, compat: []const u8) bool {
        var iter = self.compatible() orelse return false;
        while (iter.next()) |c| {
            if (std.mem.eql(u8, c, compat)) return true;
        }
        return false;
    }

    pub fn reg(self: Node) ?[]const u8 {
        const prop = self.getProperty("reg") orelse return null;
        return prop.data;
    }

    pub fn status(self: Node) ?[]const u8 {
        const prop = self.getProperty("status") orelse return null;
        return prop.asString();
    }

    pub fn isEnabled(self: Node) bool {
        const s = self.status() orelse return true;
        return std.mem.eql(u8, s, "okay") or std.mem.eql(u8, s, "ok");
    }

    pub fn addressCells(self: Node) u32 {
        const prop = self.getProperty("#address-cells") orelse return 2;
        return prop.asU32() orelse 2;
    }

    pub fn sizeCells(self: Node) u32 {
        const prop = self.getProperty("#size-cells") orelse return 1;
        return prop.asU32() orelse 1;
    }

    pub fn phandle(self: Node) ?u32 {
        // Try both 'phandle' and legacy 'linux,phandle'
        if (self.getProperty("phandle")) |prop| {
            return prop.asU32();
        }
        if (self.getProperty("linux,phandle")) |prop| {
            return prop.asU32();
        }
        return null;
    }
};

pub const PropertyIterator = struct {
    fdt: *const Fdt,
    pos: usize,

    pub fn next(self: *PropertyIterator) ?Property {
        while (self.pos < self.fdt.struct_block.len) {
            const token = self.fdt.readToken(self.pos) orelse return null;
            self.pos += 4;

            switch (token) {
                .prop => {
                    const len = self.fdt.readU32(self.pos) orelse return null;
                    const nameoff = self.fdt.readU32(self.pos + 4) orelse return null;
                    self.pos += 8;

                    const name = self.fdt.getString(nameoff) orelse return null;
                    const data = if (len > 0 and self.pos + len <= self.fdt.struct_block.len)
                        self.fdt.struct_block[self.pos..][0..len]
                    else
                        &[_]u8{};

                    self.pos += alignUp(len, 4);

                    return Property{
                        .name = name,
                        .data = data,
                    };
                },
                .nop => continue,
                .begin_node, .end_node, .end => {
                    self.pos -= 4;
                    return null;
                },
            }
        }
        return null;
    }
};

pub const ChildIterator = struct {
    fdt: *const Fdt,
    pos: usize,
    depth: i32,

    pub fn next(self: *ChildIterator) ?Node {
        while (self.pos < self.fdt.struct_block.len) {
            const token = self.fdt.readToken(self.pos) orelse return null;
            self.pos += 4;

            switch (token) {
                .begin_node => {
                    const name_start = self.pos;
                    // Find end of name
                    while (self.pos < self.fdt.struct_block.len and self.fdt.struct_block[self.pos] != 0) {
                        self.pos += 1;
                    }
                    const name = self.fdt.struct_block[name_start..self.pos];
                    self.pos += 1; // skip null
                    self.pos = alignUp(self.pos, 4);

                    if (self.depth == 0) {
                        self.depth = 1;
                        return Node{
                            .name = name,
                            .fdt = self.fdt,
                            .struct_offset = self.pos,
                        };
                    } else {
                        self.depth += 1;
                    }
                },
                .end_node => {
                    self.depth -= 1;
                    if (self.depth < 0) return null;
                },
                .prop => {
                    const len = self.fdt.readU32(self.pos) orelse return null;
                    self.pos += 8;
                    self.pos += alignUp(len, 4);
                },
                .nop => {},
                .end => return null,
            }
        }
        return null;
    }
};

pub const NodeIterator = struct {
    fdt: *const Fdt,
    pos: usize,
    depth: usize = 0,
    include_root: bool = true,
    root_seen: bool = false,

    pub fn next(self: *NodeIterator) ?Node {
        while (self.pos < self.fdt.struct_block.len) {
            const token = self.fdt.readToken(self.pos) orelse return null;
            self.pos += 4;

            switch (token) {
                .begin_node => {
                    const name_start = self.pos;
                    while (self.pos < self.fdt.struct_block.len and self.fdt.struct_block[self.pos] != 0) {
                        self.pos += 1;
                    }
                    const name = self.fdt.struct_block[name_start..self.pos];
                    self.pos += 1;
                    self.pos = alignUp(self.pos, 4);
                    self.depth += 1;

                    // Skip root node if requested
                    if (!self.root_seen) {
                        self.root_seen = true;
                        if (!self.include_root) continue;
                    }

                    return Node{
                        .name = name,
                        .fdt = self.fdt,
                        .struct_offset = self.pos,
                    };
                },
                .end_node => {
                    if (self.depth > 0) self.depth -= 1;
                },
                .prop => {
                    const len = self.fdt.readU32(self.pos) orelse return null;
                    self.pos += 8;
                    self.pos += alignUp(len, 4);
                },
                .nop => {},
                .end => return null,
            }
        }
        return null;
    }
};

pub const RegIterator = struct {
    data: []const u8,
    address_cells: u32,
    size_cells: u32,
    pos: usize,

    pub const Entry = struct {
        address: u64,
        size: u64,
    };

    pub fn next(self: *RegIterator) ?Entry {
        const entry_size = (self.address_cells + self.size_cells) * 4;
        if (self.pos + entry_size > self.data.len) return null;

        var address: u64 = 0;
        var i: usize = 0;
        while (i < self.address_cells) : (i += 1) {
            const val = std.mem.bigToNative(u32, @as(*const u32, @alignCast(@ptrCast(self.data.ptr + self.pos + i * 4))).*);
            address = (address << 32) | val;
        }

        var size: u64 = 0;
        i = 0;
        while (i < self.size_cells) : (i += 1) {
            const val = std.mem.bigToNative(u32, @as(*const u32, @alignCast(@ptrCast(self.data.ptr + self.pos + (self.address_cells + i) * 4))).*);
            size = (size << 32) | val;
        }

        self.pos += entry_size;

        return Entry{
            .address = address,
            .size = size,
        };
    }
};

base: [*]const u8,
totalsize: u32,
version: u32,
last_comp_version: u32,
boot_cpuid_phys: u32,
struct_block: []const u8,
strings_block: []const u8,
mem_rsv_block: []const u8,

pub const Error = error{
    InvalidMagic,
    UnsupportedVersion,
    InvalidStructure,
};

pub fn parse(ptr: *const anyopaque) Error!Fdt {
    const base: [*]const u8 = @ptrCast(ptr);
    const header: *const RawHeader = @alignCast(@ptrCast(ptr));

    // Validate magic
    const magic = std.mem.bigToNative(u32, header.magic);
    if (magic != MAGIC) {
        return Error.InvalidMagic;
    }

    const version = std.mem.bigToNative(u32, header.version);
    const last_comp_version = std.mem.bigToNative(u32, header.last_comp_version);

    // We support FDT version 17 (current standard)
    if (last_comp_version > 17) {
        return Error.UnsupportedVersion;
    }

    const totalsize = std.mem.bigToNative(u32, header.totalsize);
    const off_dt_struct = std.mem.bigToNative(u32, header.off_dt_struct);
    const size_dt_struct = std.mem.bigToNative(u32, header.size_dt_struct);
    const off_dt_strings = std.mem.bigToNative(u32, header.off_dt_strings);
    const size_dt_strings = std.mem.bigToNative(u32, header.size_dt_strings);
    const off_mem_rsvmap = std.mem.bigToNative(u32, header.off_mem_rsvmap);

    return Fdt{
        .base = base,
        .totalsize = totalsize,
        .version = version,
        .last_comp_version = last_comp_version,
        .boot_cpuid_phys = std.mem.bigToNative(u32, header.boot_cpuid_phys),
        .struct_block = base[off_dt_struct..][0..size_dt_struct],
        .strings_block = base[off_dt_strings..][0..size_dt_strings],
        .mem_rsv_block = base[off_mem_rsvmap..totalsize],
    };
}

pub fn root(self: *const Fdt) ?Node {
    if (self.struct_block.len < 4) return null;

    const token = self.readToken(0) orelse return null;
    if (token != .begin_node) return null;

    var pos: usize = 4;
    // Skip root node name (usually empty)
    while (pos < self.struct_block.len and self.struct_block[pos] != 0) {
        pos += 1;
    }
    pos += 1;
    pos = alignUp(pos, 4);

    return Node{
        .name = "",
        .fdt = self,
        .struct_offset = pos,
    };
}

pub fn nodes(self: *const Fdt) NodeIterator {
    return .{
        .fdt = self,
        .pos = 0,
        .include_root = true,
    };
}

pub fn nodesWithoutRoot(self: *const Fdt) NodeIterator {
    return .{
        .fdt = self,
        .pos = 0,
        .include_root = false,
    };
}

pub fn findNode(self: *const Fdt, path: []const u8) ?Node {
    if (path.len == 0 or path[0] != '/') return null;

    var current = self.root() orelse return null;

    // Handle root path
    if (path.len == 1) return current;

    // Parse path components
    var remaining = path[1..]; // skip leading '/'
    while (remaining.len > 0) {
        // Find next component
        var end: usize = 0;
        while (end < remaining.len and remaining[end] != '/') {
            end += 1;
        }

        const component = remaining[0..end];
        if (component.len == 0) break;

        // Find child with this name
        current = current.getChild(component) orelse return null;

        // Move past this component and any '/'
        remaining = if (end < remaining.len) remaining[end + 1 ..] else &[_]u8{};
    }

    return current;
}

pub fn findByPhandle(self: *const Fdt, handle: u32) ?Node {
    var iter = self.nodes();
    while (iter.next()) |node| {
        if (node.phandle()) |ph| {
            if (ph == handle) return node;
        }
    }
    return null;
}

pub fn findCompatible(self: *const Fdt, compat: []const u8) CompatibleIterator {
    return .{
        .inner = self.nodesWithoutRoot(),
        .compat = compat,
    };
}

pub fn findFirstCompatible(self: *const Fdt, compat: []const u8) ?Node {
    var iter = self.findCompatible(compat);
    return iter.next();
}

pub fn memoryReservations(self: *const Fdt) MemoryReservationIterator {
    return .{
        .data = self.mem_rsv_block,
        .pos = 0,
    };
}

pub fn chosen(self: *const Fdt) ?Node {
    return self.findNode("/chosen");
}

pub fn memory(self: *const Fdt) ?Node {
    return self.findNode("/memory") orelse self.findFirstCompatible("memory");
}

pub fn cpus(self: *const Fdt) ?Node {
    return self.findNode("/cpus");
}

// Internal helpers
fn readToken(self: *const Fdt, offset: usize) ?Token {
    if (offset + 4 > self.struct_block.len) return null;
    const val = std.mem.bigToNative(u32, @as(*const u32, @alignCast(@ptrCast(self.struct_block.ptr + offset))).*);
    return @as(Token, @enumFromInt(val));
}

fn readU32(self: *const Fdt, offset: usize) ?u32 {
    if (offset + 4 > self.struct_block.len) return null;
    return std.mem.bigToNative(u32, @as(*const u32, @alignCast(@ptrCast(self.struct_block.ptr + offset))).*);
}

fn getString(self: *const Fdt, offset: u32) ?[]const u8 {
    if (offset >= self.strings_block.len) return null;
    const start = self.strings_block[offset..];
    for (start, 0..) |c, i| {
        if (c == 0) return start[0..i];
    }
    return null;
}

pub const CompatibleIterator = struct {
    inner: NodeIterator,
    compat: []const u8,

    pub fn next(self: *CompatibleIterator) ?Node {
        while (self.inner.next()) |node| {
            if (node.isCompatible(self.compat)) {
                return node;
            }
        }
        return null;
    }
};

pub const MemoryReservationIterator = struct {
    data: []const u8,
    pos: usize,

    pub fn next(self: *MemoryReservationIterator) ?struct { address: u64, size: u64 } {
        if (self.pos + 16 > self.data.len) return null;

        const entry: *const ReserveEntry = @alignCast(@ptrCast(self.data.ptr + self.pos));
        const result = entry.toNative();
        self.pos += 16;

        // End marker is all zeros
        if (result.address == 0 and result.size == 0) return null;

        return result;
    }
};

fn alignUp(value: anytype, alignment: @TypeOf(value)) @TypeOf(value) {
    return (value + alignment - 1) & ~(alignment - 1);
}

pub fn parseReg(data: []const u8, address_cells: u32, size_cells: u32) RegIterator {
    return .{
        .data = data,
        .address_cells = address_cells,
        .size_cells = size_cells,
        .pos = 0,
    };
}