const builtin = @import("builtin");
const std = @import("std");
const Allocator = std.mem.Allocator;
const assert = std.debug.assert;

pub fn main() !void {
    var arena_instance = std.heap.ArenaAllocator.init(std.heap.page_allocator);
    defer arena_instance.deinit();
    const arena = arena_instance.allocator();

    var general_purpose_allocator: std.heap.GeneralPurposeAllocator(.{}) = .{};
    const gpa = general_purpose_allocator.allocator();

    const args = try std.process.argsAlloc(arena);
    const zig_lib_directory = args[1];
    const zig_exe_path = args[2];
    const global_cache_path = args[3];

    var lib_dir = try std.fs.cwd().openDir(zig_lib_directory, .{});
    defer lib_dir.close();

    const listen_port: u16 = 0;
    const address = std.net.Address.parseIp("127.0.0.1", listen_port) catch unreachable;
    var http_server = try address.listen(.{});
    const port = http_server.listen_address.in.getPort();
    const url = try std.fmt.allocPrint(arena, "http://127.0.0.1:{d}/\n", .{port});
    std.io.getStdOut().writeAll(url) catch {};
    openBrowserTab(gpa, url[0 .. url.len - 1 :'\n']) catch |err| {
        std.log.err("unable to open browser: {s}", .{@errorName(err)});
    };

    var context: Context = .{
        .gpa = gpa,
        .zig_exe_path = zig_exe_path,
        .global_cache_path = global_cache_path,
        .lib_dir = lib_dir,
        .zig_lib_directory = zig_lib_directory,
    };

    while (true) {
        const connection = try http_server.accept();
        _ = std.Thread.spawn(.{}, accept, .{ &context, connection }) catch |err| {
            std.log.err("unable to accept connection: {s}", .{@errorName(err)});
            connection.stream.close();
            continue;
        };
    }
}

fn accept(context: *Context, connection: std.net.Server.Connection) void {
    defer connection.stream.close();

    var read_buffer: [8000]u8 = undefined;
    var server = std.http.Server.init(connection, &read_buffer);
    while (server.state == .ready) {
        var request = server.receiveHead() catch |err| switch (err) {
            error.HttpConnectionClosing => return,
            else => {
                std.log.err("closing http connection: {s}", .{@errorName(err)});
                return;
            },
        };
        serveRequest(&request, context) catch |err| {
            std.log.err("unable to serve {s}: {s}", .{ request.head.target, @errorName(err) });
            return;
        };
    }
}

const Context = struct {
    gpa: Allocator,
    lib_dir: std.fs.Dir,
    zig_lib_directory: []const u8,
    zig_exe_path: []const u8,
    global_cache_path: []const u8,
};

fn serveRequest(request: *std.http.Server.Request, context: *Context) !void {
    if (std.mem.eql(u8, request.head.target, "/") or
        std.mem.eql(u8, request.head.target, "/debug/"))
    {
        try serveDocsFile(request, context, "docs/index.html", "text/html");
    } else if (std.mem.eql(u8, request.head.target, "/main.js") or
        std.mem.eql(u8, request.head.target, "/debug/main.js"))
    {
        try serveDocsFile(request, context, "docs/main.js", "application/javascript");
    } else if (std.mem.eql(u8, request.head.target, "/main.wasm")) {
        try serveWasm(request, context, .ReleaseFast);
    } else if (std.mem.eql(u8, request.head.target, "/debug/main.wasm")) {
        try serveWasm(request, context, .Debug);
    } else if (std.mem.eql(u8, request.head.target, "/sources.tar") or
        std.mem.eql(u8, request.head.target, "/debug/sources.tar"))
    {
        try serveSourcesTar(request, context);
    } else {
        try request.respond("not found", .{
            .status = .not_found,
            .extra_headers = &.{
                .{ .name = "content-type", .value = "text/plain" },
            },
        });
    }
}

const cache_control_header: std.http.Header = .{
    .name = "cache-control",
    .value = "max-age=0, must-revalidate",
};

