const std = @import("std");
const tree = @import("tree.zig");
const Arena = @import("arena.zig").Arena;
const reduce = @import("reduce.zig");
const codecs = @import("codecs.zig");
const kernel = @import("kernel.zig");
const bundle = @import("bundle.zig");
const io_driver = @import("io_driver.zig");

fn printNode(arena: *Arena, tag: u64, node: u32, io: std.Io) !void {
    var stdout_buf: [4096]u8 = undefined;
    var stdout = std.Io.File.stdout().writer(io, &stdout_buf);

    switch (tag) {
        codecs.HOST_STRING_TAG => {
            const s = try codecs.toString(arena, node) orelse {
                try stdout.interface.writeAll("Error: failed to decode string result\n");
                try stdout.flush();
                return error.DecodeFailed;
            };
            defer arena.allocator.free(s);
            try stdout.interface.writeAll(s);
            try stdout.interface.writeAll("\n");
        },
        codecs.HOST_NUMBER_TAG => {
            const n = try codecs.toNumber(arena, node) orelse 0;
            try stdout.interface.print("{d}\n", .{n});
        },
        codecs.HOST_BOOL_TAG => {
            const b = try codecs.toBool(arena, node) orelse {
                try stdout.interface.writeAll("Error: failed to decode bool result\n");
                try stdout.flush();
                return error.DecodeFailed;
            };
            try stdout.interface.writeAll(if (b) "true\n" else "false\n");
        },
        codecs.HOST_TREE_TAG => {
            try tree.formatTree(&stdout.interface, arena, node, 0);
            try stdout.interface.writeAll("\n");
        },
        else => {
            try stdout.interface.print("(tag={d}, payload=", .{tag});
            try tree.formatTree(&stdout.interface, arena, node, 0);
            try stdout.interface.writeAll(")\n");
        },
    }
    try stdout.flush();
}

fn runNative(arena: *Arena, tag: u64, bundle_bytes: []const u8, args_raw: []const []const u8, fuel: u64, io: std.Io) !void {
    const term = try bundle.loadBundleDefaultRoot(arena, bundle_bytes);

    var current = term;
    for (args_raw) |arg| {
        const arg_tree = try parseArg(arena, io, arg);
        current = try arena.alloc(.{ .app = .{ .func = current, .arg = arg_tree } });
    }

    const result = try reduce.reduce(current, arena, fuel);
    try printNode(arena, tag, result, io);
}

fn runIO(arena: *Arena, tag: u64, bundle_bytes: []const u8, args_raw: []const []const u8, fuel: u64, perms: io_driver.IOPerms, io: std.Io) !void {
    const term = try bundle.loadBundleDefaultRoot(arena, bundle_bytes);

    var current = term;
    for (args_raw) |arg| {
        const arg_tree = try parseArg(arena, io, arg);
        current = try arena.alloc(.{ .app = .{ .func = current, .arg = arg_tree } });
    }

    const reduced = try reduce.reduce(current, arena, fuel);

    if (try io_driver.isIOSentinel(arena, reduced) == null) {
        var stderr = std.Io.File.stderr().writer(io, &[_]u8{});
        try stderr.interface.writeAll("Error: reduced term is not a valid IO program\n");
        try stderr.flush();
        std.process.exit(1);
    }

    const result = try io_driver.runIO(arena.allocator, arena, reduced, perms);
    try printNode(arena, tag, result, io);
}

