const std = @import("std");
const tree = @import("tree.zig");
const Arena = @import("arena.zig").Arena;

pub const Error = error{
    InvalidMagic,
    InvalidVersion,
    Truncated,
    InvalidManifest,
    InvalidNodePayload,
    ExportNotFound,
    MissingChild,
    UnexpectedFormat,
    OutOfMemory,
};

const Parser = struct {
    bytes: []const u8,
    pos: usize,

    fn init(bytes: []const u8) Parser {
        return .{ .bytes = bytes, .pos = 0 };
    }

    fn remaining(self: *const Parser) usize {
        return self.bytes.len - self.pos;
    }

    fn expect(self: *Parser, n: usize) Error![]const u8 {
        if (self.remaining() < n) return error.Truncated;
        const result = self.bytes[self.pos .. self.pos + n];
        self.pos += n;
        return result;
    }

    fn readU8(self: *Parser) Error!u8 {
        const b = try self.expect(1);
        return b[0];
    }

    fn readU16(self: *Parser) Error!u16 {
        const b = try self.expect(2);
        return std.mem.readInt(u16, b[0..2], .big);
    }

    fn readU32(self: *Parser) Error!u32 {
        const b = try self.expect(4);
        return std.mem.readInt(u32, b[0..4], .big);
    }

    fn readU64(self: *Parser) Error!u64 {
        const b = try self.expect(8);
        return std.mem.readInt(u64, b[0..8], .big);
    }

    fn readLengthPrefixedBytes(self: *Parser, allocator: std.mem.Allocator) Error![]const u8 {
        const len = try self.readU32();
        const bytes = try self.expect(len);
        const copy = try allocator.alloc(u8, bytes.len);
        @memcpy(copy, bytes);
        return copy;
    }
};

const SectionEntry = struct {
    section_type: u32,
    offset: u64,
    length: u64,
};

fn parseHeader(p: *Parser) Error!struct { major: u16, minor: u16, section_count: u32, dir_offset: u64 } {
    const magic = try p.expect(8);
    if (!std.mem.eql(u8, magic, "ARBORICX")) return error.InvalidMagic;

    const major = try p.readU16();
    const minor = try p.readU16();
    const section_count = try p.readU32();
    _ = try p.readU64(); // flags
    const dir_offset = try p.readU64();

    if (major != 1) return error.InvalidVersion;

    return .{ .major = major, .minor = minor, .section_count = section_count, .dir_offset = dir_offset };
}

fn parseSectionEntries(p: *Parser, count: u32, allocator: std.mem.Allocator) Error![]SectionEntry {
    const entries = try allocator.alloc(SectionEntry, count);
    errdefer allocator.free(entries);

    for (entries) |*entry| {
        entry.section_type = try p.readU32();
        _ = try p.readU16(); // section_version
        _ = try p.readU16(); // section_flags
        const compression = try p.readU16();
        _ = try p.readU16(); // reserved (was digest_alg)
        entry.offset = try p.readU64();
        entry.length = try p.readU64();
        _ = try p.readU32(); // reserved padding

        if (compression != 0) return error.UnexpectedFormat;
    }
    return entries;
}