fn serveDocsFile(
    request: *std.http.Server.Request,
    context: *Context,
    name: []const u8,
    content_type: []const u8,
) !void {
    const gpa = context.gpa;
    // The desired API is actually sendfile, which will require enhancing std.http.Server.
    // We load the file with every request so that the user can make changes to the file
    // and refresh the HTML page without restarting this server.
    const file_contents = try context.lib_dir.readFileAlloc(gpa, name, 10 * 1024 * 1024);
    defer gpa.free(file_contents);
    try request.respond(file_contents, .{
        .extra_headers = &.{
            .{ .name = "content-type", .value = content_type },
            cache_control_header,
        },
    });
}

fn serveSourcesTar(request: *std.http.Server.Request, context: *Context) !void {
    const gpa = context.gpa;

    var send_buffer: [0x4000]u8 = undefined;
    var response = request.respondStreaming(.{
        .send_buffer = &send_buffer,
        .respond_options = .{
            .extra_headers = &.{
                .{ .name = "content-type", .value = "application/x-tar" },
                cache_control_header,
            },
        },
    });
    const w = response.writer();

    var std_dir = try context.lib_dir.openDir("std", .{ .iterate = true });
    defer std_dir.close();

    var walker = try std_dir.walk(gpa);
    defer walker.deinit();

    while (try walker.next()) |entry| {
        switch (entry.kind) {
            .file => {
                if (!std.mem.endsWith(u8, entry.basename, ".zig"))
                    continue;
                if (std.mem.endsWith(u8, entry.basename, "test.zig"))
                    continue;
            },
            else => continue,
        }

        var file = try std_dir.openFile(entry.path, .{});
        defer file.close();

        const stat = try file.stat();
        const padding = p: {
            const remainder = stat.size % 512;
            break :p if (remainder > 0) 512 - remainder else 0;
        };
        comptime assert(@sizeOf(TarHeader) == 512);

        var file_header = TarHeader.init();
        file_header.typeflag = .regular;
        try file_header.setPath("std", entry.path);
        try file_header.setSize(stat.size);
        try file_header.updateChecksum();
        try w.writeAll(std.mem.asBytes(&file_header));
        try w.writeFile(file);
        try w.writeByteNTimes(0, padding);
    }
    // intentionally omitting the pointless trailer
    //try w.writeByteNTimes(0, 512 * 2);
    try response.end();
}

fn serveWasm(
    request: *std.http.Server.Request,
    context: *Context,
    optimize_mode: std.builtin.OptimizeMode,
) !void {
    const gpa = context.gpa;

    var arena_instance = std.heap.ArenaAllocator.init(gpa);
    defer arena_instance.deinit();
    const arena = arena_instance.allocator();

    // Do the compilation every request, so that the user can edit the files
    // and see the changes without restarting the server.
    const wasm_binary_path = try buildWasmBinary(arena, context, optimize_mode);
    // std.http.Server does not have a sendfile API yet.
    const file_contents = try std.fs.cwd().readFileAlloc(gpa, wasm_binary_path, 10 * 1024 * 1024);
    defer gpa.free(file_contents);
    try request.respond(file_contents, .{
        .extra_headers = &.{
            .{ .name = "content-type", .value = "application/wasm" },
            cache_control_header,
        },
    });
}

