const std = @import("std");
const uefi = std.os.uefi;
const elf = std.elf;
const danos = @import("danos");
const BootInfo = danos.BootInfo;
const GraphicsOutput = uefi.protocol.GraphicsOutput;
const EdidActive = uefi.protocol.edid.Active;
const MemoryMapSlice = uefi.tables.MemoryMapSlice;

/// Name of the kernel ELF on the boot volume (installed to the ESP root by
/// build.zig). UEFI wants a UTF-16, null-terminated path.
const kernel_file_name = std.unicode.utf8ToUtf16LeStringLiteral("danos");

/// Physical page size, and the sentinel UEFI uses to seek to end-of-file.
const page_size = 4096;
const seek_end = 0xffff_ffff_ffff_ffff;

pub fn main() uefi.Status {
    // `boot` never returns on success — it jumps into the kernel. If it fails,
    // report the reason (boot services are still up) and park the machine so the
    // message stays on screen.
    boot() catch |err| {
        log("\r\ndanos: boot failed: ");
        logBytes(@errorName(err));
        log("\r\n");
        while (true) asm volatile ("hlt");
    };
    unreachable;
}

fn boot() !noreturn {
    const bs = uefi.system_table.boot_services orelse return error.NoBootServices;

    // Everything the kernel needs must be gathered *before* we exit boot
    // services, since afterwards none of these calls are usable.
    var boot_info: BootInfo = .{
        .framebuffer = try queryFramebuffer(bs),
        .memory_map = undefined, // filled by exitBootServices, just below
        .kernel_segments = undefined, // filled by loadKernel
        .kernel_segment_count = 0,
    };

    const entry = try loadKernel(bs, &boot_info);

    log("danos: kernel loaded, exiting boot services\r\n");
    boot_info.memory_map = try exitBootServices(bs);

    // Hand control to the kernel. `danos.kernel_abi` is SysV, so the pointer is
    // passed in RDI as the kernel expects — not RCX, which this UEFI binary's
    // default `.c` convention (Microsoft x64) would use.
    const kernel: *const fn (*const BootInfo) callconv(danos.kernel_abi) noreturn = @ptrFromInt(entry);
    kernel(&boot_info);
}

/// A display resolution in pixels.
const Resolution = struct { width: u32, height: u32 };

/// Switch the GPU to the monitor's native resolution (when we can determine it)
/// and read the resulting graphics mode into our own framebuffer description.
fn queryFramebuffer(bs: *uefi.tables.BootServices) !danos.Framebuffer {
    // Enumerate the handles carrying the Graphics Output Protocol. We go through
    // handles (rather than locateProtocol) so we can also ask them for their EDID,
    // which is what tells us the panel's native resolution.
    const handles = (try bs.locateHandleBuffer(.{ .by_protocol = &GraphicsOutput.guid })) orelse
        return error.NoGraphicsOutput;
    defer _ = bs.freePool(@ptrCast(handles.ptr)) catch {};

    const gop = (try bs.handleProtocol(GraphicsOutput, handles[0])) orelse
        return error.NoGraphicsOutput;

    // Best effort: the monitor's preferred (native) timing from its EDID.
    const native = nativeResolution(bs, handles);

    // Select the mode and switch to it if it isn't already current. setMode
    // updates gop.mode (info and frame_buffer_base) to describe the new mode.
    const target = pickMode(gop, native);
    if (target != gop.mode.mode) try gop.setMode(target);

    const info = gop.mode.info;
    return .{
        .base = @intCast(gop.mode.frame_buffer_base),
        .width = info.horizontal_resolution,
        .height = info.vertical_resolution,
        // Each pixel is 32 bits, so the byte pitch is 4 * pixels-per-row.
        .pitch = info.pixels_per_scan_line * 4,
        .format = try pixelFormat(info.pixel_format),
    };
}

