const std = @import("std");
const elfo = @import("elfo");

/// Build-system test that parses a controlled C binary and asserts correct
/// symbol resolution — static symbols, dynamic symbols, and PLT entries.
pub fn main(init: std.process.Init.Minimal) !void {
    var debug_alloc: std.heap.DebugAllocator(.{}) = .init;
    const gpa = debug_alloc.allocator();

    var threaded: std.Io.Threaded = .init(gpa, .{
        .argv0 = .init(init.args),
        .environ = init.environ,
    });
    defer threaded.deinit();
    const io = threaded.io();

    var args = init.args.iterate();
    _ = args.skip(); // argv[0]
    const path = args.next() orelse return error.MissingArg;

    const buffer = try gpa.alignedAlloc(u8, std.mem.Alignment.of(u64), 1024 * 100);
    defer gpa.free(buffer);

    const f = try std.Io.Dir.cwd().openFile(io, path, .{ .mode = .read_only });
    var reader = f.reader(io, buffer);
    const header = try std.elf.Header.read(&reader.interface);

    var sections = try elfo.collectSections(header, &reader, gpa);
    defer sections.deinit(gpa);

    const shstrtab_data: ?[]u8 = if (sections.shstrtab) |s| blk: {
        try reader.seekTo(s.sh_offset);
        break :blk try reader.interface.readAlloc(gpa, s.sh_size);
    } else null;
    defer if (shstrtab_data) |d| gpa.free(d);

    const strtab_data: ?[]u8 = if (sections.strtab) |s| blk: {
        try reader.seekTo(s.sh_offset);
        break :blk try reader.interface.readAlloc(gpa, s.sh_size);
    } else null;
    defer if (strtab_data) |d| gpa.free(d);

    const dynstr_data: ?[]u8 = if (sections.dynstr) |s| blk: {
        try reader.seekTo(s.sh_offset);
        break :blk try reader.interface.readAlloc(gpa, s.sh_size);
    } else null;
    defer if (dynstr_data) |d| gpa.free(d);

    var symbols = try elfo.collectSymbols(gpa, header, &reader, strtab_data, dynstr_data, shstrtab_data, sections);
    defer {
        for (symbols.items) |sym| gpa.free(sym.name);
        symbols.deinit(gpa);
    }

    // --- Assertions ---

    // 1. Key sections must exist
    try assert(sections.symtab != null, ".symtab must exist");
    try assert(sections.dynsym != null, ".dynsym must exist");
    try assert(sections.strtab != null, ".strtab must exist (via sh_link)");
    try assert(sections.dynstr != null, ".dynstr must exist (via sh_link)");

    // 2. Static function symbols from .symtab
    try assertSymbol(symbols.items, "main", .{ .is_func = true, .nonzero_addr = true });
    try assertSymbol(symbols.items, "a", .{ .is_func = true, .nonzero_addr = true });
    try assertSymbol(symbols.items, "bcd", .{ .is_func = true, .nonzero_addr = true });

    // 3. Dynamic symbols from .dynsym (via .dynstr - clean names without version)
    try assertSymbol(symbols.items, "printf", .{});
    try assertSymbol(symbols.items, "malloc", .{});
    try assertSymbol(symbols.items, "free", .{});

    // 4. PLT symbols (synthesized from .rela.plt)
    try assertSymbol(symbols.items, "printf@plt", .{ .is_func = true, .nonzero_addr = true });
    try assertSymbol(symbols.items, "malloc@plt", .{ .is_func = true, .nonzero_addr = true });
    try assertSymbol(symbols.items, "free@plt", .{ .is_func = true, .nonzero_addr = true });

    // 5. strtab and dynstr must be different sections (sh_link resolved correctly)
    if (sections.strtab) |st| {
        if (sections.dynstr) |ds| {
            try assert(st.sh_offset != ds.sh_offset, ".strtab and .dynstr must be different sections");
        }
    }

    // success — no output means no warnings from the build system
}

const AssertOpts = struct {
    is_func: bool = false,
    nonzero_addr: bool = false,
};

fn assertSymbol(symbols: []const elfo.SymbolRange, name: []const u8, opts: AssertOpts) !void {
    for (symbols) |sym| {
        if (std.mem.eql(u8, sym.name, name)) {
            if (opts.is_func and sym.kind != std.elf.STT_FUNC) {
                std.debug.print("FAIL: '{s}' expected STT_FUNC, got {d}\n", .{ name, sym.kind });
                return error.TestFailed;
            }
            if (opts.nonzero_addr and sym.start == 0) {
                std.debug.print("FAIL: '{s}' expected nonzero address\n", .{name});
                return error.TestFailed;
            }
            return; // found
        }
    }
    std.debug.print("FAIL: symbol '{s}' not found. available symbols:\n", .{name});
    for (symbols) |sym| {
        std.debug.print("  {x}-{x} kind={d} '{s}'\n", .{ sym.start, sym.end, sym.kind, sym.name });
    }
    return error.SymbolNotFound;
}

fn assert(condition: bool, msg: []const u8) !void {
    if (!condition) {
        std.debug.print("FAIL: {s}\n", .{msg});
        return error.TestFailed;
    }
}