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

pub fn main(init: std.process.Init) !void {
    const allocator = init.gpa;
    const io = init.io;

    const cmd = init.environ_map.get("SSH_ORIGINAL_COMMAND") orelse
        return launchTui(allocator, io, init.environ_map, &.{});
    const tokens = shlex.split(allocator, cmd, false, true) catch |err| {
        std.log.err("failed to parse SSH_ORIGINAL_COMMAND: {}", .{err});
        std.process.exit(1);
    };
    defer {
        for (tokens) |token| allocator.free(token);
        allocator.free(tokens);
    }

    if (tokens.len == 0) return launchTui(allocator, io, init.environ_map, tokens);

    const is_receive_pack = std.mem.eql(u8, tokens[0], "git-receive-pack");
    const is_upload_pack = std.mem.eql(u8, tokens[0], "git-upload-pack");
    if (!is_receive_pack and !is_upload_pack) {
        return launchTui(allocator, io, init.environ_map, tokens);
    }
    if (tokens.len != 2) {
        std.log.err("invalid git command: {s}", .{cmd});
        std.process.exit(1);
    }

    const home = init.environ_map.get("HOME") orelse {
        std.log.err("HOME not set", .{});
        std.process.exit(1);
    };
    const repo_path = try expandPath(allocator, home, tokens[1]);
    defer allocator.free(repo_path);

    if (!isUnder(repo_path, home)) {
        std.log.err("repo path must be inside HOME: {s}", .{tokens[1]});
        std.process.exit(1);
    }

    if (is_upload_pack) {
        return std.process.replace(io, .{
            .argv = &.{ "git-upload-pack", repo_path },
        });
    }

    var created_repo = false;
    if (is_receive_pack) {
        std.Io.Dir.accessAbsolute(io, repo_path, .{}) catch {
            std.log.info("auto-initializing bare repo at {s}", .{repo_path});
            std.Io.Dir.createDirPath(.cwd(), io, repo_path) catch |err| {
                std.log.err("failed to create repo directory: {}", .{err});
                std.process.exit(1);
            };
            const result = try std.process.run(allocator, io, .{
                .argv = &.{ "git", "init", "--bare", repo_path },
            });
            defer allocator.free(result.stdout);
            defer allocator.free(result.stderr);
            if (result.term != .exited or result.term.exited != 0) {
                std.log.err("git init --bare failed: {s}", .{result.stderr});
                std.process.exit(1);
            }
            created_repo = true;
        };
    }

    var receive_pack = try std.process.spawn(io, .{
        .argv = &.{ "git-receive-pack", repo_path },
    });
    const term = try receive_pack.wait(io);
    if (term != .exited) std.process.exit(1);
    if (term.exited != 0) std.process.exit(term.exited);

    if (created_repo) try fixHead(allocator, io, repo_path);
    if (is_receive_pack) refreshAfterPush(allocator, io, init.environ_map, repo_path);
}

fn launchTui(
    allocator: std.mem.Allocator,
    io: std.Io,
    environ: *const std.process.Environ.Map,
    tokens: []const []const u8,
) !void {
    const tui = environ.get("SORCERY_SSH_TUI") orelse "sorcery-ssh-tui";
    const argv = try allocator.alloc([]const u8, tokens.len + 1);
    defer allocator.free(argv);
    argv[0] = tui;
    for (tokens, argv[1..]) |token, *arg| arg.* = token;
    const err = std.process.replace(io, .{ .argv = argv });
    return switch (err) {
        error.FileNotFound => printWelcome(io, environ),
        else => err,
    };
}

fn printWelcome(io: std.Io, environ: *const std.process.Environ.Map) !void {
    var buffer: [1024]u8 = undefined;
    var stdout = std.Io.File.stdout().writer(io, &buffer);
    try writeWelcome(&stdout.interface, environ.get("SORCERY_INSTANCE_NAME") orelse "this Sorcery instance");
    try stdout.interface.flush();
}

fn writeWelcome(writer: *std.Io.Writer, instance: []const u8) !void {
    try writer.print(
        \\welcome to sorcery-ssh on {s}!
        \\
        \\usage:
        \\  ssh git@{s} describe USER/REPOSITORY
        \\  ssh git@{s} describe USER/REPOSITORY DESCRIPTION...
        \\
        \\install sorcery-ssh-tui for interactive repository management.
        \\
    , .{ instance, instance, instance });
}

