const std = @import("std");
const codegen = @import("./codegen.zig");
const target = @import("builtin").target;

pub const parse = @import("./parse.zig");
pub const Lexer = @import("./Lexer.zig");
pub const compile = @import("./compile.zig");

const blue = "\x1b[34m";
const red = "\x1b[31m";
const normal = "\x1b[0m";

pub fn main() !u8 {
    var arena: std.heap.ArenaAllocator = .init(std.heap.smp_allocator);
    defer arena.deinit();
    const allocator = arena.allocator();

    var args = std.process.args();
    _ = args.next();

    const in_path = args.next();
    const in_file = if (in_path) |path|
        try std.fs.cwd().openFile(path, .{})
    else
        std.io.getStdIn();

    const out_path = args.next();
    const out_file = if (out_path) |path|
        try std.fs.cwd().createFile(path, .{ .mode = 0o777 })
    else
        std.io.getStdOut();

    const run = if (args.next()) |arg| std.mem.eql(u8, arg, "run") else false;

    const source = try in_file.readToEndAlloc(
        allocator,
        640 * 1024, // ought to be enough for anyone
    );
    defer allocator.free(source);

    var lexer: Lexer = .{ .source = source };
    const ast = parse.file(allocator, &lexer) catch |err| {
        std.debug.print(red ++ "parsing error: {}\n" ++ normal, .{err});
        return 1;
    };
    std.debug.print(blue ++ "parse tree:" ++ normal ++ "\n{}\n", .{parse.fmt(ast, source, 0)});
    if (lexer.peek().type != .eof) {
        std.debug.print(red ++ "Unexpected token {}, expected end of file\n" ++ normal, .{lexer.next()});
        return 1;
    }
    const module = compile.compile(allocator, source, ast) catch |err| {
        std.debug.print(red ++ "compilation error: {}\n" ++ normal, .{err});
        return 1;
    };
    std.debug.print(blue ++ "bytecode instructions: " ++ normal ++ "\n{}", .{module});
    const elf = codegen.create_elf(allocator, module) catch |err| {
        std.debug.print(red ++ "code generation error: {}\n" ++ normal, .{err});
        return 1;
    };
    try out_file.writer().writeAll(elf);
    std.debug.print(blue ++ "running program:" ++ normal ++ "\n", .{});
    if (run) {
        out_file.close();

        const err = std.process.execv(
            allocator,
            if (target.cpu.arch == .riscv64 and target.os.tag == .linux)
                &.{out_path}
            else
                &.{ "qemu-riscv64", out_path.? },
        );
        std.debug.print(red ++ "{}\n" ++ normal ++ "{any}\n", .{ err, @errorReturnTrace() });
        return 1;
    } else {
        return 0;
    }
}

fn HashMapFormatter(HashMap: type) type {
    return std.fmt.Formatter(struct {
        fn formatHashMap(
            hash_map: HashMap,
            comptime fmt: []const u8,
            options: std.fmt.FormatOptions,
            writer: anytype,
        ) !void {
            const k_fmt, const v_fmt = comptime blk: {
                var fmt_it = std.mem.splitScalar(u8, fmt, ',');
                const k_fmt = fmt_it.next() orelse "";
                const v_fmt = fmt_it.next() orelse "";
                break :blk .{ k_fmt, v_fmt };
            };
            _ = options;
            try writer.writeAll("{");
            var it = hash_map.iterator();
            var first = true;
            while (it.next()) |kv| {
                try writer.print(
                    "{s} {" ++ k_fmt ++ "}: {" ++ v_fmt ++ "}",
                    .{
                        if (first) "" else ",",
                        kv.key_ptr.*,
                        kv.value_ptr.*,
                    },
                );
                first = false;
            }
            try writer.writeAll(if (first) "}" else " }");
        }
    }.formatHashMap);
}

pub fn fmtHashMap(hash_map: anytype) HashMapFormatter(@TypeOf(hash_map)) {
    return .{ .data = hash_map };
}