#include <iostream>
#include <string>
#include <thread>
#include <chrono>
#include <map>
#include <memory>
#include <pty.h>
#include <unistd.h>
#include <sys/wait.h>
#include <signal.h>
#include <termios.h>
#include <fcntl.h>
#include <poll.h>

// WebSocket and JSON libraries
#include <websocketpp/config/asio_no_tls_client.hpp>
#include <websocketpp/client.hpp>
#include <nlohmann/json.hpp>

using json = nlohmann::json;

typedef websocketpp::client<websocketpp::config::asio_client> client;

class PTYSession
{
public:
	PTYSession(const std::string &session_id)
		 : session_id_(session_id), master_fd_(-1), child_pid_(-1) {}

	~PTYSession()
	{
		cleanup();
	}

	bool start()
	{
		struct winsize ws;
		ws.ws_row = 24;
		ws.ws_col = 80;
		ws.ws_xpixel = 0;
		ws.ws_ypixel = 0;

		child_pid_ = forkpty(&master_fd_, nullptr, nullptr, &ws);

		if (child_pid_ == -1)
		{
			std::cerr << "Failed to create PTY" << std::endl;
			return false;
		}

		if (child_pid_ == 0)
		{
			// Child process - start shell
			setenv("TERM", "xterm-256color", 1);
			execl("/bin/bash", "bash", nullptr);
			exit(1);
		}

		// Parent process - set non-blocking
		int flags = fcntl(master_fd_, F_GETFL);
		fcntl(master_fd_, F_SETFL, flags | O_NONBLOCK);

		std::cout << "Started PTY session " << session_id_ << " with PID " << child_pid_ << std::endl;
		return true;
	}

	void write(const std::string &data)
	{
		if (master_fd_ != -1)
		{
			::write(master_fd_, data.c_str(), data.length());
		}
	}

	std::string read()
	{
		if (master_fd_ == -1)
			return "";

		char buffer[4096];
		ssize_t bytes = ::read(master_fd_, buffer, sizeof(buffer) - 1);

		if (bytes > 0)
		{
			buffer[bytes] = '\0';
			return std::string(buffer, bytes);
		}

		return "";
	}

	void resize(int cols, int rows)
	{
		if (master_fd_ == -1)
			return;

		struct winsize ws;
		ws.ws_row = rows;
		ws.ws_col = cols;
		ws.ws_xpixel = 0;
		ws.ws_ypixel = 0;

		ioctl(master_fd_, TIOCSWINSZ, &ws);
	}

	bool is_alive()
	{
		if (child_pid_ == -1)
			return false;

		int status;
		int result = waitpid(child_pid_, &status, WNOHANG);
		return result == 0;
	}

	int get_fd() const { return master_fd_; }

private:
	void cleanup()
	{
		if (master_fd_ != -1)
		{
			close(master_fd_);
			master_fd_ = -1;
		}

		if (child_pid_ != -1 && is_alive())
		{
			kill(child_pid_, SIGTERM);
			sleep(1);
			if (is_alive())
			{
				kill(child_pid_, SIGKILL);
			}
			waitpid(child_pid_, nullptr, 0);
			child_pid_ = -1;
		}
	}

	std::string session_id_;
	int master_fd_;
	pid_t child_pid_;
};

class PTYHost
{
public:
	PTYHost(const std::string &host_id, const std::string &server_url, const std::string &token)
		 : host_id_(host_id), server_url_(server_url), token_(token), running_(false), connection_lost_(false), reconnect_count_(0)
	{

		// Initialize WebSocket client
		client_.clear_access_channels(websocketpp::log::alevel::all);
		client_.clear_error_channels(websocketpp::log::elevel::all);
		client_.set_access_channels(websocketpp::log::alevel::connect);
		client_.set_access_channels(websocketpp::log::alevel::disconnect);
		client_.set_access_channels(websocketpp::log::alevel::app);

		client_.init_asio();

		client_.set_message_handler([this](websocketpp::connection_hdl hdl, client::message_ptr msg)
											 { handle_message(msg->get_payload()); });

		client_.set_open_handler([this](websocketpp::connection_hdl hdl)
										 {
            if (reconnect_count_ > 0) {
                std::cout << "\rSuccessfully reconnected to relay server after " 
                          << reconnect_count_ << " attempts" << std::string(20, ' ') << std::endl;
            } else {
                std::cout << "Connected to relay server" << std::endl;
            }
            reconnect_count_ = 0;  // Reset counter on successful connection
            authenticate(); });

		client_.set_close_handler([this](websocketpp::connection_hdl hdl)
										  {
            std::cout << "Disconnected from relay server, will retry..." << std::endl;
            connection_lost_ = true; });

		client_.set_fail_handler([this](websocketpp::connection_hdl hdl)
										 {
            // std::cout << "Failed to connect to relay server, will retry..." << std::endl;
            connection_lost_ = true; });
	}

