const std = @import("std");
const cs = @import("capstone");

const SymbolRange = struct {
    start: u64,
    end: u64,
    name: []u8,
    kind: u8,
};

fn iterateSymbols(
    h: std.elf.Header,
    file_reader: *std.fs.File.Reader,
    symtab: std.elf.Elf64_Shdr,
) SymbolIterator {
    return .{
        .elf_header = h,
        .file_reader = file_reader,
        .symtab = symtab,
    };
}

const SymbolIterator = struct {
    elf_header: std.elf.Header,
    file_reader: *std.fs.File.Reader,
    symtab: std.elf.Elf64_Shdr,
    index: usize = 0,

    pub fn next(it: *SymbolIterator) !?std.elf.Elf64_Sym {
        defer it.index += 1;

        const size: u64 = if (it.elf_header.is_64) @sizeOf(std.elf.Elf64_Sym) else @sizeOf(std.elf.Elf64_Sym);
        const offset = it.symtab.sh_offset + size * it.index;

        if (offset >= (it.symtab.sh_size + it.symtab.sh_offset))
            return null;

        try it.file_reader.seekTo(offset);
        return try it.file_reader.interface.takeStruct(std.elf.Elf64_Sym, it.elf_header.endian);
    }
};

pub fn main() !void {
    var args = std.process.args();
    _ = args.skip(); // skip argv[0]

    const bw = std.debug.lockStderrWriter(&.{});
    defer std.debug.unlockStderrWriter();
    const ttyconf = std.io.tty.detectConfig(.stderr());

    var gpa: std.heap.GeneralPurposeAllocator(.{}) = .init;
    const allocator = gpa.allocator();
    try printElf(
        allocator,
        args.next() orelse "./study-samples/split",
        bw,
        ttyconf,
        .{
            // .show_unaddressable_sections = true,
            // .skip_sections_content = true,
        },
    );
}

