#include "Pledge.h"
#include "memory/MemoryManager.h"
#include "sys/Syscall.h"
#include "thread/Scheduler.h"
#include <bits/atfile.h>

Result<u64> sys_getpid(Registers*, SyscallArgs)
{
    auto* current = Process::current();
    TRY(check_pledge(current, Promise::p_stdio));
    return current->id;
}

Result<u64> sys_getppid(Registers*, SyscallArgs)
{
    auto* current = Process::current();
    TRY(check_pledge(current, Promise::p_stdio));
    auto* parent = current->parent;
    return parent ? parent->id : 0;
}

Result<u64> sys_getuid(Registers*, SyscallArgs)
{
    auto* current = Process::current();
    TRY(check_pledge(current, Promise::p_stdio));
    return current->credentials().uid;
}

Result<u64> sys_geteuid(Registers*, SyscallArgs)
{
    auto* current = Process::current();
    TRY(check_pledge(current, Promise::p_stdio));
    return current->credentials().euid;
}

Result<u64> sys_getgid(Registers*, SyscallArgs)
{
    auto* current = Process::current();
    TRY(check_pledge(current, Promise::p_stdio));
    return current->credentials().gid;
}

Result<u64> sys_getegid(Registers*, SyscallArgs)
{
    auto* current = Process::current();
    TRY(check_pledge(current, Promise::p_stdio));
    return current->credentials().egid;
}

Result<u64> sys_setuid(Registers*, SyscallArgs args)
{
    u32 uid = (u32)args[0];

    auto* current = Process::current();
    TRY(check_pledge(current, Promise::p_id));
    auto auth = current->auth.lock();

    if (auth->euid == 0)
    {
        auth->uid = auth->euid = auth->suid = uid;
        return 0;
    }

    if (uid != auth->uid && uid != auth->suid) return err(EPERM);
    auth->euid = uid;

    return 0;
}

Result<u64> sys_seteuid(Registers*, SyscallArgs args)
{
    u32 uid = (u32)args[0];

    auto* current = Process::current();
    TRY(check_pledge(current, Promise::p_id));
    auto auth = current->auth.lock();

    if (auth->euid != 0 && uid != auth->uid && uid != auth->suid) return err(EPERM);
    auth->euid = uid;

    return 0;
}

Result<u64> sys_setgid(Registers*, SyscallArgs args)
{
    u32 gid = (u32)args[0];

    auto* current = Process::current();
    TRY(check_pledge(current, Promise::p_id));
    auto auth = current->auth.lock();

    if (auth->euid == 0)
    {
        auth->gid = auth->egid = auth->sgid = gid;
        return 0;
    }

    if (gid != auth->gid && gid != auth->sgid) return err(EPERM);
    auth->egid = gid;

    return 0;
}

Result<u64> sys_setegid(Registers*, SyscallArgs args)
{
    u32 gid = (u32)args[0];

    auto* current = Process::current();
    TRY(check_pledge(current, Promise::p_id));
    auto auth = current->auth.lock();

    if (auth->euid != 0 && gid != auth->gid && gid != auth->sgid) return err(EPERM);
    auth->egid = gid;

    return 0;
}

Result<u64> sys_setpgid(Registers*, SyscallArgs args)
{
    pid_t pid = (pid_t)args[0];
    pid_t pgid = (pid_t)args[1];

    auto* current = Process::current();
    TRY(check_pledge(current, Promise::p_proc));
    if (pid == 0) pid = current->id;
    if (pgid == 0) pgid = current->id;

    if (pgid < 0) return err(EINVAL);

    auto* target = TRY(Result<Process*>::from_option(Scheduler::find_by_pid(pid), ESRCH));
    if (target != current && target->parent != current) return err(ESRCH);

    if (target->is_session_leader() || target->sid != current->sid) return err(EPERM);

    if (target->has_called_exec) return err(EACCES);

    if (pgid != current->id)
    {
        bool pgid_exists = false;
        pid_t sid;
        Scheduler::for_each_in_process_group(pgid, [&pgid_exists, &sid](Process* p) {
            pgid_exists = true;
            sid = p->sid; // this should be the same for all threads in the process group
            return false;
        });
        if (!pgid_exists) return err(EPERM);
        if (sid != target->sid) return err(EPERM);
    }

    target->pgid = (u64)pgid;

    return 0;
}

