//! Minimal ELF64 executable writer for the `.rgx` native backend.
//!
//! Emits a GNU/Linux ET_EXEC (x86_64 or AArch64) that makes only raw kernel
//! syscalls (no libc). The caller supplies encoded machine-code bytes for
//! `.text` (beginning with `_start`) plus optional read-only data (string
//! literals / format tables). Everything lives in a single `RWX` PT_LOAD
//! segment so the image is self-contained and directly executable with no
//! relocations.
//!
//! The writer is restricted to the very narrow x86_64 / AArch64 Linux ABI
//! targets: it bails out with `error.UnsupportedArch` anywhere else. Machine
//! selection follows `builtin.cpu.arch` (the build host).
//!
//! This module owns only file-format assembly (headers + layout). All
//! instruction encoding lives in the value `.rgx` lowering.
const std = @import("std");
const builtin = @import("builtin");

pub const Error = error{ OutOfMemory, UnsupportedArch };

pub const ELF_HEADER_SIZE = 64;
pub const PHDR_SIZE = 56;
pub const PHDR_COUNT = 1; // a single RWX PT_LOAD; a second degenerate (empty) LOAD
// crashes the aarch64 kernel's elf_map, so it was dropped.

/// Where `.text` is placed in the address space (page-aligned; both the
/// x86-64 SysV PIE origin and a plain AArch64 static origin work here).
pub const LOAD_BASE: u64 = 0x400000;

/// ELF `e_machine` for the build-host architecture.
pub fn machineType() u16 {
    return switch (builtin.cpu.arch) {
        .x86_64 => 0x3e, // EM_X86_64
        .aarch64 => 183, // EM_AARCH64
        else => unreachable, // assemble() already rejected other arches
    };
}

/// Describes the pieces of one emitted image.
pub const Image = struct {
    /// Encoded instructions starting at the `_start` entry.
    text: []const u8,
    /// Read-only data appended (16-aligned) after the code; RIP-relative
    /// references from `text` are patched by the caller using `rodata_addr`.
    rodata: []const u8 = &.{},

    /// Filled by `assemble`: absolute virtual address of `.text`.
    text_addr: u64 = 0,
    /// Filled by `assemble`: absolute virtual address of `.rodata`.
    rodata_addr: u64 = 0,
};

const PT_LOAD = 1;
const PF_X: u32 = 1;
const PF_W: u32 = 2;
const PF_R: u32 = 4;

/// Result of `layout`: where each piece lands in the final image.
pub const Layout = struct {
    text_off: usize,
    rodata_off: usize,
    total: usize,
};

/// Computes the file offsets and absolute virtual addresses for `img` without
/// producing bytes. The `.rgx` lowering uses this to patch RIP-relative
/// displacements (which depend on the final rodata address) before assembling.
///
/// Fills `img.text_addr` and `img.rodata_addr`.
pub fn layout(img: *Image) Layout {
    const text_off = ELF_HEADER_SIZE + PHDR_COUNT * PHDR_SIZE;
    const rodata_off = alignA(text_off + img.text.len, 16);
    img.text_addr = LOAD_BASE + @as(u64, text_off);
    img.rodata_addr = LOAD_BASE + @as(u64, rodata_off);
    return .{
        .text_off = text_off,
        .rodata_off = rodata_off,
        .total = rodata_off + img.rodata.len,
    };
}

/// Builds the complete ELF byte image. Caller owns the returned block.
pub fn assemble(alloc: std.mem.Allocator, img: *Image) Error![]u8 {
    if ((builtin.cpu.arch != .x86_64 and builtin.cpu.arch != .aarch64) or builtin.os.tag != .linux) {
        return error.UnsupportedArch;
    }

    const lay = layout(img);
    const out = try alloc.alloc(u8, lay.total);
    @memset(out, 0);

    writeEhdr(out, img.text_addr);
    writePhdr(out, 0, .{
        .p_offset = 0,
        .p_vaddr = LOAD_BASE,
        .p_filesz = @as(u64, lay.total),
        .p_memsz = @as(u64, lay.total),
        .p_flags = PF_R | PF_W | PF_X,
    });

    @memcpy(out[lay.text_off .. lay.text_off + img.text.len], img.text);
    @memcpy(out[lay.rodata_off .. lay.rodata_off + img.rodata.len], img.rodata);

    return out;
}

fn putU16(out: []u8, off: usize, val: u16) void {
    const b = std.mem.toBytes(val);
    @memcpy(out[off .. off + 2], &b);
}

fn putU32(out: []u8, off: usize, val: u32) void {
    const b = std.mem.toBytes(val);
    @memcpy(out[off .. off + 4], &b);
}

fn putU64(out: []u8, off: usize, val: u64) void {
    const b = std.mem.toBytes(val);
    @memcpy(out[off .. off + 8], &b);
}

fn getU16(in: []const u8, off: usize) u16 {
    var b: [2]u8 = undefined;
    @memcpy(&b, in[off .. off + 2]);
    return std.mem.readInt(u16, &b, .little);
}

fn getU32(in: []const u8, off: usize) u32 {
    var b: [4]u8 = undefined;
    @memcpy(&b, in[off .. off + 4]);
    return std.mem.readInt(u32, &b, .little);
}

fn getU64(in: []const u8, off: usize) u64 {
    var b: [8]u8 = undefined;
    @memcpy(&b, in[off .. off + 8]);
    return std.mem.readInt(u64, &b, .little);
}

fn alignA(v: usize, a: usize) usize {
    return (v + a - 1) & ~(a - 1);
}