fn parseManifest(p: *Parser, allocator: std.mem.Allocator) Error!struct { exports: []Export, roots: []Root } {
    const magic = try p.expect(8);
    if (!std.mem.eql(u8, magic, "ARBMNFST")) return error.InvalidManifest;

    const major = try p.readU16();
    _ = try p.readU16(); // minor
    if (major != 1) return error.InvalidVersion;

    const schema = try p.readLengthPrefixedBytes(allocator);
    defer allocator.free(schema);
    if (!std.mem.eql(u8, schema, "arboricx.bundle.manifest.v1")) return error.UnexpectedFormat;

    const bundle_type = try p.readLengthPrefixedBytes(allocator);
    defer allocator.free(bundle_type);
    if (!std.mem.eql(u8, bundle_type, "tree-calculus-executable-object")) return error.UnexpectedFormat;

    const calc = try p.readLengthPrefixedBytes(allocator);
    defer allocator.free(calc);
    if (!std.mem.eql(u8, calc, "tree-calculus.v1")) return error.UnexpectedFormat;

    const hash_alg = try p.readLengthPrefixedBytes(allocator);
    defer allocator.free(hash_alg);
    if (!std.mem.eql(u8, hash_alg, "indexed")) return error.UnexpectedFormat;

    const hash_domain = try p.readLengthPrefixedBytes(allocator);
    defer allocator.free(hash_domain);
    if (!std.mem.eql(u8, hash_domain, "arboricx.indexed.node.v1")) return error.UnexpectedFormat;

    const payload_type = try p.readLengthPrefixedBytes(allocator);
    defer allocator.free(payload_type);
    if (!std.mem.eql(u8, payload_type, "arboricx.indexed.payload.v1")) return error.UnexpectedFormat;

    const sem = try p.readLengthPrefixedBytes(allocator);
    defer allocator.free(sem);
    if (!std.mem.eql(u8, sem, "tree-calculus.v1")) return error.UnexpectedFormat;

    const eval_mode = try p.readLengthPrefixedBytes(allocator);
    defer allocator.free(eval_mode);
    if (!std.mem.eql(u8, eval_mode, "normal-order")) return error.UnexpectedFormat;

    const abi = try p.readLengthPrefixedBytes(allocator);
    defer allocator.free(abi);
    if (!std.mem.eql(u8, abi, "arboricx.abi.tree.v1")) return error.UnexpectedFormat;

    const cap_count = try p.readU32();
    var i: u32 = 0;
    while (i < cap_count) : (i += 1) {
        const cap = try p.readLengthPrefixedBytes(allocator);
        defer allocator.free(cap);
        if (cap.len != 0) return error.UnexpectedFormat;
    }

    const closure = try p.readU8();
    if (closure != 0) return error.UnexpectedFormat;

    const root_count = try p.readU32();
    const roots = try allocator.alloc(Root, root_count);
    errdefer allocator.free(roots);
    for (roots) |*r| {
        r.index = try p.readU32();
        r.role = try p.readLengthPrefixedBytes(allocator);
    }

    const export_count = try p.readU32();
    const exports = try allocator.alloc(Export, export_count);
    errdefer {
        for (exports) |*e| {
            allocator.free(e.name);
            allocator.free(e.kind);
            allocator.free(e.abi);
        }
        allocator.free(exports);
    }
    for (exports) |*e| {
        e.name = try p.readLengthPrefixedBytes(allocator);
        e.root = try p.readU32();
        e.kind = try p.readLengthPrefixedBytes(allocator);
        e.abi = try p.readLengthPrefixedBytes(allocator);
        if (!std.mem.eql(u8, e.abi, "arboricx.abi.tree.v1")) return error.UnexpectedFormat;
    }

    const metadata_count = try p.readU32();
    var m: u32 = 0;
    while (m < metadata_count) : (m += 1) {
        _ = try p.readU16(); // tag
        const len = try p.readU32();
        _ = try p.expect(len);
    }

    const ext_count = try p.readU32();
    var e_idx: u32 = 0;
    while (e_idx < ext_count) : (e_idx += 1) {
        _ = try p.readU16(); // tag
        const len = try p.readU32();
        _ = try p.expect(len);
    }

    return .{ .exports = exports, .roots = roots };
}

const Export = struct {
    name: []const u8,
    root: u32,
    kind: []const u8,
    abi: []const u8,
};

const Root = struct {
    index: u32,
    role: []const u8,
};

/// Parse the node section and build nodes directly into the arena.
/// Returns a slice mapping node-section index -> arena index.
/// The caller owns the returned slice and must free it with the arena's allocator.
fn parseNodeSection(p: *Parser, arena: *Arena) Error![]u32 {
    const node_count = try p.readU64();
    const indices = try arena.allocator.alloc(u32, node_count);
    errdefer arena.allocator.free(indices);

    var i: u64 = 0;
    while (i < node_count) : (i += 1) {
        const plen = try p.readU32();
        const payload = try p.expect(plen);

        if (payload.len == 0) return error.InvalidNodePayload;

        const idx: u32 = switch (payload[0]) {
            0x00 => blk: {
                if (plen != 1) return error.InvalidNodePayload;
                break :blk try arena.alloc(.leaf);
            },
            0x01 => blk: {
                if (plen != 5) return error.InvalidNodePayload;
                const child_idx = std.mem.readInt(u32, payload[1..5], .big);
                if (child_idx >= i) return error.InvalidNodePayload;
                break :blk try arena.alloc(.{ .stem = .{ .child = indices[child_idx] } });
            },
            0x02 => blk: {
                if (plen != 9) return error.InvalidNodePayload;
                const left_idx = std.mem.readInt(u32, payload[1..5], .big);
                const right_idx = std.mem.readInt(u32, payload[5..9], .big);
                if (left_idx >= i or right_idx >= i) return error.InvalidNodePayload;
                break :blk try arena.alloc(.{ .fork = .{ .left = indices[left_idx], .right = indices[right_idx] } });
            },
            else => return error.InvalidNodePayload,
        };
        indices[i] = idx;
    }

    return indices;
}