/// Map a GOP pixel format to ours. bit_mask / blt_only have no linear 32bpp
/// layout we can paint into, so they're rejected.
fn pixelFormat(fmt: GraphicsOutput.PixelFormat) !danos.PixelFormat {
    return switch (fmt) {
        .red_green_blue_reserved_8_bit_per_color => .rgbx,
        .blue_green_red_reserved_8_bit_per_color => .bgrx,
        else => error.UnsupportedPixelFormat,
    };
}

/// Choose the graphics mode to boot with. If we learned the monitor's native
/// resolution from EDID and a mode offers it (with a layout we can paint into),
/// use that. Otherwise keep whatever mode the firmware already selected: with a
/// valid EDID present the firmware normally defaults to the native mode itself,
/// so its default is a far safer bet than second-guessing it with, say, the
/// largest advertised mode (which is often a huge non-native surface).
fn pickMode(gop: *GraphicsOutput, native: ?Resolution) u32 {
    const n = native orelse return gop.mode.mode;
    var id: u32 = 0;
    while (id < gop.mode.max_mode) : (id += 1) {
        const info = gop.queryMode(id) catch continue;
        _ = pixelFormat(info.pixel_format) catch continue; // must be paintable
        if (info.horizontal_resolution == n.width and
            info.vertical_resolution == n.height) return id;
    }
    return gop.mode.mode; // native not on offer; trust the firmware's default
}

/// The monitor's native resolution, read from an EDID's preferred timing. We try
/// every GOP handle and both EDID protocols (Active first, then Discovered),
/// since firmware installs them inconsistently — and many, including OVMF with
/// QEMU's stdvga, don't expose them at all. Returns null when none is found, in
/// which case pickMode keeps the firmware's default mode.
fn nativeResolution(bs: *uefi.tables.BootServices, handles: []uefi.Handle) ?Resolution {
    for (handles) |h| {
        if (bs.handleProtocol(EdidActive, h) catch null) |e| {
            if (e.edid) |p| if (edidNative(p[0..e.size_of_edid])) |r| return r;
        }
        if (bs.handleProtocol(uefi.protocol.edid.Discovered, h) catch null) |e| {
            if (e.edid) |p| if (edidNative(p[0..e.size_of_edid])) |r| return r;
        }
    }
    return null;
}

/// Parse the native resolution from a raw EDID block. The first Detailed Timing
/// Descriptor (at byte 54) is the preferred — i.e. native — mode by convention;
/// its active pixel counts are split across low bytes and the high nibbles of
/// later bytes.
fn edidNative(edid: []const u8) ?Resolution {
    if (edid.len < 128) return null;
    // Every EDID begins with this fixed 8-byte header.
    const header = [_]u8{ 0x00, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x00 };
    if (!std.mem.eql(u8, edid[0..8], &header)) return null;

    const dtd = edid[54..][0..18];
    // A zero pixel clock marks a display (not timing) descriptor: no resolution.
    if (dtd[0] == 0 and dtd[1] == 0) return null;

    const w = @as(u32, dtd[2]) | (@as(u32, dtd[4] & 0xf0) << 4);
    const h = @as(u32, dtd[5]) | (@as(u32, dtd[7] & 0xf0) << 4);
    if (w == 0 or h == 0) return null;
    return .{ .width = w, .height = h };
}

/// Open the kernel on the volume we booted from, read it into a pool buffer,
/// load its segments, and return the physical entry-point address.
fn loadKernel(bs: *uefi.tables.BootServices, boot_info: *BootInfo) !usize {
    const loaded = (try bs.handleProtocol(uefi.protocol.LoadedImage, uefi.handle)) orelse
        return error.NoLoadedImage;
    const device = loaded.device_handle orelse return error.NoBootDevice;
    const fs = (try bs.handleProtocol(uefi.protocol.SimpleFileSystem, device)) orelse
        return error.NoFileSystem;

    const root = try fs.openVolume();
    defer _ = root.close() catch {};

    const file = try root.open(kernel_file_name, .read, .{});
    defer _ = file.close() catch {};

    // Seek to the end to learn the size, then rewind.
    try file.setPosition(seek_end);
    const size: usize = @intCast(try file.getPosition());
    try file.setPosition(0);

    const image = try bs.allocatePool(.loader_data, size);
    defer _ = bs.freePool(image.ptr) catch {};

    // `read` may return short; loop until the whole file is in memory.
    var read_total: usize = 0;
    while (read_total < size) {
        const n = try file.read(image[read_total..]);
        if (n == 0) return error.UnexpectedEof;
        read_total += n;
    }

    return loadElf(bs, image, boot_info);
}

