const std = @import("std");
const mime = @import("mime/mime.zig");

const Address = std.net.Address;
const ArrayList = std.ArrayList;
const debug = std.debug;
const fmt = std.fmt;
const heap = std.heap;
const http = std.http;
const fs = std.fs;
const log = std.log.scoped(.server);
const mem = std.mem;
const bufferedReader = std.io.bufferedReader;
const getcwd = std.os.getcwd;

const server_addr = "127.0.0.1";
const server_port = 8080;
const MAX_PATH_BYTES = fs.MAX_PATH_BYTES;
const BUFFER_LIMIT = 1 << 21;

fn read_files(target: []const u8, buffer: []u8, allocator: mem.Allocator) !usize {
    var file_path = try allocator.alloc(u8, MAX_PATH_BYTES);
    defer allocator.free(file_path);
    var cwd_buffer = [_]u8{0} ** MAX_PATH_BYTES;
    const cwd = try getcwd(&cwd_buffer);
    file_path = try fmt.allocPrint(allocator, "{s}{s}", .{ cwd, target });
    log.info("Loading file {s}...", .{file_path});
    // Determine if the requested file exists
    const file = try fs.cwd().openFile(file_path, .{});
    defer file.close();
    var stat = try file.stat();
    switch (stat.kind) {
        .directory => {
            if (mem.endsWith(u8, target, "/")) {
                file_path = try fmt.allocPrint(allocator, "{s}index.html", .{target});
            } else {
                file_path = try fmt.allocPrint(allocator, "{s}/index.html", .{target});
            }
            return read_files(file_path, buffer, allocator);
        },
        .file => {
            return try file.readAll(buffer);
        },
        else => {
            return 0;
        },
    }
}

fn handle_request(response: *http.Server.Response, allocator: mem.Allocator) !void {
    // Log the request details
    log.info("{s} {s} {s}", .{ @tagName(response.request.method), @tagName(response.request.version), response.request.target });

    // Create a []u8 to read up to BUFFER_LIMIT characters
    var read = [_]u8{0} ** BUFFER_LIMIT;

    // Set "connection" header to "keep-alive" if present in request headers
    if (response.request.headers.contains("connection")) {
        try response.headers.append("connection", "keep-alive");
    }

    const size = read_files(response.request.target, &read, allocator) catch |err| {
        switch (err) {
            error.AccessDenied => {
                response.status = .forbidden;
            },
            error.FileNotFound => {
                response.status = .not_found;
            },
            error.OutOfMemory => {
                response.status = .uri_too_long;
            },
            else => {
                return err;
            },
        }
        response.transfer_encoding = .{ .content_length = 0 };
        try response.do();
        return;
    };
    if (size >= BUFFER_LIMIT) {
        response.status = .payload_too_large;
        response.transfer_encoding = .{ .content_length = 0 };
        try response.do();
        return;
    } else if (size > 0) {
        // Get the file extension, and set the content-type header if it exists
        // (using mime.zig)
        // To do this, we iterate through response.request.target in reverse
        // looking for a '/' or a '.'
        // a '/' indicates a directory, so there is no extension
        // a '.' indicates a file extension
        var i: usize = response.request.target.len;
        var mime_type: ?mime.Type = undefined;
        if (mem.endsWith(u8, response.request.target, "/")) {
            log.warn("Dir requested, returning index.html!", .{});
            try response.headers.append("content-type", "text/html");
        } else {
            while (i > 0) {
                i -= 1;
                switch (response.request.target[i]) {
                    '/' => {
                        i = 0;
                        break;
                    },
                    '.' => {
                        break;
                    },
                    else => {
                        continue;
                    },
                }
            }
            if (i <= 0) {
                log.warn("No extension detected!", .{});
                if (mem.indexOf(u8, &read, "<html")) |_| {
                    try response.headers.append("content-type", "text/html");
                } else {
                    try response.headers.append("content-type", "text/plain");
                }
            } else {
                log.info("Extension {s} detected!", .{response.request.target[i..]});
                mime_type = mime.extension_map.get(response.request.target[i..]);
                if (mime_type) |mime_val| {
                    try response.headers.append("content-type", @tagName(mime_val));
                } else {
                    try response.headers.append("content-type", "text/plain");
                }
            }
        }
        // Create a []u8 that is exactly the intended size and no larger
        const body = read[0..size];
        log.info("{}", .{body.len});
        // Check if the request target contains "?chunked"
        if (mem.indexOf(u8, response.request.target, "?chunked") != null) {
            response.transfer_encoding = .chunked;
        } else {
            response.transfer_encoding = .{ .content_length = size };
        }
        // Transmit the response
        try response.do();
        if (response.request.method != .HEAD) {
            try response.writeAll(body);
            try response.finish();
        }
    } else {
        // If the file was not found, return error 404
        response.status = .not_found;
        response.transfer_encoding = .{ .content_length = 0 };
        try response.do();
    }
}

fn run_server(server: *http.Server, allocator: mem.Allocator) !void {
    outer: while (true) {
        // Accept incoming connection
        var response = try server.accept(.{
            .allocator = allocator,
        });
        defer response.deinit();

        while (response.reset() != .closing) {
            // Handle errors during request processing
            response.wait() catch |err| switch (err) {
                error.HttpHeadersInvalid => continue :outer,
                error.EndOfStream => continue,
                else => return err,
            };

            // Process the request
            handle_request(&response, allocator) catch |err| {
                response.status = .internal_server_error;
                response.transfer_encoding = .{ .content_length = 0 };
                try response.do();
                return err;
            };
        }
    }
}

pub fn main() !void {
    // Define allocator
    var gpa = heap.GeneralPurposeAllocator(.{}){};
    defer debug.assert(gpa.deinit() == .ok);
    const allocator = gpa.allocator();

    // Initialize the server
    var server = http.Server.init(allocator, .{ .reuse_address = true });
    defer server.deinit();

    // Log the server address and port
    log.info("Server is running at {s}:{d}", .{ server_addr, server_port });

    // Parse the server address
    const address = Address.parseIp(server_addr, server_port) catch unreachable;
    try server.listen(address);

    // Run the server
    run_server(&server, allocator) catch |err| {
        // Handle server errors
        log.err("server error: {}\n", .{err});
        if (@errorReturnTrace()) |trace| {
            debug.dumpStackTrace(trace.*);
        }
        std.os.exit(1);
    };
}