#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <limits.h>
#include <errno.h>
#include <stdint.h>
#include <fcntl.h>

#define MAX_CMD_LEN 4096
#define MAX_KEY_LEN 8192
#define GIT_BASE_PATH "/srv/git"
#define AUTH_FILE "/etc/marrow/authorized_keys"

typedef enum {
    CMD_UNKNOWN,
    CMD_GIT_UPLOAD_PACK,
    CMD_GIT_RECEIVE_PACK,
    CMD_GIT_UPLOAD_ARCHIVE,
    CMD_MARROW_KEY
} git_command_t;

static const char *
get_git_user(void) {
    return getenv("GIT_USER");
}

static int32_t
is_anonymous(void) {
    return NULL == get_git_user();
}

static git_command_t
parse_command(const char *cmd) {
    if (0 == strncmp(cmd, "git-upload-pack", 15)) {
        return CMD_GIT_UPLOAD_PACK;
    } else if (0 == strncmp(cmd, "git-receive-pack", 16)) {
        return CMD_GIT_RECEIVE_PACK;
    } else if (0 == strncmp(cmd, "git-upload-archive", 18)) {
        return CMD_GIT_UPLOAD_ARCHIVE;
    } else if (0 == strcmp(cmd, "marrow-key")) {
        return CMD_MARROW_KEY;
    }
    return CMD_UNKNOWN;
}

static int32_t
is_read_command(git_command_t cmd) {
    return CMD_GIT_UPLOAD_PACK == cmd || CMD_GIT_UPLOAD_ARCHIVE == cmd;
}

static int32_t
is_write_command(git_command_t cmd) {
    return CMD_GIT_RECEIVE_PACK == cmd;
}

static char *
extract_repo_path(const char *cmd) {
    char *path = NULL;
    const char *start = strchr(cmd, '\'');
    
    if (start) {
        start++;
        const char *end = strchr(start, '\'');
        if (end) {
            size_t len = end - start;
            path = malloc(len + 1);
            if (path) {
                strncpy(path, start, len);
                path[len] = '\0';
            }
        }
    } else {
        const char *space = strchr(cmd, ' ');
        if (space) {
            space++;
            while (' ' == *space) space++;
            
            size_t len = strlen(space);
            path = malloc(len + 1);
            if (path) {
                strcpy(path, space);
                char *end = strchr(path, ' ');
                if (end) *end = '\0';
            }
        }
    }
    
    return path;
}

static int32_t
validate_repo_path(const char *repo_path, const char *user) {
    char real_path[PATH_MAX];
    char expected_base[PATH_MAX];
    
    if ('/' == repo_path[0]) {
        snprintf(real_path, sizeof(real_path), "%s", repo_path);
    } else {
        snprintf(real_path, sizeof(real_path), "%s/%s", GIT_BASE_PATH, repo_path);
    }
    
    char *resolved = realpath(real_path, NULL);
    if (!resolved) {
        if (ENOENT == errno) {
            resolved = strdup(real_path);
        } else {
            return 0;
        }
    }
    
    if (!user) {
        snprintf(expected_base, sizeof(expected_base), "%s/", GIT_BASE_PATH);
    } else {
        snprintf(expected_base, sizeof(expected_base), "%s/%s/", GIT_BASE_PATH, user);
    }
    
    int32_t valid = (0 == strncmp(resolved, expected_base, strlen(expected_base))) ||
                    (0 == strcmp(resolved, GIT_BASE_PATH) && !user);
    
    free(resolved);
    return valid;
}

static int32_t
check_permissions(const char *cmd) {
    git_command_t git_cmd = parse_command(cmd);
    
    if (CMD_UNKNOWN == git_cmd) {
        fputs("Error: Unknown command\n", stderr);
        return 0;
    }
    
    /* marrow-key command requires authentication */
    if (CMD_MARROW_KEY == git_cmd) {
        if (is_anonymous()) {
            fputs("Error: marrow-key requires authentication\n", stderr);
            return 0;
        }
        return 1;
    }
    
    const char *user = get_git_user();
    char *repo_path = extract_repo_path(cmd);
    
    if (!repo_path) {
        fputs("Error: Could not extract repository path\n", stderr);
        return 0;
    }
    
    if (is_anonymous()) {
        if (is_write_command(git_cmd)) {
            fputs("Error: Anonymous users cannot perform write operations\n", stderr);
            free(repo_path);
            return 0;
        }
    } else {
        if (is_write_command(git_cmd)) {
            if (!validate_repo_path(repo_path, user)) {
                fputs("Error: User '", stderr);
                fputs(user, stderr);
                fputs("' can only write to repositories under ", stderr);
                fputs(GIT_BASE_PATH, stderr);
                fputs("/", stderr);
                fputs(user, stderr);
                fputs("/\n", stderr);
                free(repo_path);
                return 0;
            }
        }
    }
    
    if (is_read_command(git_cmd)) {
        if (!validate_repo_path(repo_path, NULL)) {
            fputs("Error: Repository path must be under ", stderr);
            fputs(GIT_BASE_PATH, stderr);
            fputs("/\n", stderr);
            free(repo_path);
            return 0;
        }
    }
    
    free(repo_path);
    return 1;
}