/// Validate the ELF, copy every PT_LOAD segment to its physical address, and
/// record each segment's layout so the kernel can re-map itself with the right
/// permissions.
fn loadElf(bs: *uefi.tables.BootServices, image: []u8, boot_info: *BootInfo) !usize {
    if (image.len < @sizeOf(elf.Elf64_Ehdr)) return error.NotElf;
    const ehdr: *const elf.Elf64_Ehdr = @ptrCast(@alignCast(image.ptr));

    if (ehdr.e_ident[0] != 0x7f or ehdr.e_ident[1] != 'E' or
        ehdr.e_ident[2] != 'L' or ehdr.e_ident[3] != 'F') return error.NotElf;
    if (ehdr.e_machine != .X86_64) return error.WrongArchitecture;

    var i: usize = 0;
    while (i < ehdr.e_phnum) : (i += 1) {
        const phdr: *const elf.Elf64_Phdr = @ptrCast(@alignCast(
            image.ptr + ehdr.e_phoff + i * ehdr.e_phentsize,
        ));
        if (phdr.p_type != elf.PT_LOAD) continue;

        // Reserve the exact physical pages this segment is linked at. This
        // requires the segment's p_paddr to be free in the firmware memory map;
        // if it collides, adjust `image_base` in build.zig.
        const mem_sz: usize = @intCast(phdr.p_memsz);
        const pages = (mem_sz + page_size - 1) / page_size;
        const dest: [*]align(page_size) uefi.Page = @ptrFromInt(phdr.p_paddr);
        _ = try bs.allocatePages(.{ .address = dest }, .loader_data, pages);

        // Copy the file-backed part, then zero the .bss tail (memsz > filesz).
        const bytes: [*]u8 = @ptrFromInt(phdr.p_paddr);
        const file_sz: usize = @intCast(phdr.p_filesz);
        const off: usize = @intCast(phdr.p_offset);
        @memcpy(bytes[0..file_sz], image[off..][0..file_sz]);
        @memset(bytes[file_sz..mem_sz], 0);

        // Record it (identity-loaded: virtual == physical) for the kernel's VMM.
        const n = boot_info.kernel_segment_count;
        if (n < boot_info.kernel_segments.len) {
            boot_info.kernel_segments[n] = .{
                .virt = phdr.p_vaddr,
                .pages = pages,
                .flags = phdr.p_flags,
            };
            boot_info.kernel_segment_count = n + 1;
        }
    }

    return @intCast(ehdr.e_entry);
}

/// Fetch the memory map, exit boot services, and hand back the map in danos's
/// neutral form. Allocating the buffers can itself change the map (invalidating
/// the key), so retry until it takes. Both buffers are LoaderData, which survives
/// the exit, so the returned map stays valid for the kernel.
fn exitBootServices(bs: *uefi.tables.BootServices) !danos.MemoryMap {
    var attempts: usize = 0;
    while (attempts < 8) : (attempts += 1) {
        const info = try bs.getMemoryMapInfo();
        // Spare descriptors to absorb the growth from the allocations below.
        const cap = info.len + 8;
        const map_buf = try bs.allocatePool(.loader_data, cap * info.descriptor_size);
        const regions_buf = try bs.allocatePool(.loader_data, cap * @sizeOf(danos.MemoryRegion));
        const map = bs.getMemoryMap(map_buf) catch {
            _ = bs.freePool(map_buf.ptr) catch {};
            _ = bs.freePool(regions_buf.ptr) catch {};
            continue;
        };
        bs.exitBootServices(uefi.handle, map.info.key) catch {
            _ = bs.freePool(map_buf.ptr) catch {};
            _ = bs.freePool(regions_buf.ptr) catch {};
            continue;
        };
        // Boot services are gone; do not touch `bs` again. Converting the map is
        // pure computation on memory we already hold, so it's safe here.
        return convertMemoryMap(map, regions_buf);
    }
    return error.ExitBootServicesFailed;
}