Result<u64> sys_getpgid(Registers*, SyscallArgs args)
{
    pid_t pid = (pid_t)args[0];

    auto* current = Process::current();
    TRY(check_pledge(current, Promise::p_stdio));
    if (pid == 0) pid = current->id;

    if (pid < 0) return err(EINVAL);

    auto* process = TRY(Result<Process*>::from_option(Scheduler::find_by_pid(pid), ESRCH));

    return (u64)process->pgid.load();
}

Result<u64> sys_setsid(Registers*, SyscallArgs)
{
    auto* current = Process::current();
    TRY(check_pledge(current, Promise::p_proc));

    if (current->pgid == current->id) return err(EPERM);

    current->sid = current->pgid = current->id;
    current->controlling_terminal = {};

    return current->sid.load();
}

Result<u64> sys_getsid(Registers*, SyscallArgs args)
{
    pid_t pid = (pid_t)args[0];

    auto* current = Process::current();
    TRY(check_pledge(current, Promise::p_stdio));

    if (pid == 0) pid = current->id;

    auto* p = TRY(Result<Process*>::from_option(Scheduler::find_by_pid(pid), ESRCH));

    return p->sid.load();
}

Result<u64> sys_fchmodat(Registers*, SyscallArgs args)
{
    int dirfd = (int)args[0];
    auto path = TRY(MemoryManager::strdup_from_user(args[1]));
    mode_t mode = (mode_t)args[2];
    int flags = (int)args[3];

    auto* current = Process::current();
    TRY(check_pledge(current, Promise::p_wpath));
    auto credentials = current->credentials();

    auto inode = TRY(current->resolve_atfile(dirfd, path, flags & AT_EMPTY_PATH, !(flags & AT_SYMLINK_NOFOLLOW)));

    if (credentials.euid != 0 && credentials.euid != inode->metadata().uid) return err(EPERM);

    auto metadata = inode->metadata();
    metadata.mode = mode;
    TRY(inode->set_metadata(metadata));

    return 0;
}

Result<u64> sys_fchownat(Registers*, SyscallArgs args)
{
    int dirfd = (int)args[0];
    auto path = TRY(MemoryManager::strdup_from_user(args[1]));
    uid_t uid = (u32)args[2];
    gid_t gid = (u32)args[3];
    int flags = (int)args[4];

    auto* current = Process::current();
    TRY(check_pledge(current, Promise::p_chown));
    auto credentials = current->credentials();

    auto inode = TRY(current->resolve_atfile(dirfd, path, flags & AT_EMPTY_PATH, !(flags & AT_SYMLINK_NOFOLLOW)));

    if (credentials.euid != 0) return err(EPERM);

    auto metadata = inode->metadata();
    if (uid != (uid_t)-1) metadata.uid = uid;
    if (gid != (gid_t)-1) metadata.gid = gid;
    TRY(inode->set_metadata(metadata));

    return 0;
}

Result<u64> sys_getgroups(Registers*, SyscallArgs args)
{
    int ngroups = (int)args[0];
    gid_t* grouplist = (gid_t*)args[1];

    auto* current = Process::current();

    TRY(check_pledge(current, Promise::p_stdio));

    auto groups = current->extra_groups.lock();

    if (!ngroups) return groups->size();
    if (ngroups < 0) return err(EINVAL);

    if (static_cast<usize>(ngroups) < groups->size()) return err(EINVAL);

    if (!MemoryManager::copy_to_user(grouplist, groups->data(), groups->size() * sizeof(gid_t))) return err(EFAULT);

    return groups->size();
}

Result<u64> sys_setgroups(Registers*, SyscallArgs args)
{
    int ngroups = (int)args[0];
    const gid_t* grouplist = (const gid_t*)args[1];

    auto* current = Process::current();
    TRY(check_pledge(current, Promise::p_id));

    auto credentials = current->credentials();
    if (credentials.euid != 0) return err(EPERM);

    auto groups = current->extra_groups.lock();

    if (!ngroups)
    {
        groups->clear();
        return 0;
    }

    if (ngroups < 0 || ngroups > 32) return err(EINVAL);

    TRY(groups->try_reserve(ngroups));

    groups->mutate([&](gid_t* list, usize) -> usize {
        if (MemoryManager::copy_from_user(grouplist, list, ngroups * sizeof(gid_t))) return ngroups;
        return groups->size();
    });

    return 0;
}