#include "net/Socket.h"
#include "Pledge.h"
#include "memory/MemoryManager.h"
#include "net/UnixSocket.h"
#include "sys/Syscall.h"
#include "thread/Scheduler.h"
#include <bits/open-flags.h>

Result<u64> sys_socket(Registers*, SyscallArgs args)
{
    int domain = (int)args[0];
    int type = (int)args[1];
    // protocol is not used for now.

    if (type != SOCK_STREAM) return err(EPROTOTYPE);
    if (domain != AF_UNIX) return err(EAFNOSUPPORT);

    auto* current = Scheduler::current();
    TRY(check_pledge(current, Promise::p_unix));

    auto socket = TRY(make_shared<UnixSocket>());

    int fd = TRY(current->allocate_fd(0));

    current->fd_table[fd] = FileDescriptor { TRY(make_shared<OpenFileDescription>(socket, O_RDWR)), 0 };

    return fd;
}

Result<u64> sys_bind(Registers*, SyscallArgs args)
{
    int sockfd = (int)args[0];
    struct sockaddr* addr = (struct sockaddr*)args[1];
    socklen_t addrlen = (socklen_t)args[2];

    struct sockaddr_storage storage;
    if ((usize)addrlen > sizeof(storage)) return err(EINVAL);
    if (!MemoryManager::copy_from_user(addr, &storage, addrlen)) return err(EFAULT);

    auto* current = Scheduler::current();
    TRY(check_pledge(current, Promise::p_unix));

    auto inode = TRY(current->resolve_fd(sockfd))->inode();

    if (inode->type() != VFS::InodeType::Socket) return err(ENOTSOCK);

    auto socket = (SharedPtr<Socket>)inode;

    TRY(socket->bind((struct sockaddr*)&storage, addrlen));

    return 0;
}

Result<u64> sys_connect(Registers* regs, SyscallArgs args)
{
    int sockfd = (int)args[0];
    struct sockaddr* addr = (struct sockaddr*)args[1];
    socklen_t addrlen = (socklen_t)args[2];

    struct sockaddr_storage storage;
    if ((usize)addrlen > sizeof(storage)) return err(EINVAL);
    if (!MemoryManager::copy_from_user(addr, &storage, addrlen)) return err(EFAULT);

    auto* current = Scheduler::current();
    TRY(check_pledge(current, Promise::p_unix));

    auto description = TRY(current->resolve_fd(sockfd))->description;

    if (description->inode->type() != VFS::InodeType::Socket) return err(ENOTSOCK);

    auto socket = (SharedPtr<Socket>)description->inode;

    TRY(socket->connect(regs, description->flags, (struct sockaddr*)&storage, addrlen));

    return 0;
}

Result<u64> sys_listen(Registers*, SyscallArgs args)
{
    int sockfd = (int)args[0];
    int backlog = (int)args[1];

    auto* current = Scheduler::current();
    TRY(check_pledge(current, Promise::p_unix));

    auto inode = TRY(current->resolve_fd(sockfd))->inode();

    if (inode->type() != VFS::InodeType::Socket) return err(ENOTSOCK);

    auto socket = (SharedPtr<Socket>)inode;

    TRY(socket->listen(backlog));

    return 0;
}

Result<u64> sys_accept(Registers* regs, SyscallArgs args)
{
    int sockfd = (int)args[0];
    struct sockaddr* addr = (struct sockaddr*)args[1];
    socklen_t* addrlen = (socklen_t*)args[2];

    if (addr && !addrlen) return err(EINVAL);

    socklen_t len;
    if (addr)
    {
        if (!MemoryManager::copy_from_user_typed(addrlen, &len)) return err(EFAULT);
    }

    auto* current = Scheduler::current();
    TRY(check_pledge(current, Promise::p_unix));

    auto description = TRY(current->resolve_fd(sockfd))->description;

    if (description->inode->type() != VFS::InodeType::Socket) return err(ENOTSOCK);

    auto socket = (SharedPtr<Socket>)description->inode;

    struct sockaddr* client;
    socklen_t client_len;
    auto new_description = TRY(socket->accept(regs, description->flags, &client, &client_len));

    int fd = TRY(current->allocate_fd(0));
    current->fd_table[fd] = FileDescriptor { new_description, 0 };

    if (client_len < len) len = client_len;
    if (addr)
    {
        MemoryManager::copy_to_user(addr, client, len);
        MemoryManager::copy_to_user_typed(addrlen, &client_len);
    }

    return fd;
}