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

// ---------------------------------------------------------------------------
// Number encoding/decoding
// ---------------------------------------------------------------------------

pub fn ofNumber(arena: *Arena, n: u64) !u32 {
    if (n == 0) {
        return try arena.alloc(.leaf);
    }
    const bit = if (n % 2 == 1) try arena.alloc(.{ .stem = .{ .child = try arena.alloc(.leaf) } }) else try arena.alloc(.leaf);
    const rest = try ofNumber(arena, n / 2);
    return try arena.alloc(.{ .fork = .{ .left = bit, .right = rest } });
}

pub fn toNumber(arena: *Arena, idx: u32) !?u64 {
    const node = try reduce.reduce(idx, arena, 10_000);
    const n = arena.get(node);
    return switch (n.*) {
        .leaf => 0,
        .stem => return null,
        .fork => |f| blk: {
            const bit_node = try reduce.reduce(f.left, arena, 10_000);
            const bit = arena.get(bit_node);
            const bit_val: u64 = switch (bit.*) {
                .leaf => 0,
                .stem => |s| if (arena.get(s.child).* == .leaf) 1 else return null,
                else => return null,
            };
            const rest = try toNumber(arena, f.right) orelse return null;
            break :blk bit_val + 2 * rest;
        },
        .app => return null,
    };
}

// ---------------------------------------------------------------------------
// List encoding/decoding
// ---------------------------------------------------------------------------

pub fn ofList(arena: *Arena, items: []const u32) !u32 {
    var result = try arena.alloc(.leaf);
    var i: usize = items.len;
    while (i > 0) {
        i -= 1;
        result = try arena.alloc(.{ .fork = .{ .left = items[i], .right = result } });
    }
    return result;
}

pub fn toList(arena: *Arena, idx: u32) !?std.ArrayList(u32) {
    var result = std.ArrayList(u32).empty;
    errdefer result.deinit(arena.allocator);

    var current = idx;
    while (true) {
        const node = try reduce.reduce(current, arena, 10_000);
        const n = arena.get(node);
        switch (n.*) {
            .leaf => return result,
            .stem => return null,
            .fork => |f| {
                try result.append(arena.allocator, f.left);
                current = f.right;
            },
            .app => return null,
        }
    }
}

// ---------------------------------------------------------------------------
// String / Bytes encoding/decoding
// Strings are lists of byte values (each character encoded as a number tree).
// ---------------------------------------------------------------------------

pub fn ofString(arena: *Arena, s: []const u8) !u32 {
    var bytes = try arena.allocator.alloc(u32, s.len);
    defer arena.allocator.free(bytes);
    for (s, 0..) |c, i| {
        bytes[i] = try ofNumber(arena, c);
    }
    return try ofList(arena, bytes);
}

pub fn toString(arena: *Arena, idx: u32) !?[]u8 {
    var list = try toList(arena, idx) orelse return null;
    defer list.deinit(arena.allocator);
    var result = try arena.allocator.alloc(u8, list.items.len);
    errdefer arena.allocator.free(result);
    for (list.items, 0..) |elem_idx, i| {
        const num = try toNumber(arena, elem_idx) orelse {
            arena.allocator.free(result);
            return null;
        };
        if (num > 255) {
            arena.allocator.free(result);
            return null;
        }
        result[i] = @intCast(num);
    }
    return result;
}

pub fn ofBytes(arena: *Arena, bytes: []const u8) !u32 {
    return try ofString(arena, bytes);
}

pub fn toBytes(arena: *Arena, idx: u32) !?[]u8 {
    return try toString(arena, idx);
}

// ---------------------------------------------------------------------------
// Result unwrapping (ok/err protocol)
// ok value rest = pair true (pair value rest)
// err code rest = pair false (pair code rest)
// ---------------------------------------------------------------------------

pub const UnwrapResult = struct {
    ok: bool,
    value: u32,
    rest: u32,
};

pub fn unwrapResult(arena: *Arena, idx: u32) !?UnwrapResult {
    const node = try reduce.reduce(idx, arena, 10_000);
    const n = arena.get(node);
    switch (n.*) {
        .fork => |f| {
            const tag = try reduce.reduce(f.left, arena, 10_000);
            const rest_pair = try reduce.reduce(f.right, arena, 10_000);
            const rp = arena.get(rest_pair);
            switch (rp.*) {
                .fork => |rf| {
                    const is_ok = tree.sameTree(arena, tag, try arena.alloc(.{ .stem = .{ .child = try arena.alloc(.leaf) } }));
                    return UnwrapResult{
                        .ok = is_ok,
                        .value = rf.left,
                        .rest = rf.right,
                    };
                },
                else => return null,
            }
        },
        else => return null,
    }
}

// ---------------------------------------------------------------------------
// Host ABI value unwrapping
// A host ABI value is: pair tag payload
// ---------------------------------------------------------------------------

pub const HostValue = struct {
    tag: u64,
    payload: u32,
};

pub fn unwrapHostValue(arena: *Arena, idx: u32) !?HostValue {
    const node = try reduce.reduce(idx, arena, 10_000);
    const n = arena.get(node);
    switch (n.*) {
        .fork => |f| {
            const tag_num = try toNumber(arena, f.left) orelse return null;
            return HostValue{ .tag = tag_num, .payload = f.right };
        },
        else => return null,
    }
}

/// Returns true if the tree is a valid boolean (Leaf=false, Stem Leaf=true).
pub fn isBool(arena: *Arena, idx: u32) !bool {
    const node = try reduce.reduce(idx, arena, 10_000);
    const n = arena.get(node);
    return switch (n.*) {
        .leaf => true,
        .stem => |s| arena.get(s.child).* == .leaf,
        else => false,
    };
}

/// Extract the boolean value: false for Leaf, true for Stem Leaf.
/// Returns null if the tree is not a valid boolean.
pub fn toBool(arena: *Arena, idx: u32) !?bool {
    const node = try reduce.reduce(idx, arena, 10_000);
    const n = arena.get(node);
    return switch (n.*) {
        .leaf => false,
        .stem => |s| if (arena.get(s.child).* == .leaf) true else null,
        else => null,
    };
}

// ---------------------------------------------------------------------------
// Host ABI tag constants
// ---------------------------------------------------------------------------

pub const HOST_TREE_TAG: u64 = 0;
pub const HOST_STRING_TAG: u64 = 1;
pub const HOST_NUMBER_TAG: u64 = 2;
pub const HOST_BOOL_TAG: u64 = 3;
pub const HOST_LIST_TAG: u64 = 4;
pub const HOST_BYTES_TAG: u64 = 5;