fn writeEhdr(out: []u8, entry: u64) void {
    out[0] = 0x7f;
    out[1] = 'E';
    out[2] = 'L';
    out[3] = 'F';
    out[4] = 2; // 64-bit
    out[5] = 1; // little-endian
    out[6] = 1; // ELF version
    putU16(out, 16, 2); // ET_EXEC
    putU16(out, 18, machineType()); // e_machine (host arch)
    putU32(out, 20, 1); // e_version
    putU64(out, 24, entry); // e_entry
    putU64(out, 32, ELF_HEADER_SIZE); // e_phoff
    putU64(out, 40, 0); // e_shoff
    putU32(out, 48, 0); // e_flags
    putU16(out, 52, ELF_HEADER_SIZE); // e_ehsize
    putU16(out, 54, PHDR_SIZE); // e_phentsize
    putU16(out, 56, PHDR_COUNT); // e_phnum
    putU16(out, 58, 0); // e_shentsize
    putU16(out, 60, 0); // e_shnum
    putU16(out, 62, 0); // e_shstrndx
}

fn writePhdr(out: []u8, i: usize, p: struct {
    p_offset: u64,
    p_vaddr: u64,
    p_filesz: u64,
    p_memsz: u64,
    p_flags: u32,
}) void {
    const h = ELF_HEADER_SIZE + i * PHDR_SIZE;
    putU32(out, h, PT_LOAD);
    putU32(out, h + 4, p.p_flags);
    putU64(out, h + 8, p.p_offset);
    putU64(out, h + 16, p.p_vaddr);
    putU64(out, h + 24, p.p_vaddr); // p_paddr
    putU64(out, h + 32, p.p_filesz);
    putU64(out, h + 40, p.p_memsz);
    putU64(out, h + 48, 0x1000); // p_align
}

const testing = std.testing;

test "elf: headers and single load segment are structurally sound" {
    // A trivial `_start` that immediately exits(42), for the host arch.
    const text = switch (builtin.cpu.arch) {
        .x86_64 => [_]u8{ 0x48, 0xc7, 0xc0, 0x3c, 0x00, 0x00, 0x00, 0x48, 0xc7, 0xc7, 0x2a, 0x00, 0x00, 0x00, 0x0f, 0x05 },
        .aarch64 => [_]u8{ 0x40, 0x05, 0x80, 0xd2, 0xa8, 0x0b, 0x80, 0xd2, 0x01, 0x00, 0x00, 0xd4 },
        else => return error.SkipZigTest,
    };
    var img = Image{ .text = &text };
    const blob = try assemble(testing.allocator, &img);
    defer testing.allocator.free(blob);

    try testing.expect(std.mem.eql(u8, blob[0..4], &.{ 0x7f, 'E', 'L', 'F' }));
    try testing.expectEqual(@as(u16, 2), getU16(blob, 16)); // ET_EXEC
    try testing.expectEqual(machineType(), getU16(blob, 18)); // host arch machine type
    try testing.expectEqual(img.text_addr, getU64(blob, 24)); // e_entry
    const phoff = getU64(blob, 32);
    try testing.expectEqual(@as(u16, 1), getU16(blob, 56)); // e_phnum
    try testing.expectEqual(@as(u32, PT_LOAD), getU32(blob, phoff));
    // Code bytes are present in the image after the ELF+phdr headers.
    const text_off = ELF_HEADER_SIZE + PHDR_COUNT * PHDR_SIZE;
    try testing.expect(std.mem.eql(u8, blob[text_off .. text_off + text.len], &text));
}

test "elf: emitted image is a runnable no-libc executable (exit 42)" {
    if ((builtin.cpu.arch != .x86_64 and builtin.cpu.arch != .aarch64) or builtin.os.tag != .linux)
        return error.SkipZigTest;

    // `_start` exits(42) with raw syscalls for the host arch.
    const text = switch (builtin.cpu.arch) {
        // mov rax,60; mov rdi,42; syscall
        .x86_64 => [_]u8{ 0x48, 0xc7, 0xc0, 0x3c, 0x00, 0x00, 0x00, 0x48, 0xc7, 0xc7, 0x2a, 0x00, 0x00, 0x00, 0x0f, 0x05 },
        // movz x0,#42; movz x8,#93; svc #0
        .aarch64 => [_]u8{ 0x40, 0x05, 0x80, 0xd2, 0xa8, 0x0b, 0x80, 0xd2, 0x01, 0x00, 0x00, 0xd4 },
        else => return error.SkipZigTest,
    };
    var img = Image{ .text = &text };
    const blob = try assemble(testing.allocator, &img);
    defer testing.allocator.free(blob);

    const path = "/tmp/rg_elf_exit42_rinnescript";
    const io = std.testing.io;
    const cwd = std.Io.Dir.cwd();
    cwd.writeFile(io, .{ .sub_path = path, .data = blob }) catch return error.TestUnexpectedResult;
    defer {
        cwd.deleteFile(io, path) catch {};
    }

    const path_z = try testing.allocator.dupeZ(u8, path);
    defer testing.allocator.free(path_z);
    try testing.expect(std.os.linux.chmod(path_z, 0o755) == 0);

    var child = std.process.spawn(io, .{ .argv = &.{ path }, .stdout = .ignore, .stderr = .ignore }) catch return error.TestUnexpectedResult;
    const term = child.wait(io) catch return error.TestUnexpectedResult;
    switch (term) {
        .exited => |code| try testing.expectEqual(@as(u8, 42), code),
        else => return error.TestUnexpectedResult,
    }
}