fn buildWasmBinary(
    arena: Allocator,
    context: *Context,
    optimize_mode: std.builtin.OptimizeMode,
) ![]const u8 {
    const gpa = context.gpa;

    const main_src_path = try std.fs.path.join(arena, &.{
        context.zig_lib_directory, "docs", "wasm", "main.zig",
    });

    var argv: std.ArrayListUnmanaged([]const u8) = .{};

    try argv.appendSlice(arena, &.{
        context.zig_exe_path,
        "build-exe",
        "-fno-entry",
        "-O",
        @tagName(optimize_mode),
        "-target",
        "wasm32-freestanding",
        "-mcpu",
        "baseline+atomics+bulk_memory+multivalue+mutable_globals+nontrapping_fptoint+reference_types+sign_ext",
        "--cache-dir",
        context.global_cache_path,
        "--global-cache-dir",
        context.global_cache_path,
        "--name",
        "autodoc",
        "-rdynamic",
        main_src_path,
        "--listen=-",
    });

    var child = std.ChildProcess.init(argv.items, gpa);
    child.stdin_behavior = .Pipe;
    child.stdout_behavior = .Pipe;
    child.stderr_behavior = .Pipe;
    try child.spawn();

    var poller = std.io.poll(gpa, enum { stdout, stderr }, .{
        .stdout = child.stdout.?,
        .stderr = child.stderr.?,
    });
    defer poller.deinit();

    try sendMessage(child.stdin.?, .update);
    try sendMessage(child.stdin.?, .exit);

    const Header = std.zig.Server.Message.Header;
    var result: ?[]const u8 = null;
    var result_error_bundle = std.zig.ErrorBundle.empty;

    const stdout = poller.fifo(.stdout);

    poll: while (true) {
        while (stdout.readableLength() < @sizeOf(Header)) {
            if (!(try poller.poll())) break :poll;
        }
        const header = stdout.reader().readStruct(Header) catch unreachable;
        while (stdout.readableLength() < header.bytes_len) {
            if (!(try poller.poll())) break :poll;
        }
        const body = stdout.readableSliceOfLen(header.bytes_len);

        switch (header.tag) {
            .zig_version => {
                if (!std.mem.eql(u8, builtin.zig_version_string, body)) {
                    return error.ZigProtocolVersionMismatch;
                }
            },
            .error_bundle => {
                const EbHdr = std.zig.Server.Message.ErrorBundle;
                const eb_hdr = @as(*align(1) const EbHdr, @ptrCast(body));
                const extra_bytes =
                    body[@sizeOf(EbHdr)..][0 .. @sizeOf(u32) * eb_hdr.extra_len];
                const string_bytes =
                    body[@sizeOf(EbHdr) + extra_bytes.len ..][0..eb_hdr.string_bytes_len];
                // TODO: use @ptrCast when the compiler supports it
                const unaligned_extra = std.mem.bytesAsSlice(u32, extra_bytes);
                const extra_array = try arena.alloc(u32, unaligned_extra.len);
                @memcpy(extra_array, unaligned_extra);
                result_error_bundle = .{
                    .string_bytes = try arena.dupe(u8, string_bytes),
                    .extra = extra_array,
                };
            },
            .emit_bin_path => {
                const EbpHdr = std.zig.Server.Message.EmitBinPath;
                const ebp_hdr = @as(*align(1) const EbpHdr, @ptrCast(body));
                if (!ebp_hdr.flags.cache_hit) {
                    std.log.info("source changes detected; rebuilt wasm component", .{});
                }
                result = try arena.dupe(u8, body[@sizeOf(EbpHdr)..]);
            },
            else => {}, // ignore other messages
        }

        stdout.discard(body.len);
    }

    const stderr = poller.fifo(.stderr);
    if (stderr.readableLength() > 0) {
        const owned_stderr = try stderr.toOwnedSlice();
        defer gpa.free(owned_stderr);
        std.debug.print("{s}", .{owned_stderr});
    }

    // Send EOF to stdin.
    child.stdin.?.close();
    child.stdin = null;

    switch (try child.wait()) {
        .Exited => |code| {
            if (code != 0) {
                std.log.err(
                    "the following command exited with error code {d}:\n{s}",
                    .{ code, try std.Build.Step.allocPrintCmd(arena, null, argv.items) },
                );
                return error.WasmCompilationFailed;
            }
        },
        .Signal, .Stopped, .Unknown => {
            std.log.err(
                "the following command terminated unexpectedly:\n{s}",
                .{try std.Build.Step.allocPrintCmd(arena, null, argv.items)},
            );
            return error.WasmCompilationFailed;
        },
    }

    if (result_error_bundle.errorMessageCount() > 0) {
        const color = std.zig.Color.auto;
        result_error_bundle.renderToStdErr(color.renderOptions());
        std.log.err("the following command failed with {d} compilation errors:\n{s}", .{
            result_error_bundle.errorMessageCount(),
            try std.Build.Step.allocPrintCmd(arena, null, argv.items),
        });
        return error.WasmCompilationFailed;
    }

    return result orelse {
        std.log.err("child process failed to report result\n{s}", .{
            try std.Build.Step.allocPrintCmd(arena, null, argv.items),
        });
        return error.WasmCompilationFailed;
    };
}