pub fn printElf(
    allocator: std.mem.Allocator,
    path: []const u8,
    bw: *std.Io.Writer,
    ttyconf: std.io.tty.Config,
    options: struct {
        show_unaddressable_sections: bool = false,
        skip_sections_content: bool = false,
    },
) !void {
    const f = try std.fs.cwd().openFile(path, .{ .mode = .read_only });
    var buffer = try allocator.alignedAlloc(u8, std.mem.Alignment.of(u64), 1024 * 100);
    // const buffer = try allocator.alloc(u8, 1024 * 10000000);
    var reader = f.reader(buffer);
    const header = try std.elf.Header.read(&reader.interface);

    var handle: usize = undefined;
    std.debug.assert(cs.cs_open(cs.CS_ARCH_X86, cs.CS_MODE_64, @ptrCast(&handle)) == cs.CS_ERR_OK);

    const shstrtab = blk: {
        var section_it = header.iterateSectionHeaders(&reader);
        var section_idx: u32 = 0;
        while (try section_it.next()) |s| {
            defer section_idx += 1;
            if (section_idx == header.shstrndx) {
                std.debug.assert(s.sh_type == std.elf.SHT_STRTAB);
                break :blk s;
            }
        }
        break :blk null;
    };

    const elf_shstrtab_slice = blk: {
        if (shstrtab == null)
            break :blk null;

        try reader.seekTo(shstrtab.?.sh_offset);
        const slice = try reader.interface.readAlloc(allocator, shstrtab.?.sh_size);
        break :blk slice;
    };

    const strtab = blk: {
        if (elf_shstrtab_slice == null)
            break :blk null;
        var section_it = header.iterateSectionHeaders(&reader);
        while (try section_it.next()) |s| {
            if (s.sh_type == std.elf.SHT_STRTAB and std.mem.eql(
                u8,
                ".strtab",
                std.mem.sliceTo(elf_shstrtab_slice.?[s.sh_name..], 0),
            ))
                // if (s.sh_type == std.elf.SHT_STRTAB and s.sh_name != shstrtab.?.sh_name and s.sh_addr == 0)
                break :blk s;
        }
        break :blk null;
    };

    const elf_strtab_slice = blk: {
        if (strtab == null)
            break :blk null;

        try reader.seekTo(strtab.?.sh_offset);
        const slice = try reader.interface.readAlloc(allocator, strtab.?.sh_size);
        break :blk slice;
    };

    var strs: std.ArrayList([]const u8) = try .initCapacity(allocator, 8);
    {
        if (elf_shstrtab_slice != null) {
            var str_it = std.mem.splitScalar(u8, elf_shstrtab_slice.?, 0);
            while (str_it.next()) |str| {
                const owned_str = try allocator.alloc(u8, str.len);
                @memcpy(owned_str, str);
                try strs.append(allocator, owned_str);
            }
        }
    }

    var sections: std.ArrayList(std.elf.Elf64_Shdr) = try .initCapacity(allocator, 8);
    {
        var section_it = header.iterateSectionHeaders(&reader);
        while (try section_it.next()) |section| {
            try sections.append(allocator, section);
        }
        std.mem.sort(std.elf.Elf64_Shdr, sections.items, {}, struct {
            pub fn inner(_: void, x: std.elf.Elf64_Shdr, y: std.elf.Elf64_Shdr) bool {
                return x.sh_addr < y.sh_addr;
            }
        }.inner);
    }

    const symtab = blk: {
        var section_it = header.iterateSectionHeaders(&reader);
        while (try section_it.next()) |s| {
            if (s.sh_type == std.elf.SHT_SYMTAB) {
                try bw.print("sym: {any}\n", .{s});
                break :blk s;
            }
        }
        break :blk null;
    };

    const dynsym = blk: {
        var section_it = header.iterateSectionHeaders(&reader);
        while (try section_it.next()) |s| {
            if (s.sh_type == std.elf.SHT_DYNSYM) {
                try bw.print("sym: {any}\n", .{s});
                break :blk s;
            }
        }
        break :blk null;
    };
    try bw.print("dynsym: {any}\n", .{dynsym});

    const symbols_index = blk: {
        var syms: std.ArrayList(SymbolRange) = try .initCapacity(allocator, 8);
        if (symtab != null) {
            var sym_it = iterateSymbols(header, &reader, symtab.?);
            while (try sym_it.next()) |s| {
                const t = s.st_info & 0xf;
                const name = std.mem.sliceTo(elf_strtab_slice.?[s.st_name..], 0);
                const owned_name = try allocator.alloc(u8, name.len);
                @memcpy(owned_name, name);
                try syms.append(allocator, .{
                    .start = s.st_value,
                    .end = s.st_value + s.st_size,
                    .name = owned_name,
                    .kind = t,
                });
            }
        }

        // the check on elf_strtab_slice might not be necessary
        if (dynsym != null and elf_strtab_slice != null) {
            var sym_it = iterateSymbols(header, &reader, dynsym.?);
            while (try sym_it.next()) |s| {
                const t = s.st_info & 0xf;
                const name = std.mem.sliceTo(elf_strtab_slice.?[s.st_name..], 0);
                const owned_name = try allocator.alloc(u8, name.len);
                @memcpy(owned_name, name);
                try syms.append(allocator, .{
                    .start = s.st_value,
                    .end = s.st_value + s.st_size,
                    .name = owned_name,
                    .kind = t,
                });
            }
        }
        std.mem.sort(SymbolRange, syms.items, {}, struct {
            fn inner(_: void, x: SymbolRange, y: SymbolRange) bool {
                return x.start < y.start;
            }
        }.inner);
        break :blk syms.items;
    };

    for (symbols_index) |sym| {
        if (sym.kind == std.elf.STT_FUNC and sym.name.len > 0)
            try bw.print("{s} {x}-{x}\n", .{ sym.name, sym.start, sym.end });
    }

    for (sections.items) |section| {
        if (section.sh_size > 0 and section.sh_addr > 0) {
            try ttyconf.setColor(bw, .reset);
            try ttyconf.setColor(bw, .dim);
            try bw.print("\n{x}-{x} (t: {x})  --  ", .{
                section.sh_addr,
                section.sh_addr + section.sh_size,
                section.sh_type,
            });
            try ttyconf.setColor(bw, .bright_green);
            if (elf_shstrtab_slice != null)
                try bw.print("{s}", .{std.mem.sliceTo(elf_shstrtab_slice.?[section.sh_name..], 0)});
            try bw.print("\n", .{});
            try ttyconf.setColor(bw, .reset);

            // --
            try reader.seekTo(section.sh_offset);

            if (buffer.len < section.sh_size) {
                buffer = try allocator.realloc(buffer, section.sh_size);
                reader = f.reader(buffer);
            }
            const section_slice = reader.interface.take(section.sh_size) catch |e| blk: {
                switch (e) {
                    error.EndOfStream => {
                        try bw.print("failed\n", .{});
                        break :blk null;
                    },
                    error.ReadFailed => unreachable,
                }
            };
            // TODO: this heuristic is probably wrong
            if (section_slice != null and !options.skip_sections_content) {
                if (section.sh_type == std.elf.SHT_PROGBITS and (section.sh_flags & (std.elf.SHF_ALLOC | std.elf.SHF_EXECINSTR)) != 0) {
                    const instrs: []cs.cs_insn = blk: {
                        var insn: [*]cs.cs_insn = undefined;
                        const count = cs.cs_disasm(handle, section_slice.?.ptr, section_slice.?.len, section.sh_addr, 0, @ptrCast(&insn));
                        break :blk insn[0..count];
                    };

                    try dumpInstr(allocator, bw, ttyconf, instrs, symbols_index);
                } else {
                    try dumpHexFallible(u64, bw, ttyconf, section_slice.?, section.sh_addr);
                }
            }
        }
    }

    if (options.show_unaddressable_sections) {
        for (sections.items) |section| {
            if (section.sh_size > 0 and section.sh_addr == 0) {
                try ttyconf.setColor(bw, .reset);
                try ttyconf.setColor(bw, .dim);
                try bw.print("{x}-{x} (t: {x})  --  ", .{
                    section.sh_addr,
                    section.sh_addr + section.sh_size,
                    section.sh_type,
                });
                try ttyconf.setColor(bw, .bright_cyan);
                if (elf_shstrtab_slice != null)
                    try bw.print("{s}", .{std.mem.sliceTo(elf_shstrtab_slice.?[section.sh_name..], 0)});
                try bw.print("\n", .{});
                try ttyconf.setColor(bw, .reset);
                // --

                try reader.seekTo(section.sh_offset);

                if (buffer.len < section.sh_size) {
                    buffer = try allocator.realloc(buffer, section.sh_size);
                    reader = f.reader(buffer);
                }
                const section_slice = reader.interface.take(section.sh_size) catch |e| blk: {
                    switch (e) {
                        error.EndOfStream => {
                            break :blk null;
                        },
                        error.ReadFailed => unreachable,
                    }
                };
                if (section_slice != null and !options.skip_sections_content) {
                    try dumpHexFallible(u64, bw, ttyconf, section_slice.?, section.sh_addr);
                }
            }
        }
    }
}