fn runBundle(arena: *Arena, tag: u64, bundle_bytes: []const u8, args_raw: []const []const u8, fuel: u64, io: std.Io) !void {
    const kernel_root = try kernel.loadKernel(arena);

    const tag_tree = try codecs.ofNumber(arena, tag);
    const bundle_tree = try codecs.ofBytes(arena, bundle_bytes);

    var arg_items = try arena.allocator.alloc(u32, args_raw.len);
    defer arena.allocator.free(arg_items);
    for (args_raw, 0..) |arg, i| {
        arg_items[i] = try parseArg(arena, io, arg);
    }
    const args_tree = try codecs.ofList(arena, arg_items);

    // Build: (((runArboricxTyped tag) bundle_bytes) args)
    const app0 = try arena.alloc(.{ .app = .{ .func = kernel_root, .arg = tag_tree } });
    const app1 = try arena.alloc(.{ .app = .{ .func = app0, .arg = bundle_tree } });
    const app2 = try arena.alloc(.{ .app = .{ .func = app1, .arg = args_tree } });

    const result = try reduce.reduce(app2, arena, fuel);

    const unwrapped = try codecs.unwrapResult(arena, result) orelse {
        var stderr = std.Io.File.stderr().writer(io, &[_]u8{});
        try stderr.interface.writeAll("Error: result is not a valid ok/err pair\n");
        try stderr.flush();
        return error.InvalidResult;
    };

    if (!unwrapped.ok) {
        var stderr = std.Io.File.stderr().writer(io, &[_]u8{});
        const code = try codecs.toNumber(arena, unwrapped.value) orelse 0;
        try stderr.interface.print("Error: kernel returned err, code={d}\n", .{code});
        try stderr.flush();
        return error.KernelError;
    }

    const hv = try codecs.unwrapHostValue(arena, unwrapped.value) orelse {
        var stderr = std.Io.File.stderr().writer(io, &[_]u8{});
        try stderr.interface.writeAll("Error: result is not a valid host ABI value\n");
        try stderr.flush();
        return error.InvalidHostValue;
    };

    try printNode(arena, hv.tag, hv.payload, io);
}

fn parseArg(arena: *Arena, io: std.Io, s: []const u8) !u32 {
    if (std.mem.endsWith(u8, s, ".arboricx")) {
        const bundle_bytes = try std.Io.Dir.cwd().readFileAlloc(io, s, arena.allocator, .limited(10 * 1024 * 1024));
        defer arena.allocator.free(bundle_bytes);
        return try bundle.loadBundleDefaultRoot(arena, bundle_bytes);
    }

    if (std.fmt.parseInt(u64, s, 10)) |n| {
        return try codecs.ofNumber(arena, n);
    } else |_| {}

    if (s.len >= 2 and s[0] == '"' and s[s.len - 1] == '"') {
        return try codecs.ofString(arena, s[1 .. s.len - 1]);
    }

    return try codecs.ofString(arena, s);
}