fn sendMessage(file: std.fs.File, tag: std.zig.Client.Message.Tag) !void {
    const header: std.zig.Client.Message.Header = .{
        .tag = tag,
        .bytes_len = 0,
    };
    try file.writeAll(std.mem.asBytes(&header));
}

fn openBrowserTab(gpa: Allocator, url: []const u8) !void {
    // Until https://github.com/ziglang/zig/issues/19205 is implemented, we
    // spawn a thread for this child process.
    _ = try std.Thread.spawn(.{}, openBrowserTabThread, .{ gpa, url });
}

fn openBrowserTabThread(gpa: Allocator, url: []const u8) !void {
    const main_exe = switch (builtin.os.tag) {
        .windows => "explorer",
        else => "xdg-open",
    };
    var child = std.ChildProcess.init(&.{ main_exe, url }, gpa);
    child.stdin_behavior = .Ignore;
    child.stdout_behavior = .Ignore;
    child.stderr_behavior = .Ignore;
    try child.spawn();
    _ = try child.wait();
}

/// Forked from https://github.com/mattnite/tar/blob/main/src/main.zig which is
/// MIT licensed.
pub const TarHeader = extern struct {
    name: [100]u8,
    mode: [7:0]u8,
    uid: [7:0]u8,
    gid: [7:0]u8,
    size: [11:0]u8,
    mtime: [11:0]u8,
    checksum: [7:0]u8,
    typeflag: FileType,
    linkname: [100]u8,
    magic: [5:0]u8,
    version: [2]u8,
    uname: [31:0]u8,
    gname: [31:0]u8,
    devmajor: [7:0]u8,
    devminor: [7:0]u8,
    prefix: [155]u8,
    pad: [12]u8,

    const FileType = enum(u8) {
        regular = '0',
        hard_link = '1',
        symbolic_link = '2',
        character = '3',
        block = '4',
        directory = '5',
        fifo = '6',
        reserved = '7',
        pax_global = 'g',
        extended = 'x',
        _,
    };

    fn init() TarHeader {
        var ret = std.mem.zeroes(TarHeader);
        ret.magic = [_:0]u8{ 'u', 's', 't', 'a', 'r' };
        ret.version = [_:0]u8{ '0', '0' };
        return ret;
    }

    fn setPath(self: *TarHeader, prefix: []const u8, path: []const u8) !void {
        if (prefix.len + 1 + path.len > 100) {
            var i: usize = 0;
            while (i < path.len and path.len - i > 100) {
                while (path[i] != '/') : (i += 1) {}
            }

            _ = try std.fmt.bufPrint(&self.prefix, "{s}/{s}", .{ prefix, path[0..i] });
            _ = try std.fmt.bufPrint(&self.name, "{s}", .{path[i + 1 ..]});
        } else {
            _ = try std.fmt.bufPrint(&self.name, "{s}/{s}", .{ prefix, path });
        }
    }

    fn setSize(self: *TarHeader, size: u64) !void {
        _ = try std.fmt.bufPrint(&self.size, "{o:0>11}", .{size});
    }

    fn updateChecksum(self: *TarHeader) !void {
        const offset = @offsetOf(TarHeader, "checksum");
        var checksum: usize = 0;
        for (std.mem.asBytes(self), 0..) |val, i| {
            checksum += if (i >= offset and i < offset + @sizeOf(@TypeOf(self.checksum)))
                ' '
            else
                val;
        }

        _ = try std.fmt.bufPrint(&self.checksum, "{o:0>7}", .{checksum});
    }
};