const std = @import("std");
const stdout = std.io.getStdOut().writer();

// ANSI colors!
const cdone = "\x1b[0m";
const red = "\x1b[31m";
const green = "\x1b[32m";
const yellow = "\x1b[33m";
const magenta = "\x1b[35m";
const cyan = "\x1b[36m";
const gray = "\x1b[37m";
const gray2 = "\x1b[90m";
const bright_magenta = "\x1b[95m";
const bright_yellow = "\x1b[93m";
const bright_cyan = "\x1b[96m";
const bgmagenta = "\x1b[30;105m";

// A global slice will "point" to the data buffer
// so functions can access it directly.
var buffer: []u8 = undefined;

// Gets a bit of buffer as an integer of given type
fn get(addr: usize, T: anytype) T {
    return std.mem.readIntSliceNative(T, buffer[addr..]);
}

// Verifies that we found what we expected at the given address.
fn expect(addr: usize, expected: anytype, desc: []const u8) void {
    var val = get(addr, @TypeOf(expected));

    if (expected != val) {
        stdout.print("\n{s}At 0x{x} expected '{x}', but found '{x}' ({s}){s}\n", .{red, addr, expected, val, desc, cdone})
            catch unreachable;
        std.os.exit(1);
    }
}

fn printByte(addr: usize) void {
    stdout.print("{X:0>2} ", .{buffer[addr]})
        catch unreachable;
}

fn printByteColor(addr: usize, color: []const u8) void {
    stdout.print("{s}{X:0>2}{s} ", .{color, buffer[addr], cdone})
        catch unreachable;
}

fn printBytes(addr: usize, count: usize) void {
    for (0..count) |i| {
        printByte(addr+i);
    }
}

fn printBytesColor(addr: usize, count: usize, color: []const u8) void {
    for (0..count) |i| {
        printByteColor(addr+i, color);
    }
}

fn printCharColor(addr: usize, color: []const u8) void {
    stdout.print("{s}{c}{s}  ", .{color, buffer[addr], cdone})
        catch unreachable;
}

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

    const args = try std.process.argsAlloc(allocator);
    if(args.len < 2){
        std.debug.print("Usage: mez <ELF BINARY>\n", .{});
        std.os.exit(1);
    }

    var file = try std.fs.cwd().openFile(args[1], .{});
    defer file.close();

    const file_size = (try file.stat()).size;
    buffer = try allocator.alloc(u8, file_size);
//    var actual_buffer = try allocator.alloc(u8, file_size);

    // Set global slice "pointer"