/// Translate UEFI's memory map into danos's neutral `MemoryRegion` array, written
/// into `out` (sized for at least `map.info.len` regions). Adjacent regions of
/// the same kind are coalesced. This is the loader's job precisely so the kernel
/// never sees UEFI's vocabulary — the same seam the framebuffer already uses.
fn convertMemoryMap(map: MemoryMapSlice, out: []u8) danos.MemoryMap {
    const regions: [*]danos.MemoryRegion = @ptrCast(@alignCast(out.ptr));
    var count: usize = 0;
    var i: usize = 0;
    while (i < map.info.len) : (i += 1) {
        // Stride by descriptor_size, NOT @sizeOf — firmware descriptors may be
        // larger than the struct.
        const d: *const uefi.tables.MemoryDescriptor =
            @ptrCast(@alignCast(map.ptr + i * map.info.descriptor_size));
        if (d.number_of_pages == 0) continue;
        const kind = classify(d);

        // Coalesce with the previous region if it's the same kind and contiguous.
        if (count > 0) {
            const prev = &regions[count - 1];
            if (prev.kind == kind and
                prev.base + prev.pages * danos.page_size == d.physical_start)
            {
                prev.pages += d.number_of_pages;
                continue;
            }
        }
        regions[count] = .{
            .base = d.physical_start,
            .pages = d.number_of_pages,
            .kind = kind,
        };
        count += 1;
    }
    return .{ .regions = @intFromPtr(regions), .len = count };
}

/// Map a UEFI descriptor to danos's neutral kind. A region that isn't
/// writeback-cacheable (`wb`) isn't backed by real RAM — it's device registers or
/// a reserved address-space window (e.g. PCIe config space) — so it's `mmio`
/// regardless of type. UEFI overloads `reserved_memory_type` for both reserved RAM
/// and such holes, and the cache attribute is what actually tells them apart.
/// Among RAM regions, anything we don't recognise is `reserved` — the safe
/// default; our own LoaderData (kernel image, these buffers) lands there too and
/// stays reserved until the kernel reclaims it.
fn classify(d: *const uefi.tables.MemoryDescriptor) danos.MemoryKind {
    if (!d.attribute.wb) return .mmio;
    return switch (d.@"type") {
        .conventional_memory => .usable,
        .boot_services_code, .boot_services_data => .reclaimable,
        .acpi_reclaim_memory => .acpi_tables,
        .acpi_memory_nvs => .acpi_nvs,
        .memory_mapped_io, .memory_mapped_io_port_space => .mmio,
        else => .reserved,
    };
}

/// Write a compile-time string to the console (best effort).
fn log(comptime msg: []const u8) void {
    const out = uefi.system_table.con_out orelse return;
    _ = out.outputString(std.unicode.utf8ToUtf16LeStringLiteral(msg)) catch {};
}

/// Write a runtime ASCII byte string (e.g. an @errorName) by widening to UTF-16.
fn logBytes(bytes: []const u8) void {
    const out = uefi.system_table.con_out orelse return;
    var buf: [128]u16 = undefined;
    var i: usize = 0;
    for (bytes) |b| {
        if (i + 1 >= buf.len) break;
        buf[i] = b;
        i += 1;
    }
    buf[i] = 0;
    _ = out.outputString(buf[0..i :0].ptr) catch {};
}