//! Script to generate ksyms
//! Usage: ksyms.zig <output> <input>
const std = @import("std");
const builtin = @import("builtin");
const elf = std.elf;
const Io = std.Io;
const Allocator = std.mem.Allocator;

const Symbol = struct {
    addr: usize,
    name: []const u8,
};

const file_hdr =
    \\/* Auto-generated by build/ksyms.zig - do not edit */
    \\.section .rodata
    \\.global ksyms_array
    \\.global ksyms_count
    \\.align 8
    \\
;

const empty_str = file_hdr ++
    \\ksyms_array:
    \\ksyms_count:
    \\.quad 0
;

fn getElfString(strtab: []const u8, offset: u32) []const u8 {
    const remainder = strtab[offset..];
    const str_ptr: [*:0]const u8 = @ptrCast(remainder.ptr);
    const len = std.mem.len(str_ptr);
    return remainder[0..len];
}

fn getSymbols(allocator: Allocator, bytes: []u8) !std.ArrayList(Symbol) {
    var reader = Io.Reader.fixed(bytes);

    const header = try std.elf.Header.read(&reader);

    const shstrhdr_addr = header.shoff + (header.shentsize * header.shstrndx);

    const shstrhdr = std.mem.bytesAsValue(elf.Shdr, bytes[shstrhdr_addr..]);
    const section_names: []const u8 = bytes[shstrhdr.sh_offset..];

    reader.seek = header.shoff;

    var text_section_index: usize = undefined;
    var symtab: elf.Shdr = undefined;

    const shdrs = std.mem.bytesAsSlice(
        elf.Shdr,
        bytes[header.shoff..][0 .. header.shnum * header.shentsize],
    );

    for (0..header.shnum) |i| {
        const name = getElfString(section_names, shdrs[i].sh_name);

        if (std.mem.eql(u8, name, ".text")) {
            text_section_index = i;
        } else if (shdrs[i].sh_type == elf.SHT_SYMTAB) {
            symtab = shdrs[i];
        }
    }

    const sym_strtab_shdr_addr = header.shoff + (header.shentsize * symtab.sh_link);
    const sym_strtab_shdr = std.mem.bytesAsValue(elf.Shdr, bytes[sym_strtab_shdr_addr..]);
    const symbol_names: []const u8 = bytes[sym_strtab_shdr.sh_offset..];

    const how_many = symtab.sh_size / symtab.sh_entsize;
    const symbols: [*]elf.Sym = @ptrCast(@alignCast(bytes[symtab.sh_offset..].ptr));

    var ret = try std.ArrayList(Symbol).initCapacity(allocator, how_many);

    for (0..how_many) |i| {
        const sym = symbols[i];

        if (sym.st_shndx == elf.SHN_UNDEF)
            continue;

        if (sym.st_name == 0)
            continue;

        if (sym.st_shndx != text_section_index) {
            continue;
        }

        const name = getElfString(symbol_names, sym.st_name);

        try ret.append(allocator, .{
            .addr = sym.st_value,
            .name = name,
        });
    }

    return ret;
}

fn lessThan(_: void, a: Symbol, b: Symbol) bool {
    return a.addr < b.addr;
}

fn generateSymbols(alloc: Allocator, file: []u8, writer: *Io.Writer) !void {
    const symbols = try getSymbols(alloc, file);

    std.mem.sort(Symbol, symbols.items, {}, lessThan);

    _ = try writer.write(file_hdr);
    _ = try writer.write("ksyms_array:\n");

    for (symbols.items, 0..) |sym, i| {
        try writer.print(".quad 0x{x}\n", .{sym.addr});
        try writer.print(".quad .Lstr{}\n", .{i});
        try writer.print(".quad {}\n", .{sym.name.len});
    }

    try writer.print("ksyms_count:\n .quad {}\n", .{symbols.items.len});

    for (symbols.items, 0..) |sym, i| {
        if (std.mem.indexOfScalar(u8, sym.name, '"') == null) {
            try writer.print(".Lstr{}: .asciz \"{s}\"\n", .{ i, sym.name });
        } else {
            const size = std.mem.replacementSize(u8, sym.name, "\"", "\\\"");
            const escaped = try alloc.alloc(u8, size);

            _ = std.mem.replace(u8, sym.name, "\"", "\\\"", escaped);
            try writer.print(".Lstr{}: .asciz \"{s}\" \n", .{ i, escaped });
        }
    }
}

pub fn main(init: std.process.Init) !void {
    const arena = init.arena.allocator();
    const io = init.io;
    const args = init.minimal.args;
    const args_slice = try args.toSlice(arena);

    const outfile = try std.Io.Dir.cwd().createFile(io, args_slice[1], .{});
    defer outfile.close(io);

    const write_buffer = try arena.alloc(u8, 64 * 1024);
    var writer = outfile.writer(io, write_buffer);

    if (args_slice.len == 2) {
        // Only output, generate an empty file.
        try writer.interface.writeAll(empty_str);
    }

    if (args_slice.len == 3) {
        // Got an input, parse the ELF.
        const buf = try std.Io.Dir.cwd().readFileAlloc(
            io,
            args_slice[2],
            arena,
            .unlimited,
        );
        try generateSymbols(arena, buf, &writer.interface);
    }

    try writer.flush();
}