fn refreshAfterPush(
    allocator: std.mem.Allocator,
    io: std.Io,
    environ: *const std.process.Environ.Map,
    repo_path: []const u8,
) void {
    const repositories_raw = environ.get("SORCERY_REPOSITORIES") orelse {
        std.log.warn("SORCERY_REPOSITORIES not set; site refresh skipped", .{});
        return;
    };
    const repositories = std.fs.path.resolve(allocator, &.{repositories_raw}) catch |err| {
        std.log.warn("failed to resolve SORCERY_REPOSITORIES: {}", .{err});
        return;
    };
    defer allocator.free(repositories);

    const coordinates = repositoryCoordinates(repositories, repo_path) orelse return;
    const socket = environ.get("SORCERY_SOCKET") orelse {
        std.log.warn("SORCERY_SOCKET not set; site refresh skipped", .{});
        return;
    };
    const token_file = environ.get("SORCERY_REFRESH_TOKEN_FILE") orelse {
        std.log.warn("SORCERY_REFRESH_TOKEN_FILE not set; site refresh skipped", .{});
        return;
    };
    const token_alloc = std.Io.Dir.readFileAlloc(
        .cwd(),
        io,
        token_file,
        allocator,
        .limited(4096),
    ) catch |err| {
        std.log.warn("failed to read refresh token: {}", .{err});
        return;
    };
    defer allocator.free(token_alloc);
    const token = std.mem.trimEnd(u8, token_alloc, "\r\n");
    if (token.len == 0) {
        std.log.warn("refresh token must contain visible ASCII without whitespace", .{});
        return;
    }
    for (token) |byte| {
        if (!std.ascii.isPrint(byte) or std.ascii.isWhitespace(byte)) {
            std.log.warn("refresh token must contain visible ASCII without whitespace", .{});
            return;
        }
    }

    refresh(io, socket, token, coordinates.user, coordinates.repo) catch |err| {
        std.log.warn("site refresh failed: {}", .{err});
        return;
    };

    if (environ.get("SORCERY_CLONE_URL_BASE")) |base| {
        std.debug.print(
            "sorcery: {s}/{s}/{s}/\n",
            .{ std.mem.trimEnd(u8, base, "/"), coordinates.user, coordinates.repo },
        );
    }
}

const Coordinates = struct {
    user: []const u8,
    repo: []const u8,
};

fn repositoryCoordinates(root: []const u8, path: []const u8) ?Coordinates {
    if (!isUnder(path, root)) return null;
    const offset = if (std.mem.eql(u8, root, &.{std.fs.path.sep})) root.len else root.len + 1;
    var parts = std.mem.splitScalar(u8, path[offset..], std.fs.path.sep);
    const user = parts.next() orelse return null;
    const repo_raw = parts.next() orelse return null;
    if (user.len == 0 or repo_raw.len == 0 or parts.next() != null) return null;
    const repo = if (std.mem.endsWith(u8, repo_raw, ".git"))
        repo_raw[0 .. repo_raw.len - ".git".len]
    else
        repo_raw;
    if (repo.len == 0) return null;
    return .{ .user = user, .repo = repo };
}

/// Whether `path` is strictly below `root`; `root` itself does not count.
fn isUnder(path: []const u8, root: []const u8) bool {
    if (std.mem.eql(u8, root, &.{std.fs.path.sep})) {
        return path.len > 1 and path[0] == std.fs.path.sep;
    }
    return std.mem.startsWith(u8, path, root) and
        path.len > root.len and path[root.len] == std.fs.path.sep;
}

fn refresh(
    io: std.Io,
    socket_path: []const u8,
    token: []const u8,
    user: []const u8,
    repo: []const u8,
) !void {
    const address = try std.Io.net.UnixAddress.init(socket_path);
    const stream = try address.connect(io);
    defer stream.close(io);

    var write_buffer: [1024]u8 = undefined;
    var stream_writer = stream.writer(io, &write_buffer);
    const writer = &stream_writer.interface;
    try writer.writeAll("POST /-/refresh/");
    try writeUrlComponent(writer, user);
    try writer.writeByte('/');
    try writeUrlComponent(writer, repo);
    try writer.writeAll(" HTTP/1.1\r\nHost: sorcery\r\nAuthorization: Bearer ");
    try writer.writeAll(token);
    try writer.writeAll("\r\nContent-Length: 0\r\nConnection: close\r\n\r\n");
    try writer.flush();

    var read_buffer: [1024]u8 = undefined;
    var stream_reader = stream.reader(io, &read_buffer);
    const status = try stream_reader.interface.takeDelimiter('\n') orelse return error.EmptyResponse;
    if (!std.mem.startsWith(u8, status, "HTTP/1.1 200 ") and
        !std.mem.startsWith(u8, status, "HTTP/1.0 200 "))
    {
        return error.RefreshRejected;
    }
}

fn writeUrlComponent(writer: *std.Io.Writer, component: []const u8) !void {
    const hex = "0123456789ABCDEF";
    for (component) |byte| {
        if (std.ascii.isAlphanumeric(byte) or std.mem.indexOfScalar(u8, "-._~", byte) != null) {
            try writer.writeByte(byte);
        } else {
            try writer.writeAll(&.{ '%', hex[byte >> 4], hex[byte & 0xf] });
        }
    }
}