fn allocComment(
    gpa: std.mem.Allocator,
    code: []u8,
    symbols: []SymbolRange,
) !?[]u8 {
    var iter = std.mem.splitAny(u8, code, " \t[],+-");
    while (iter.next()) |s| {
        if (std.mem.startsWith(u8, s, "0x")) {
            // todo split at the zero char at the end of string
            const v = std.fmt.parseInt(u64, std.mem.sliceTo(s[2..], 0), 16) catch |e| blk: {
                std.debug.print("{any}\n", .{e});
                std.debug.dumpHex(s);
                break :blk 0;
            };
            if (v > 0) {
                const idx = std.sort.lowerBound(SymbolRange, symbols, v, struct {
                    fn inner(a: u64, sym: SymbolRange) std.math.Order {
                        return std.math.order(a, sym.start);
                    }
                }.inner);
                if (idx < symbols.len and v >= symbols[idx].start and v <= symbols[idx].end) {
                    // if (idx > 0)
                    //     idx -= 1;
                    const d = v - symbols[idx].start;
                    if (d > 0)
                        return try std.fmt.allocPrint(gpa, "{s}+0x{x}", .{ symbols[idx].name, d });
                    return try std.fmt.allocPrint(gpa, "{s}", .{symbols[idx].name});
                }
            }
        }
    }

    return null;
}