static void
execute_git_command(const char *cmd) {
    char command[MAX_CMD_LEN];
    git_command_t git_cmd = parse_command(cmd);
    char *repo_path = extract_repo_path(cmd);
    
    if (!repo_path) {
        fputs("Error: Invalid command format\n", stderr);
        exit(1);
    }
    
    char full_path[PATH_MAX];
    if ('/' == repo_path[0]) {
        snprintf(full_path, sizeof(full_path), "%s", repo_path);
    } else {
        snprintf(full_path, sizeof(full_path), "%s/%s", GIT_BASE_PATH, repo_path);
    }
    
    switch (git_cmd) {
        case CMD_GIT_UPLOAD_PACK:
            snprintf(command, sizeof(command), "git-upload-pack '%s'", full_path);
            break;
        case CMD_GIT_RECEIVE_PACK:
            snprintf(command, sizeof(command), "git-receive-pack '%s'", full_path);
            break;
        case CMD_GIT_UPLOAD_ARCHIVE:
            snprintf(command, sizeof(command), "git-upload-archive '%s'", full_path);
            break;
        default:
            fputs("Error: Unknown git command\n", stderr);
            free(repo_path);
            exit(1);
    }
    
    free(repo_path);
    
    execl("/bin/sh", "sh", "-c", command, NULL);
    perror("execl");
    exit(1);
}

static int32_t
handle_marrow_key(void) {
    const char *user = get_git_user();
    if (!user) {
        fputs("Error: marrow-key requires authentication\n", stderr);
        return 1;
    }
    
    fputs("Paste your SSH public key (press Ctrl+D when done):\n", stdout);
    fflush(stdout);
    
    char key[MAX_KEY_LEN];
    size_t total_read = 0;
    ssize_t bytes_read;
    
    while (total_read < sizeof(key) - 1) {
        bytes_read = read(STDIN_FILENO, key + total_read, sizeof(key) - total_read - 1);
        if (bytes_read <= 0) break;
        total_read += bytes_read;
    }
    key[total_read] = '\0';
    
    /* Remove trailing newline if present */
    if (total_read > 0 && '\n' == key[total_read - 1]) {
        key[total_read - 1] = '\0';
    }
    
    /* Validate key format (basic check) */
    if (0 != strncmp(key, "ssh-", 4)) {
        fputs("Error: Invalid SSH key format\n", stderr);
        return 1;
    }
    
    /* Open auth file for appending */
    FILE *fp = fopen(AUTH_FILE, "a");
    if (!fp) {
        fputs("Error: Cannot open authorized_keys file\n", stderr);
        return 1;
    }
    
    /* Write key with username as comment */
    fputs(key, fp);
    fputs(" ", fp);
    fputs(user, fp);
    fputs("\n", fp);
    fclose(fp);
    
    fputs("SSH key added successfully for user: ", stdout);
    fputs(user, stdout);
    fputs("\n", stdout);
    
    return 0;
}

static void
execute_command(const char *cmd) {
    git_command_t command = parse_command(cmd);
    
    if (CMD_MARROW_KEY == command) {
        exit(handle_marrow_key());
    } else {
        execute_git_command(cmd);
    }
}

int32_t
main(int32_t argc, char *argv[]) {
    const char *username = NULL;
    int32_t arg_offset = 0;
    
    /* Check for -u USER flag for authenticated access */
    if (argc > 2 && 0 == strcmp(argv[1], "-u")) {
        username = argv[2];
        arg_offset = 2;
        /* Set GIT_USER environment variable */
        setenv("GIT_USER", username, 1);
    }
    
    /* Check for SSH_ORIGINAL_COMMAND first (SSH forced command) */
    const char *ssh_cmd = getenv("SSH_ORIGINAL_COMMAND");
    const char *cmd = NULL;
    
    if (ssh_cmd) {
        cmd = ssh_cmd;
    } else if (argc >= (3 + arg_offset) && 0 == strcmp(argv[1 + arg_offset], "-c")) {
        /* Fallback to command line argument */
        cmd = argv[2 + arg_offset];
    } else {
        fputs("Usage: ", stderr);
        fputs(argv[0], stderr);
        fputs(" [-u USER] -c <command>\n", stderr);
        fputs("This shell is designed to work with git SSH access\n", stderr);
        fputs("Usually invoked via SSH with SSH_ORIGINAL_COMMAND set\n", stderr);
        fputs("Supported commands:\n", stderr);
        fputs("  git-upload-pack <repo>   - Clone/fetch repository\n", stderr);
        fputs("  git-receive-pack <repo>  - Push to repository\n", stderr);
        fputs("  git-upload-archive <repo> - Archive repository\n", stderr);
        fputs("  marrow-key              - Add SSH key for authenticated user\n", stderr);
        return 1;
    }
    
    const char *user = get_git_user();
    
    if (user) {
        fputs("Authenticated as: ", stderr);
        fputs(user, stderr);
        fputs("\n", stderr);
    } else {
        fputs("Anonymous access\n", stderr);
    }
    
    if (!check_permissions(cmd)) {
        return 1;
    }
    
    execute_command(cmd);
    
    return 0;
}