pub fn main(init: std.process.Init) !void {
    const gpa = init.gpa;
    const io = init.io;

    const args = try init.minimal.args.toSlice(init.arena.allocator());
    if (args.len < 2) {
        var stderr = std.Io.File.stderr().writer(io, &[_]u8{});
        try stderr.interface.writeAll("Usage: tricu-zig [--type TYPE] [--kernel] [--io] [--unsafe-io] [--fuel N] <bundle.arboricx> [arg1 arg2 ...]\n");
        try stderr.flush();
        std.process.exit(1);
    }

    // Parse options before bundle path
    var tag = codecs.HOST_STRING_TAG;
    var bundle_idx: usize = 1;
    var arg_start: usize = 2;

    var use_kernel = false;
    var use_io = false;
    var io_perms = io_driver.IOPerms{};
    var fuel: u64 = std.math.maxInt(u64);

    var i: usize = 1;
    while (i < args.len) : (i += 1) {
        if (std.mem.eql(u8, args[i], "--type")) {
            if (i + 1 >= args.len) {
                var stderr = std.Io.File.stderr().writer(io, &[_]u8{});
                try stderr.interface.writeAll("Usage: tricu-zig --type <tree|number|bool|string|list|bytes> [--io] [--unsafe-io] [--fuel N] <bundle> [args...]\n");
                try stderr.flush();
                std.process.exit(1);
            }
            const type_str = args[i + 1];
            tag = if (std.mem.eql(u8, type_str, "tree")) codecs.HOST_TREE_TAG
                else if (std.mem.eql(u8, type_str, "number")) codecs.HOST_NUMBER_TAG
                else if (std.mem.eql(u8, type_str, "bool")) codecs.HOST_BOOL_TAG
                else if (std.mem.eql(u8, type_str, "string")) codecs.HOST_STRING_TAG
                else if (std.mem.eql(u8, type_str, "list")) codecs.HOST_LIST_TAG
                else if (std.mem.eql(u8, type_str, "bytes")) codecs.HOST_BYTES_TAG
                else blk: {
                    var stderr = std.Io.File.stderr().writer(io, &[_]u8{});
                    try stderr.interface.print("Unknown type: {s}\n", .{type_str});
                    try stderr.flush();
                    std.process.exit(1);
                    break :blk codecs.HOST_STRING_TAG;
                };
            i += 1;
        } else if (std.mem.eql(u8, args[i], "--kernel")) {
            use_kernel = true;
        } else if (std.mem.eql(u8, args[i], "--io")) {
            use_io = true;
        } else if (std.mem.eql(u8, args[i], "--unsafe-io")) {
            io_perms.allow_read_all = true;
            io_perms.allow_write_all = true;
        } else if (std.mem.eql(u8, args[i], "--fuel")) {
            if (i + 1 >= args.len) {
                var stderr = std.Io.File.stderr().writer(io, &[_]u8{});
                try stderr.interface.writeAll("Usage: tricu-zig --fuel <N> [--io] [--unsafe-io] <bundle> [args...]\n");
                try stderr.flush();
                std.process.exit(1);
            }
            const n = std.fmt.parseInt(u64, args[i + 1], 10) catch {
                var stderr = std.Io.File.stderr().writer(io, &[_]u8{});
                try stderr.interface.print("Invalid fuel: {s}\n", .{args[i + 1]});
                try stderr.flush();
                std.process.exit(1);
            };
            fuel = std.math.mul(u64, n, 1_000_000) catch std.math.maxInt(u64);
            i += 1;
        } else {
            bundle_idx = i;
            arg_start = i + 1;
            break;
        }
    }

    if (bundle_idx >= args.len) {
        var stderr = std.Io.File.stderr().writer(io, &[_]u8{});
        try stderr.interface.writeAll("Usage: tricu-zig [--type TYPE] [--kernel] [--io] [--unsafe-io] [--fuel N] <bundle.arboricx> [arg1 arg2 ...]\n");
        try stderr.flush();
        std.process.exit(1);
    }

    const bundle_path = args[bundle_idx];
    const bundle_bytes = try std.Io.Dir.cwd().readFileAlloc(io, bundle_path, gpa, .limited(10 * 1024 * 1024));
    defer gpa.free(bundle_bytes);

    var arena = Arena.init(gpa);
    defer arena.deinit();

    const call_args = if (arg_start < args.len) args[arg_start..] else &[_][]const u8{};

    if (use_io) {
        runIO(&arena, tag, bundle_bytes, call_args, fuel, io_perms, io) catch |err| {
            var stderr = std.Io.File.stderr().writer(io, &[_]u8{});
            try stderr.interface.print("Execution failed: {s}\n", .{@errorName(err)});
            try stderr.flush();
            std.process.exit(1);
        };
    } else if (use_kernel) {
        runBundle(&arena, tag, bundle_bytes, call_args, fuel, io) catch |err| {
            var stderr = std.Io.File.stderr().writer(io, &[_]u8{});
            try stderr.interface.print("Execution failed: {s}\n", .{@errorName(err)});
            try stderr.flush();
            std.process.exit(1);
        };
    } else {
        runNative(&arena, tag, bundle_bytes, call_args, fuel, io) catch |err| {
            var stderr = std.Io.File.stderr().writer(io, &[_]u8{});
            try stderr.interface.print("Execution failed: {s}\n", .{@errorName(err)});
            try stderr.flush();
            std.process.exit(1);
        };
    }
}