fn dumpInstr(
    gpa: std.mem.Allocator,
    bw: *std.Io.Writer,
    ttyconf: std.io.tty.Config,
    instrs: []cs.cs_insn,
    symbols: []SymbolRange,
) !void {
    for (instrs) |instr| {
        const addr = instr.address;
        const idx = std.sort.lowerBound(SymbolRange, symbols, addr, struct {
            fn inner(a: u64, sym: SymbolRange) std.math.Order {
                return std.math.order(a, sym.start);
            }
        }.inner);

        if (idx < symbols.len and symbols[idx].start == addr and symbols[idx].name.len > 0) {
            try ttyconf.setColor(bw, .blue);
            try bw.print("\n{x:0>16}  {s}:\n", .{
                addr,
                symbols[idx].name,
            });
            try ttyconf.setColor(bw, .reset);
        }
        try ttyconf.setColor(bw, .dim);
        try bw.print("{x:0>[1]}  ", .{ addr, @sizeOf(usize) * 2 });
        try ttyconf.setColor(bw, .reset);
        // if(instr.detail.)
        try ttyconf.setColor(bw, .bright_green);
        try bw.print("{s} ", .{instr.mnemonic});
        try ttyconf.setColor(bw, .reset);
        try bw.print("{s}", .{instr.op_str});
        const asm_comment = try allocComment(gpa, @ptrCast(@constCast(&instr.op_str)), symbols);
        if (asm_comment != null and asm_comment.?.len > 0) {
            try ttyconf.setColor(bw, .blue);
            try bw.print(" <{s}>", .{asm_comment.?});
        }
        try bw.print("\n", .{});
        try ttyconf.setColor(bw, .reset);
    }
}

/// Prints a hexadecimal view of the bytes, returning any error that occurs.
pub fn dumpHexFallible(_: type, bw: *std.Io.Writer, ttyconf: std.io.tty.Config, bytes: []const u8, offset: u64) !void {
    // @breakpoint();
    const nbytes = 16;
    var chunks = std.mem.window(u8, @ptrCast(@alignCast(bytes)), nbytes, nbytes);
    while (chunks.next()) |window| {
        // 1. Print the address.
        const address = ((0x10 * (std.math.divCeil(usize, chunks.index orelse bytes.len, nbytes) catch unreachable)) - 0x10) + offset;
        try ttyconf.setColor(bw, .dim);
        // We print the address in lowercase and the bytes in uppercase hexadecimal to distinguish them more.
        // Also, make sure all lines are aligned by padding the address.
        try bw.print("{x:0>[1]}  ", .{ address, @sizeOf(usize) * 2 });
        try ttyconf.setColor(bw, .reset);

        // 2. Print the bytes.
        for (window, 0..) |byte, index| {
            try bw.print("{X:0>2} ", .{byte});
            if (index == 7) try bw.writeByte(' ');
        }
        try bw.writeByte(' ');
        if (window.len < 16) {
            var missing_columns = (16 - window.len) * 3;
            if (window.len < 8) missing_columns += 1;
            try bw.splatByteAll(' ', missing_columns);
        }

        const window_bytes: []const u8 = @ptrCast(@alignCast(window));

        // 3. Print the characters.
        for (window_bytes) |byte| {
            if (std.ascii.isPrint(byte)) {
                try bw.writeByte(byte);
            } else {
                // Related: https://github.com/ziglang/zig/issues/7600
                if (ttyconf == .windows_api) {
                    try bw.writeByte('.');
                    continue;
                }

                // Let's print some common control codes as graphical Unicode symbols.
                // We don't want to do this for all control codes because most control codes apart from
                // the ones that Zig has escape sequences for are likely not very useful to print as symbols.
                switch (byte) {
                    '\n' => try bw.writeAll("␊"),
                    '\r' => try bw.writeAll("␍"),
                    '\t' => try bw.writeAll("␉"),
                    else => try bw.writeByte('.'),
                }
            }
        }
        try bw.writeByte('\n');
    }
}