fn findSection(entries: []SectionEntry, section_type: u32) ?SectionEntry {
    for (entries) |entry| {
        if (entry.section_type == section_type) return entry;
    }
    return null;
}

/// Parse an Arboricx bundle and load the named export into the arena.
/// Returns the arena index of the exported term tree.
pub fn loadBundleExport(
    arena: *Arena,
    bundle_bytes: []const u8,
    export_name: []const u8,
) Error!u32 {
    var p = Parser.init(bundle_bytes);

    const header = try parseHeader(&p);

    p.pos = @intCast(header.dir_offset);
    const allocator = arena.allocator;
    const entries = try parseSectionEntries(&p, header.section_count, allocator);
    defer allocator.free(entries);

    const manifest_section = findSection(entries, 1) orelse return error.InvalidManifest;
    const nodes_section = findSection(entries, 2) orelse return error.InvalidNodePayload;

    const manifest_bytes = bundle_bytes[@intCast(manifest_section.offset)..@intCast(manifest_section.offset + manifest_section.length)];
    const nodes_bytes = bundle_bytes[@intCast(nodes_section.offset)..@intCast(nodes_section.offset + nodes_section.length)];

    var mp = Parser.init(manifest_bytes);
    const manifest = try parseManifest(&mp, allocator);
    defer {
        for (manifest.exports) |e| {
            allocator.free(e.name);
            allocator.free(e.kind);
            allocator.free(e.abi);
        }
        allocator.free(manifest.exports);
        for (manifest.roots) |r| {
            allocator.free(r.role);
        }
        allocator.free(manifest.roots);
    }

    var export_root: ?u32 = null;
    for (manifest.exports) |e| {
        if (std.mem.eql(u8, e.name, export_name)) {
            export_root = e.root;
            break;
        }
    }
    const root_index = export_root orelse return error.ExportNotFound;

    var np = Parser.init(nodes_bytes);
    const node_indices = try parseNodeSection(&np, arena);
    defer allocator.free(node_indices);

    if (root_index >= node_indices.len) return error.InvalidNodePayload;
    return node_indices[root_index];
}

/// Parse an Arboricx bundle and load the default (first) root into the arena.
pub fn loadBundleDefaultRoot(
    arena: *Arena,
    bundle_bytes: []const u8,
) Error!u32 {
    var p = Parser.init(bundle_bytes);

    const header = try parseHeader(&p);

    p.pos = @intCast(header.dir_offset);
    const allocator = arena.allocator;
    const entries = try parseSectionEntries(&p, header.section_count, allocator);
    defer allocator.free(entries);

    const manifest_section = findSection(entries, 1) orelse return error.InvalidManifest;
    const nodes_section = findSection(entries, 2) orelse return error.InvalidNodePayload;

    const manifest_bytes = bundle_bytes[@intCast(manifest_section.offset)..@intCast(manifest_section.offset + manifest_section.length)];
    const nodes_bytes = bundle_bytes[@intCast(nodes_section.offset)..@intCast(nodes_section.offset + nodes_section.length)];

    var mp = Parser.init(manifest_bytes);
    const manifest = try parseManifest(&mp, allocator);
    defer {
        for (manifest.exports) |e| {
            allocator.free(e.name);
            allocator.free(e.kind);
            allocator.free(e.abi);
        }
        allocator.free(manifest.exports);
        for (manifest.roots) |r| {
            allocator.free(r.role);
        }
        allocator.free(manifest.roots);
    }

    if (manifest.roots.len == 0) return error.ExportNotFound;
    const root_index = manifest.roots[0].index;

    var np = Parser.init(nodes_bytes);
    const node_indices = try parseNodeSection(&np, arena);
    defer allocator.free(node_indices);

    if (root_index >= node_indices.len) return error.InvalidNodePayload;
    return node_indices[root_index];
}