#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <string.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <time.h>
#include <unistd.h>
#include <stdint.h>
#include <dirent.h>
#include "process.h"

#ifndef PF_KTHREAD
# define PF_KTHREAD 0x00200000
#endif

static inline uint64_t fast_str2ull(char** str) {
   uint64_t result = 0;
   int maxlen = 20; // length of maximum value of 18446744073709551615

   while (maxlen-- && **str >= '0' && **str <= '9') {
      result *= 10;
      result += **str - '0';
      (*str)++;
   }

   return result;
}

static inline int64_t fast_str2ll(char** str) {
   bool neg = false;

   if (**str == '-') {
      neg = true;
      (*str)++;
   }

   uint64_t res = fast_str2ull(str);
   int64_t result = (int64_t) res;

   return neg ? -result : result;
}

#define skipRange(inclusive,exclusive)     \
    for (int i##inclusive=0; i##inclusive<exclusive-inclusive; i##inclusive++) \
        location = strchr(location, ' ') + 1;

int parseStatFile(Process *proc, char *filedata) {
    char *location = filedata;
    // (1) pid - %d
    proc->pid = fast_str2ull(&location);
    location += 2;
     
    // (2) comm - %s
    char *end = strrchr(location, ')');
    size_t length = end - location;
    memcpy(proc->label, location, length);
    proc->label[length] = 0;
    location = end + 2;

    // (3) State - %c
    location += 2;

    // (4) ppid - %d
    proc->ppid = fast_str2ull(&location);
    location += 1;

    // skip [5 - 9)
    skipRange(5, 9);

    // (9) flags - %u
    proc->flags = fast_str2ull(&location);
    location += 1;

    proc->iskernel = (proc->flags & PF_KTHREAD) ? true : false;

    // skip [10 - 14)
    skipRange(10, 14);

    // (14) utime - %lu
    proc->cpuTime = fast_str2ull(&location);
    location += 1;

    // (15) stime - %lu
    proc->cpuTime += fast_str2ull(&location);
    location += 1;

    // (16) cutime - %lu
    proc->childCpuTime = fast_str2ull(&location);
    location += 1;

    // (17) cstime - %lu
    proc->childCpuTime += fast_str2ull(&location);
    location += 1;

    // skip 18-19
    skipRange(18, 20);

    // (20) num_threads - %ld
    proc->nThreads = fast_str2ull(&location);
    location += 1;

    return 0;
}

static inline uint64_t findAndParseField(char **filedata, char *fld) {
    char *field = strstr(*filedata, fld);
    if (!field) return -1;
    while (*field < '0' || *field > '9') field++;
    *filedata = field;
    uint64_t val = fast_str2ull(filedata);
    return val;
}

int parseStatusFile(Process *proc, char *filedata) {
    proc->tgid = findAndParseField(&filedata, "Tgid");
    proc->pid = findAndParseField(&filedata, "Pid");

    if (proc->pid != proc->tgid) return -1; // ignore child threads, as their stats are the same as the parent

    proc->resident = 1024 * findAndParseField(&filedata, "VmRSS");
    proc->swap = 1024 * findAndParseField(&filedata, "VmSwap");
    return 0;
}

int parseIOFile(Process *proc, char *filedata) {
    proc->read = findAndParseField(&filedata, "read_bytes");
    proc->written = findAndParseField(&filedata, "write_bytes");
    return 0;
}

char *readProcesses(char *procdir) {
    int len;
    struct dirent *pDirent;
    DIR *pDir;

    // Ensure we can open directory.

    pDir = opendir(procdir);
    if (pDir == NULL) {
        printf("Cannot open directory '%s'\n", procdir);
        return NULL;
    }

    // Process each entry.

    char *buffer = malloc(4096);
    char *fname = malloc(1024);
    char first;
    FILE *file;

    Process *cur = malloc(sizeof(Process));
    cur->prev = NULL;
    cur->next = NULL;
    Process *head = cur;

    while ((pDirent = readdir(pDir)) != NULL) {
        first = pDirent->d_name[0];
        if (first < '0' || first > '9') continue;
        sprintf(fname, "/proc/%s/status", pDirent->d_name);
        file = fopen(fname, "rb");
        if (file == NULL) continue;
        fread(buffer, 1, 4096, file); 
        fclose(file);

        if (parseStatusFile(cur, buffer)) continue;

        // truncate 'status' to 'stat'
        len = strlen(fname);
        fname[len-2] = 0; 
        file = fopen(fname, "rb");
        if (file == NULL) continue;
        fread(buffer, 1, 4096, file);
        fclose(file);

        if (parseStatFile(cur, buffer)) continue;
        
	sprintf(fname, "/proc/%s/io", pDirent->d_name);
        file = fopen(fname, "rb");
        if (file == NULL) continue;
        fread(buffer, 1, 4096, file);
        fclose(file);

        if (parseIOFile(cur, buffer)) continue;

	sprintf(fname, "/proc/%s/cpuset", pDirent->d_name);
	file = fopen(fname, "rb");
	if (file == NULL) goto nextProcess;
	fread(cur->cpuset, 1, sizeof(cur->cpuset), file);
	fclose(file);

	strtok(cur->cpuset, "\n\t ");
	char *lxcTag = "/lxc/";
	if (memcmp(cur->cpuset, lxcTag, strlen(lxcTag)) == 0) {
	    // Resides in LXC -- read file and tag
	    sscanf(cur->cpuset, "/lxc/%d/ns", &cur->lxc);
	    sprintf(fname, "/etc/pve/lxc/%d.conf", cur->lxc);
	    file = fopen(fname, "rb");
	    if (file == NULL) goto nextProcess;
	    fread(buffer, 1, 4096, file);
	    fclose(file);

	    char *ptr = strstr(buffer, "hostname:");
	    ptr += strlen("hostname:");
	    while (*ptr == ' ') ptr++;
	    strtok(ptr, "\n\t ");
	    strcpy(cur->lxcname, ptr);
	} else {
	    cur->lxcname[0] = 0;
            cur->lxc = -1;
	}

nextProcess:
        cur->next = malloc(sizeof(Process));
        cur->next->prev = cur;
        cur = cur->next;
    }

    // clean up unused last node
    cur = cur->prev;
    free(cur->next);
    cur->next = NULL;

    int clocks = sysconf(_SC_CLK_TCK);
    char *output = malloc(8*1024*1024);
    char *ptr = output;
    
    cur = head;
    while (cur != NULL) {
	    // create process descriptor
	sprintf(buffer, "pid=\"%d\",ppid=\"%d\",label=\"%s\",kernel=\"%d\",ctid=\"%d\",lxc=\"%s\"", cur->pid, cur->ppid, cur->label, cur->iskernel, cur->lxc, cur->lxcname);
        ptr += sprintf(ptr, "proc_cpu_time{%s} %f\n", buffer, (double) cur->cpuTime / clocks);
        ptr += sprintf(ptr, "proc_child_cpu_time{%s} %f\n", buffer, (double) cur->childCpuTime / clocks);
        ptr += sprintf(ptr, "proc_num_threads{%s} %llu\n", buffer, cur->nThreads);
        ptr += sprintf(ptr, "proc_resident_bytes{%s} %llu\n", buffer, cur->resident);
        ptr += sprintf(ptr, "proc_swap_bytes{%s} %llu\n", buffer, cur->swap);
        ptr += sprintf(ptr, "proc_fileio_bytes_written{%s} %llu\n", buffer, cur->written);
        ptr += sprintf(ptr, "proc_fileio_bytes_read{%s} %llu\n", buffer, cur->read);
        cur = cur->next;
    }

    cur = head;
    while (cur != NULL) {
        Process *prev = cur;
        cur = cur->next;
        free(prev);
    }


    closedir(pDir);
    free(buffer);
    free(fname);
    return output;
}

#define BUFFER_SIZE 1024

void handle_client(int client_socket) {
    char *data = readProcesses("/proc");
    int length = strlen(data);
    printf("Got data of length: %d\n", length);

    // HTTP response
    const char *http_headers_fmt =
        "HTTP/1.1 200 OK\r\n"
        "Content-Type: text/plain\r\n"
        "Content-Length: %u\r\n"
        "Connection: close\r\n"
        "\r\n";

    char headers[BUFFER_SIZE];
    sprintf(headers, http_headers_fmt, length);

    printf("Sending headers\n");
    if (send(client_socket, headers, strlen(headers), 0) == -1) {
	    printf("Failed to send headers\n");
	    free(data);
	    return;
	} else {
	printf("Sent headers\n");
	}

    // Send the response to the client
    int total_sent = 0;

    while (total_sent < length) {
      int amt = length - total_sent;
      int sent = send(client_socket, data+total_sent, amt, 0);
      printf("Tried sending %d actually sent %d\n", amt, sent);
      total_sent += sent;
    }

    printf("Closing connection\n");
    // Close the client socket
    shutdown(client_socket, SHUT_WR);
    sleep(1);
    close(client_socket);

    free(data);
}

int main(int argc, char *argv[]) {

    int port = 9101;

    if (argc == 1) {
        char *buf = readProcesses("/proc");
	if (buf == NULL) return 4;
	printf(buf);
	return 0;

    } else {
	port = atoi(argv[1]);
    }

    int server_socket, client_socket;
    struct sockaddr_in server_addr, client_addr;
    socklen_t addr_len = sizeof(client_addr);

    // Create the server socket
    if ((server_socket = socket(AF_INET, SOCK_STREAM, 0)) == -1) {
        perror("Socket creation failed");
        exit(EXIT_FAILURE);
    }

    // Configure the server address
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = INADDR_ANY;
    server_addr.sin_port = htons(port);

    // Bind the socket to the specified port
    if (bind(server_socket, (struct sockaddr *)&server_addr, sizeof(server_addr)) == -1) {
        perror("Bind failed");
        close(server_socket);
        exit(EXIT_FAILURE);
    }

    // Listen for incoming connections
    if (listen(server_socket, 5) == -1) {
        perror("Listen failed");
        close(server_socket);
        exit(EXIT_FAILURE);
    }

    printf("Server is listening on port %d\n", port);

    // Accept and handle incoming connections
    while (1) {
        if ((client_socket = accept(server_socket, (struct sockaddr *)&client_addr, &addr_len)) == -1) {
            perror("Accept failed");
            continue;
        }

        // Handle the client in a separate function
        handle_client(client_socket);
    }

    // Close the server socket (this won't actually be reached in this example)
    close(server_socket);

    return 0;
}