//    buffer = actual_buffer; 

    try file.reader().readNoEof(buffer);

    // ===========================
    // || Print Main ELF Header ||
    // ===========================
    try stdout.print("+-[ELF Header]------------------------------------------+\n", .{});

    // ROW 1
    // ===========================
    try stdout.print("| ", .{});

    printByteColor(0x00, green);
    expect(0x00, @as(u8, 0x7F), "Magic number 0x7F");

    printCharColor(0x01, green);
    expect(0x01, @as(u8, 'E'), "Magic number 'E'");

    printCharColor(0x02, green);
    expect(0x02, @as(u8, 'L'), "Magic number 'L'");

    printCharColor(0x03, green);
    expect(0x03, @as(u8, 'F'), "Magic number 'F'");

    printByte(0x04);
    expect(0x04, @as(u8, 1), "1=32 bit arch");

    printByte(0x05);
    expect(0x05, @as(u8, 1), "1=little endian");

    printByte(0x06);
    expect(0x06, @as(u8, 1), "1=ELF version is current");

    printByte(0x07);
    expect(0x07, @as(u8, 0), "0=System V ABI");

    inline for(0x08..0x10) |addr| {
        // print padding zeroes as gray to de-emph
        printByteColor(addr, gray);
        expect(addr, @as(u8, 0), "Padding 0s");
    }

    printBytes(0x10, 2);
    expect(0x10, @as(u16, 2), "2=Executable file");
    try stdout.print("|\n", .{});

    // ROW 2
    // ===========================
    try stdout.print("| ", .{});

    printBytes(0x12, 2);
    expect(0x12, @as(u16, 3), "3=x386 ISA (arch)");

    printBytes(0x14, 4);
    expect(0x14, @as(u32, 1), "1=ELF version again");

    // e_entry - entry address (in memory)
    printBytesColor(0x18, 4, bright_magenta);
    const e_entry = get(0x18, u32);

    // e_phoff - program header offset (in file)
    printBytesColor(0x1C, 4, bright_yellow);
    const e_phoff = get(0x1C, u32);

    // e_shoff - section header offset (in file), ignore
    printBytesColor(0x20, 4, gray);
    try stdout.print("|\n", .{});

    // ROW 3
    // ===========================
    try stdout.print("| ", .{});

    // e_flags - totally ignoring
    printBytesColor(0x24, 4, gray);

    // e_hsize - don't need to display, but will check
    printBytes(0x28, 2);
    expect(0x28, @as(u16, 52), "header size should be 52 bytes (0x34) for 32-bit");

    // e_phentsize - the size of program header entries
    printBytesColor(0x2A, 2, bright_yellow);
    expect(0x2A, @as(u16, 32), "expecting program headers to be 32 bytes (0x20)");
    const e_phentsize = 32;

    // e_phnum - the count of program header entries
    printBytesColor(0x2C, 2, bright_yellow);
    const e_phnum = get(0x2C, u16);

    // e_shentsize - section header size, ignoring
    printBytesColor(0x2E, 2, gray);

    // e_shnum - section header count, ignoring
    printBytesColor(0x30, 2, gray);

    // e_shstrndx - section header string table offset, ignoring
    printBytesColor(0x32, 2, gray);
    try stdout.print("/-----+\n", .{});
    try stdout.print("+--+---------------------------------------------/\n", .{});

    // Main ELF Header Summary (decoded)
    // =================================
    try stdout.print("   +-- Entry point address: {s}0x{X:0>8}{s}\n", .{bright_magenta, e_entry, cdone});
    try stdout.print("   +-- Program header file offset: {s}0x{X}{s}\n", .{bright_yellow, e_phoff, cdone});
    try stdout.print("   +-- Program header size: {s}{d} (0x{x}){s}\n", .{bright_yellow, e_phentsize, e_phentsize, cdone});
    try stdout.print("   \\-- Program header count: {s}{d}{s}\n", .{bright_yellow, e_phnum, cdone});
    try stdout.print("    |\n", .{});
    for (0..e_phnum) |ph_num| {
        const ph_offset: usize = e_phoff + (ph_num * e_phentsize);
        const ph_type = switch(get(ph_offset, u32)){
            0 => cyan ++ "PT_NULL" ++ cdone,
            1 => bright_cyan ++ "PT_LOAD" ++ cdone,
            2 => cyan ++ "PT_DYNAMIC" ++ cdone,
            3 => cyan ++ "PT_INTERP" ++ cdone,
            else => red ++ "PT_WTF (a surprise)" ++ cdone,
        };

        try stdout.print("    +-- Program Header {d} at 0x{x}, type: {s}\n", .{ph_num, ph_offset,ph_type});
    }

    // =========================================
    // || Print all LOAD-type program headers ||
    // =========================================
    for (0..e_phnum) |ph_num| {
        const ph_offset: usize = e_phoff + (ph_num * e_phentsize);
        if (get(ph_offset, u32) != 1) {
            continue;
        }

        // Program Header Row 1
        // =========================
        try stdout.print("\n+-[{s}Program Header {d}{s}]------------------------------+\n", .{bright_yellow, ph_num, cdone});
        try stdout.print("| ", .{});

        // p_type (decoded in summary above)
        printBytesColor(ph_offset, 4, bright_cyan);

        // p_offset - file offset of data image to load
        printBytesColor(ph_offset + 0x04, 4, yellow);
        const p_offset = get(ph_offset + 0x04, u32);

        // p_vaddr - memory segment start address
        printBytesColor(ph_offset + 0x08, 4, magenta);
        const p_vaddr = get(ph_offset + 0x08, u32);

        // p_paddr - physical address, ignore
        printBytesColor(ph_offset + 0x0C, 4, gray);
        try stdout.print("|\n", .{});

        // Program Header Row 2
        // =========================
        try stdout.print("| ", .{});

        // p_filesz - data image size to load from file
        printBytesColor(ph_offset + 0x10, 4, yellow);
        const p_filesz = get(ph_offset + 0x10, u32);

        // p_memsz - memory segment size
        printBytesColor(ph_offset + 0x14, 4, magenta);
        const p_memsz = get(ph_offset + 0x14, u32);

        const p_vaddr_end = p_vaddr + p_memsz;

        // p_flags - memory segment flags
        printBytesColor(ph_offset + 0x18, 4, magenta);
        const p_flags = get(ph_offset + 0x18, u32);

        const rwx_decode = switch(p_flags) {
            1 => "X",
            2 => "W",
            3 => "W+X",
            4 => "R",
            5 => "R+X",
            6 => "R+W",
            7 => "R+W+X",
            else => "WTF",
        };

        // p_align - memory segment alignment
        printBytesColor(ph_offset + 0x1C, 4, gray);
        try stdout.print("|\n", .{});

        // Program Header Summary (decoded)
        // ================================
        try stdout.print("+--+----------------------------------------------+\n", .{});
        try stdout.print("   +-- File data start offset: {s}0x{x}{s}\n", .{yellow, p_offset, cdone});
        try stdout.print("   +-- File data bytes to load: {s}{d} (0x{x}){s}\n", .{yellow, p_filesz, p_filesz, cdone});
        try stdout.print("   +-- Memory segment start addr: {s}0x{x:0>8}{s}\n", .{magenta, p_vaddr, cdone});
        try stdout.print("   +-- Memory segment byte size: {s}{d} (0x{x}){s}\n", .{magenta, p_memsz, p_memsz, cdone});
        try stdout.print("   +-- Memory segment flags: {s}{s} (0x{x}){s}\n", .{magenta, rwx_decode, p_flags, cdone});

        // If we contain the entry point address, display it!
        var need_to_print_entry = (e_entry >= p_vaddr and e_entry < p_vaddr_end);
        if (need_to_print_entry) {
            try stdout.print("   +-- {s}Contains entry point 0x{X:0>8}{s}\n", .{bright_magenta, e_entry, cdone});
        }

        // Hex Dump data!
        // ================================
        const dump_lines = 4;
        const dump_columns = 12;
        var data_start = p_offset;
        var data_end = p_offset + p_filesz;
        var mem_start = p_vaddr;
        var line: usize = 0;
try stdout.print("from {x} to {x}...\n", .{mem_start, mem_start+p_filesz});

        while (line < dump_lines) {

            var mybyte = data_start + (line * dump_columns);
            var myaddr = mem_start + (line * dump_columns);

            // We've printed it all!
            if (mybyte >= data_end) {
                break;
            }

            // Memory address at start of line
            try stdout.print("0x{x:0>8} ", .{myaddr});

            // Print columns of bytes
            try stdout.print("{s}", .{gray});
            for (0..dump_columns) |i| {
                if (mybyte + i >= data_end) {
                    // Print spaces to complete row
                    try stdout.print("   ", .{});
                    continue;
                }

                if(myaddr + i == e_entry) {
                    try stdout.print("{s}{x:0>2}{s} ", .{bgmagenta, buffer[mybyte + i], cdone});
                    need_to_print_entry = false;
                    continue;
                }

                try stdout.print("{x:0>2} ", .{buffer[mybyte + i]});
            }
            try stdout.print("{s}", .{cdone});

            // Print columns of ASCII
            for (0..dump_columns) |i| {
                if (mybyte + i >= data_end) {
                    break;
                }

                const m = buffer[mybyte + i];

                if(m > 32 and m < 127){
                    try stdout.print("{s}{c}{s}", .{cyan, m, cdone});
                }
                else {
                    try stdout.print(".", .{});
                }
            }

            line += 1;

            // If that was the last line and we need to print
            // the entry address, reset lines and set the start
            // address so it'll start with the entry's line.
            if (line == dump_lines and need_to_print_entry){
                var skip = e_entry - p_vaddr;
                skip -= skip % dump_columns;
                try stdout.print("\n  ...Skipping to entry point...", .{});
                data_start = p_offset + skip;
                mem_start = p_vaddr + skip;
                line = 0;
            }

            try stdout.print("\n", .{});
        }

        // Did we print it all? If not, show how much 'til end
        const displayed_to = data_start + dump_lines * dump_columns;
        if (displayed_to < data_end) {
            try stdout.print("  ...{d} more bytes to load...\n", .{data_end - displayed_to});
        }
    }
}