	void start()
	{
		running_ = true;

		// Start background thread for reading PTY output
		pty_thread_ = std::thread(&PTYHost::pty_reader_loop, this);

		// Main reconnection loop
		while (running_)
		{
			connection_lost_ = false;

			try
			{
				if (reconnect_count_ == 0)
				{
					std::cout << "Attempting to connect to relay server..." << std::endl;
				}
				else
				{
					std::cout << "\rReconnection attempt #" << reconnect_count_ << "..." << std::flush;
				}

				websocketpp::lib::error_code ec;
				auto con = client_.get_connection(server_url_, ec);

				if (ec)
				{
					reconnect_count_++;
					std::cout << "\rError creating connection (attempt " << reconnect_count_ << "): "
								 << ec.message() << std::string(20, ' ') << std::flush;
					std::this_thread::sleep_for(std::chrono::seconds(5));
					continue;
				}

				connection_ = con;
				client_.connect(con);

				// Run WebSocket client - this blocks until connection is lost
				client_.run();

				// Reset the client for next connection attempt
				client_.reset();
			}
			catch (const std::exception &e)
			{
				reconnect_count_++;
				std::cout << "\rException (attempt " << reconnect_count_ << "): "
							 << e.what() << std::string(20, ' ') << std::flush;
			}

			// If we get here, connection was lost
			if (running_)
			{
				reconnect_count_++;
				std::cout << "\rConnection lost. Retrying in 5 seconds... (attempt "
							 << reconnect_count_ << ")" << std::string(10, ' ') << std::flush;
				std::this_thread::sleep_for(std::chrono::seconds(5));
			}
		}
	}

	void stop()
	{
		running_ = false;

		if (pty_thread_.joinable())
		{
			pty_thread_.join();
		}

		// Clean up all sessions
		sessions_.clear();

		client_.stop();
	}

private:
	void authenticate()
	{
		json auth_msg;
		auth_msg["type"] = "auth";
		auth_msg["token"] = token_;

		send_message(auth_msg);
	}

	void handle_message(const std::string &payload)
	{
		try
		{
			json msg = json::parse(payload);
			std::string type = msg["type"];

			if (type == "shell_start")
			{
				std::string session_id = msg["session_id"];
				start_shell_session(session_id);
			}
			else if (type == "shell_data")
			{
				std::string session_id = msg["session_id"];
				std::string data = msg["data"];
				write_to_session(session_id, data);
			}
			else if (type == "shell_resize")
			{
				std::string session_id = msg["session_id"];
				int cols = msg["cols"];
				int rows = msg["rows"];
				resize_session(session_id, cols, rows);
			}
			else if (type == "error")
			{
				std::string error = msg["message"];
				std::cout << "Server error: " << error << std::endl;
			}
		}
		catch (const std::exception &e)
		{
			std::cout << "Error parsing message: " << e.what() << std::endl;
		}
	}

	void start_shell_session(const std::string &session_id)
	{
		auto session = std::make_shared<PTYSession>(session_id);
		if (session->start())
		{
			sessions_[session_id] = session;
			std::cout << "Started shell session: " << session_id << std::endl;
		}
	}

	void write_to_session(const std::string &session_id, const std::string &data)
	{
		auto it = sessions_.find(session_id);
		if (it != sessions_.end())
		{
			it->second->write(data);
		}
	}

	void resize_session(const std::string &session_id, int cols, int rows)
	{
		auto it = sessions_.find(session_id);
		if (it != sessions_.end())
		{
			it->second->resize(cols, rows);
		}
	}

	void send_message(const json &msg)
	{
		try
		{
			std::string payload = msg.dump();
			client_.send(connection_, payload, websocketpp::frame::opcode::text);
		}
		catch (const std::exception &e)
		{
			std::cout << "Error sending message: " << e.what() << std::endl;
		}
	}

	void pty_reader_loop()
	{
		while (running_)
		{
			// Check all active sessions for output
			for (auto it = sessions_.begin(); it != sessions_.end();)
			{
				auto &session = it->second;

				if (!session->is_alive())
				{
					std::cout << "Session " << it->first << " terminated" << std::endl;
					it = sessions_.erase(it);
					continue;
				}

				std::string output = session->read();
				if (!output.empty())
				{
					json msg;
					msg["type"] = "shell_data";
					msg["data"] = output;
					msg["session_id"] = it->first;
					send_message(msg);
				}

				++it;
			}

			std::this_thread::sleep_for(std::chrono::milliseconds(10));
		}
	}

	std::string host_id_;
	std::string server_url_;
	std::string token_;
	bool running_;
	bool connection_lost_;
	int reconnect_count_;

	client client_;
	client::connection_ptr connection_;

	std::map<std::string, std::shared_ptr<PTYSession>> sessions_;
	std::thread pty_thread_;
};

// Signal handler for graceful shutdown
PTYHost *g_host = nullptr;

void signal_handler(int signal)
{
	std::cout << "\nReceived signal " << signal << ", shutting down..." << std::endl;
	if (g_host)
	{
		g_host->stop();
	}
	exit(0);
}

int main(int argc, char *argv[])
{
	if (argc < 4)
	{
		std::cout << "Usage: " << argv[0] << " <host_id> <server_url> <token>" << std::endl;
		std::cout << "Example: " << argv[0] << " host-1 ws://localhost:8081 <jwt_token>" << std::endl;
		return 1;
	}

	std::string host_id = argv[1];
	std::string server_url = argv[2];
	std::string token = argv[3];

	// Set up signal handlers
	signal(SIGINT, signal_handler);
	signal(SIGTERM, signal_handler);

	std::cout << "Starting PTY Host: " << host_id << std::endl;
	std::cout << "Connecting to: " << server_url << std::endl;

	PTYHost host(host_id, server_url, token);
	g_host = &host;

	host.start();

	return 0;
}