#include <iostream>
#include <fstream>
#include <string>
#include <format>
#include <filesystem>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <sys/stat.h>
#include <time.h>

#include "server.h"
#include "filer.cpp"
#include "util.h"

namespace fs = std::filesystem;

class Client {
public:
	int control_sock;
	/* == State
	 * -1 - Delete
	 *  0 - Visitor / Unauthenticated
	 *  1 - Authenticated
	 *  2 - Waiting for data
	 */
	int state = 0;

	Client(int &_sock) {
		control_sock = _sock;
		submit(230, "Init");
	};
	
	int receive(std::string cmd, std::string argstr) {
		if (cmd == "QUIT") {
			submit(250, "Goodbye!");
			return -1;
		} else if (state > 0) {
			if (cmd == "SYST") {
				submit(215, "UNIX");
			} else if (cmd == "PWD") {
				submit(257, "'"+filer->cwd.string()+"'");
			} else if (cmd == "CWD") {
				int rc = filer->traverse(argstr);
				if (rc == 0) {
					submit(250, "OK");
				} else if (rc == -1) {
					submit(500, "Not found!");
				} else {
					submit(500, std::to_string(rc));
				}
			} else if (cmd == "CDUP") {
				int rc = filer->traverse("..");
				if (rc == 0) {
					submit(250, "OK");
				} else if (rc == -1) {
					submit(500, "Not found!");
				} else {
					submit(500, std::to_string(rc));
				}
			} else if (cmd == "MKD") {
				int rc = filer->createDirectory(argstr);
				if (rc == 0) {
					submit(250, "OK");
				} else {
					submit(500, std::to_string(rc));
				}
			} else if (cmd == "TYPE") {
				sscanf(argstr.c_str(), "%c", &(filer->type));
				submit(226, "OK");
			} else if (cmd == "PASV") {
				submit(502, "Command not implemented");
			} else if (cmd == "PORT") {
				data_sock = socket(AF_INET, SOCK_STREAM, 0);
				
				unsigned char act_port[2];
				int act_ip[4], port_dec;
				char ip_decimal[40];
				sscanf(argstr.c_str(), "%d,%d,%d,%d,%d,%d", &act_ip[0], &act_ip[1], &act_ip[2], &act_ip[3], (int*)&act_port[0], (int*)&act_port[1]);
				sprintf(ip_decimal, "%d.%d.%d.%d", act_ip[0], act_ip[1], act_ip[2], act_ip[3]);
				port_dec = act_port[0]*256+act_port[1];
				printf("[d] D(%i) Initialized data transport: %s:%d\n", control_sock, ip_decimal, port_dec);
				
				data_address.sin_family = AF_INET;
				data_address.sin_addr.s_addr = inet_addr(ip_decimal);
				data_address.sin_port = htons(port_dec);
						
				if (connect(data_sock, (struct sockaddr *)&data_address, sizeof(data_address)) != 0) {
					close(data_sock);
					data_sock = 0;
					submit(425, "Unknown Error");
				} else {
					state = 2;
					submit(200, "OK");
				}
			} else if (cmd == "SIZE") {
				int ret = filer->fileSize(argstr);
				if (ret >= 0) {
					submit(213, std::to_string(ret));
				} else {
					submit(550, "Access Denied "+std::to_string(ret));
				}
			} else if (cmd == "FEAT") {
				submit(211, ":\r\n211 END");
			} else if (cmd == "NOOP") {
				submit(226, "OK");
			} else if (cmd == "DELE" || cmd == "RMD") {
				int ret = filer->deleteFile(argstr);
				if (ret == 0) {
					submit(226, "OK");
				} else {
					submit(550, "Access Denied "+std::to_string(ret));
				}
			} else if (state == 2) {
				if (data_sock <= 0) return -1;
				if (cmd == "LIST" || cmd == "NLST") {
					submit(150, "Transferring");
					
					std::string output = filer->list();
					char out[output.size()] = {0};
					strncpy(out, output.c_str(), output.size());
					data_submit(out, output.size());
					
					submit(226, "OK");
				} else if (cmd == "RETR") {
					struct file_data fd = filer->readFile(argstr);
					if (fd.ecode == 0) {
						submit(150, "Transferring");
						data_submit(fd.data, fd.size);
					} else {
						if (fd.ecode == -3) submit(550, "I refuse to transmit in ASCII mode!");
						else submit(550, "Access Denied");
					}
					// Segfault on freeing ASCII data, memory leak on binary.
					if (filer->type != 'A') free(fd.data);
					submit(226, "OK");
				} else if (cmd == "STOR") {
					unsigned char inbuf[BUFFERSIZE] = {0};
					submit(150, "Transferring");
					int psize;
					// Write nothing to file to make sure it exists.
					// Also so I don't have to write a more complex writer.
					int ret = filer->writeFile(argstr, inbuf, 0);
					while ((psize = recv(data_sock, inbuf, sizeof(inbuf), 0)) > 0) {
						ret = filer->appendFile(argstr, inbuf, psize);
						if (ret < 0) {
							submit(550, "Access Denied "+std::to_string(ret));
							data_close();
							return -1;
						}
						for (int i = 0; i < BUFFERSIZE; i++) inbuf[i] = 0;
					}
					submit(226, "OK");
				}
				if (data_close() != 0) {
					printf("[!] Data Sock returned %i instead of 0! Not closed!!!\n", data_sock);
				}
			} else {
				submit(502, "Command not implemented");
			}
		} else {
			state = 0; // If we reach here, force state to zero just in case.
			if (cmd == "USER") {
				name = argstr;
				submit(331, "Password required");
			} else if (cmd == "PASS") {
				if (argstr == "123") {
					printf("[d] (%i) logged in as '%s'\n", control_sock, name.c_str());
					// We can now set the root safely (I hope).
					filer->setRoot("/data/ftp-test/"+name);
					state = 1;
					submit(230, "Login successful");
				} else {
					submit(530, "Authentication refused.");
					return -1;
				}
			} else if (cmd == "AUTH") {
				submit(502, "Command not implemented");
			} else {
				submit(332, "Login Required");
			}
		}
		return 0;
	}
	
	int submit(int code, std::string msg) {
		std::string out = std::to_string(code)+" "+msg+"\r\n";
		int bytes = send(control_sock, out.c_str(), out.size(), 0);
		if (bytes < 0) {
			printf("[d] C(%i) !< %i %s\n", control_sock, code, msg.c_str());
			return 1;
		}
		printf("[d] C(%i) << %i %s\n", control_sock, code, msg.c_str());
		return 0;
	}
	
	int data_submit(char* out, int size) {
		int bytes = send(data_sock, out, size, 0);
		if (bytes < 0) {
			printf("[d] D(%i) !< %i\n", data_sock, size);
			return 1;
		}
		printf("[d] D(%i) << %i\n", data_sock, size);
		return 0;
	}
	
	int data_close() {
		printf("[d] D(%i) Closing...\n", data_sock);
		close(data_sock);
		data_sock = 0;
		state = 1;
		return data_sock; // Should be 0 if successful.
	}
	
private:
	std::string name;
	Filer* filer = new Filer();
	
	int data_sock;
	struct sockaddr_in data_address;
};