fn fixHead(allocator: std.mem.Allocator, io: std.Io, repo_path: []const u8) !void {
    const head_result = try std.process.run(allocator, io, .{
        .argv = &.{ "git", "--git-dir", repo_path, "rev-parse", "--verify", "--quiet", "HEAD" },
    });
    defer allocator.free(head_result.stdout);
    defer allocator.free(head_result.stderr);
    if (head_result.term == .exited and head_result.term.exited == 0) return;
    if (head_result.term != .exited or head_result.term.exited != 1) {
        std.log.err("failed to resolve repository HEAD: {s}", .{head_result.stderr});
        return error.GitCommandFailed;
    }

    const refs_result = try std.process.run(allocator, io, .{
        .argv = &.{ "git", "--git-dir", repo_path, "for-each-ref", "--format=%(refname)", "refs/heads/" },
    });
    defer allocator.free(refs_result.stdout);
    defer allocator.free(refs_result.stderr);
    if (refs_result.term != .exited or refs_result.term.exited != 0) {
        std.log.err("failed to list repository branches: {s}", .{refs_result.stderr});
        return error.GitCommandFailed;
    }

    var refs = std.mem.tokenizeScalar(u8, refs_result.stdout, '\n');
    var sole_ref: ?[]const u8 = null;
    var main_ref: ?[]const u8 = null;
    var ref_count: usize = 0;
    while (refs.next()) |ref| {
        sole_ref = ref;
        ref_count += 1;
        if (std.mem.eql(u8, ref, "refs/heads/main")) main_ref = ref;
    }

    const head_ref = main_ref orelse if (ref_count == 1) sole_ref else null;
    if (head_ref) |ref| {
        const set_head_result = try std.process.run(allocator, io, .{
            .argv = &.{ "git", "--git-dir", repo_path, "symbolic-ref", "HEAD", ref },
        });
        defer allocator.free(set_head_result.stdout);
        defer allocator.free(set_head_result.stderr);
        if (set_head_result.term != .exited or set_head_result.term.exited != 0) {
            std.log.err("failed to set repository HEAD: {s}", .{set_head_result.stderr});
            return error.GitCommandFailed;
        }
    } else if (ref_count > 1) {
        std.log.warn("HEAD remains unresolved because the repository has multiple branches and no main branch", .{});
    }
}

// im pretty sure this is what git expand user path does. wordexp from posix is overkill
fn expandPath(allocator: std.mem.Allocator, home: []const u8, path: []const u8) ![]const u8 {
    if (std.mem.startsWith(u8, path, "~/")) {
        return std.fs.path.resolve(allocator, &.{ home, path[2..] });
    } else if (std.mem.eql(u8, path, "~")) {
        return allocator.dupe(u8, home);
    } else {
        return std.fs.path.resolve(allocator, &.{ home, path });
    }
}

test "pushed paths normalise into HOME or are rejected" {
    const home = "/home/git";
    for ([_][]const u8{ "x", "~/x", "/home/git/x", "~/a/../x", "/home/git/../git/x" }) |path| {
        const expanded = try expandPath(std.testing.allocator, home, path);
        defer std.testing.allocator.free(expanded);
        try std.testing.expectEqualStrings("/home/git/x", expanded);
        try std.testing.expect(isUnder(expanded, home));
    }
    for ([_][]const u8{ "~", ".", "../x", "~/../x", "/etc/passwd", "/home/gitolite/x", "/home" }) |path| {
        const expanded = try expandPath(std.testing.allocator, home, path);
        defer std.testing.allocator.free(expanded);
        try std.testing.expect(!isUnder(expanded, home));
    }
}

test "refresh url components escape everything outside the unreserved set" {
    var output = std.Io.Writer.Allocating.init(std.testing.allocator);
    defer output.deinit();
    try writeUrlComponent(&output.writer, "a b/c?d#\u{e9}-._~");
    try std.testing.expectEqualStrings("a%20b%2Fc%3Fd%23%C3%A9-._~", output.written());
}

test "repository coordinates require exactly user and repo under root" {
    const expectEqualStrings = std.testing.expectEqualStrings;
    const root = "/home/git/public";

    const plain = repositoryCoordinates(root, "/home/git/public/char/sorcery").?;
    try expectEqualStrings("char", plain.user);
    try expectEqualStrings("sorcery", plain.repo);

    const dotted = repositoryCoordinates(root, "/home/git/public/char/sorcery.git").?;
    try expectEqualStrings("sorcery", dotted.repo);

    try std.testing.expect(repositoryCoordinates(root, "/home/git/private/char/sorcery") == null);
    try std.testing.expect(repositoryCoordinates(root, "/home/git/public/char") == null);
    try std.testing.expect(repositoryCoordinates(root, "/home/git/public/char/sorcery/extra") == null);
    try std.testing.expect(repositoryCoordinates(root, "/home/git/publicity/char/sorcery") == null);
}