#include <sol/types.hpp>
#include <sstream>

#include "command_processor.h"
#include "vfs.h"
#include "lua_processor.h"
#include "machine_manager.h"
#include "machine.h"

CommandProcessor::CommandProcessor(Machine* home_machine,
                                   std::map<std::string, Machine*>& world_machines) :
    _home_machine(home_machine),
    _session_machine(home_machine),
    _world_machines(world_machines) {
  _lua = new LuaProcessor();
  _current_dir = _session_machine->vfs_root;
}

CommandProcessor::~CommandProcessor(void) {
  delete _lua;
}

vfs_node* CommandProcessor::get_current_dir(void) {
  return _current_dir;
}

std::string CommandProcessor::process_command(const std::string& command) {
  /* Trim trailing whitespace. */
  std::string cmd = command;
  size_t end = cmd.find_last_not_of(" \t\n\r");
  cmd = (end == std::string::npos) ? "" : cmd.substr(0, end+1);

  /* === Generic script executer. === */
  /* Parse the command and it's arguments. */
  std::stringstream ss(cmd);
  std::string command_name;
  ss >> command_name;

  std::vector<std::string> args;
  std::string arg;
  while(ss >> arg) {
    args.push_back(arg);
  }

  /* Search for script in the /bin directory of the current session machine. */
  std::string script_filename = command_name + ".lua";
  vfs_node* root = _session_machine->vfs_root;
  if(root->children.count("bin") && root->children["bin"]->children.count(script_filename)) {
    vfs_node* script_node = root->children["bin"]->children[script_filename];
    bool is_remote = (_session_machine != _home_machine);
    sol::object result = _lua->execute(script_node->content, _current_dir, args, is_remote);
    if(result.is<std::string>()) {
      return result.as<std::string>();
    } else if(result.is<sol::table>()) {
      return _handle_vfs_action(result.as<sol::table>());
    }
    return "[Script returned an unexpected type]";
  }
  return "Unknown command: " + command_name + "\n";
}

/* Find a VFS node by it's absolute path. */
vfs_node* find_node_by_path(vfs_node* root, const std::string& path) {
  if(path == "/") {
    return root;
  }
  vfs_node* current = root;
  std::stringstream ss(path);
  std::string segment;

  /* Discard the first empty segment that comes form the leading '/'. */
  if(path[0] == '/') {
    std::getline(ss, segment, '/');
  }

  while(std::getline(ss, segment, '/')) {
    if(current->type == DIR_NODE && current->children.count(segment)) {
      current = current->children[segment];
    } else {
      return nullptr; /* Path segmenet not found. */
    }
  }
  return current;
}

/* Recursively set all nodes in a VFS tree to writeable. */
void make_vfs_writable(vfs_node* node) {
  if(!node) return;
  node->read_only = false;
  for(auto const& [key, child] : node->children) {
    make_vfs_writable(child);
  }
}

std::string CommandProcessor::_handle_vfs_action(sol::table action) {
  std::string action_name = action["action"].get_or<std::string>("");
  /* Make the CoW check universal for any write operation. */
  if(action_name == "rm" || action_name == "write_file") {
    if(_current_dir->read_only) {
      std::string remote_ip = "";
      /* Check if we are on a known remote system. */
      if(!_session_machine->is_vfs_a_copy) {
        for(auto const& [ip, machine] : _world_machines) {
          if(machine == _session_machine) {
            remote_ip = ip;
            break;
          }
        }
      }

      if(!remote_ip.empty()) {
        /* This machine is using a shared VFS. Copy it! ;) */
        fprintf(stderr, "CoW: Write attempt on read-only remote system '%s'."
                "Creating persistant copy.\n", remote_ip.c_str());
        std::string original_path = get_full_path(_current_dir);
        /* Copy the sh.t out of the VFS. */
        vfs_node* new_vfs_root = copy_vfs_node(_session_machine->vfs_root, nullptr);
        make_vfs_writable(new_vfs_root);

        /* Create the new machine that will hold the copied VFS. */
        auto* new_machine = new Machine(_session_machine->id, _session_machine->hostname);
        new_machine->vfs_root = new_vfs_root;
        new_machine->services = _session_machine->services;
        new_machine->is_vfs_a_copy = true; /* Mark as copy. */

        /* Update the world map and the current session. */
        _world_machines[remote_ip] = new_machine;
        _session_machine = new_machine;

        /* Update _current_dir to point to the equivalent dir in the new VFS. */
        _current_dir = find_node_by_path(_session_machine->vfs_root, original_path);
        if(!_current_dir) { _current_dir = _session_machine->vfs_root; }
      } else {
        return "Permission denied: Filesystem is read-only and not part of the world.";
      }
    }
  }
  if(action_name == "rm") {
    std::string path = action["target"].get_or<std::string>("");
    if(path.empty()) return "rm: missing operand";

    if(!_current_dir->children.count(path)) {
      return std::string("rm: cannot remove '") + path + "': No such file or directory.";
    }
    vfs_node* target_node = _current_dir->children[path];
    if(target_node->type == DIR_NODE) {
      return std::string("rm: cannot remove '") + path + "': Is a directory.";
    }
    _current_dir->children.erase(path);
    delete target_node;
  } else if(action_name == "write_file") {
    std::string filename  = action["target"].get_or<std::string>("");
    std::string content   = action["content"].get_or<std::string>("");
    if(filename.empty()) {
      return "write_file: missing filename";
    }
 
    if(_current_dir->children.count(filename)) {
      vfs_node* target_node = _current_dir->children[filename];
      if(target_node->type == DIR_NODE) {
        return std::string("cannot write to '") + filename + "': Is a directory.";
      }
      target_node->content = content;
    } else {
      vfs_node* new_file = new vfs_node {
        .name = filename, .type = FILE_NODE, .read_only = false, .content = content,
        .parent = _current_dir
      };
      _current_dir->children[filename] = new_file;
    }
    return ""; /* Success. */
  } else if(action_name == "ssh") {
    std::string target_ip = action["target"].get_or<std::string>("");
    if(_world_machines.count(target_ip)) {
      _session_machine = _world_machines[target_ip];
      _current_dir = _session_machine->vfs_root;
      return "Connected to " + target_ip;
    }
    return "ssh: Could not resolve hostname " + target_ip + ": Name or service not known";
  } else if(action_name == "cd") {
    std::string target_dir_name = action["target"].get_or<std::string>("");
    if(target_dir_name == "..") {
      if(_current_dir->parent) {
        _current_dir = _current_dir->parent;
      }
    } else if(_current_dir->children.count(target_dir_name)) {
      vfs_node* target_node = _current_dir->children[target_dir_name];
      if(target_node->type == DIR_NODE) {
        _current_dir = target_node;
      } else {
        return std::string("cd: not a directory: ") + target_dir_name;
      }
    } else {
        return std::string("cd: no such file or directory: ") + target_dir_name;
    }
      return ""; /* Success. */
  } else if(action_name == "disconnect") {
    _session_machine = _home_machine;
    _current_dir = _home_machine->vfs_root;
    return "Connection closed.";
  } else if(action_name == "close_terminal") {
    return "__CLOSE_CONNECTION__";
  } else if(action_name == "scan") {
    std::string target_ip = action["target"].get_or<std::string>("");
    if(!_world_machines.count(target_ip)) {
      return std::string("nmap: Could not resolve host: ") + target_ip;
    }

    Machine* target_machine = _world_machines[target_ip];
    if(target_machine->services.empty()) {
      return std::string("No open ports for ") + target_ip;
    }

    std::stringstream ss;
    ss<<"Host: "<<target_ip<<"\n";
    ss<<"PORT\tSTATE\tSERVICE\n";
    for(auto const& [port, service_name] : target_machine->services) {
      ss<<port<<"/tcp\t"<<"open\t"<<service_name<<"\n";
    }
    return ss.str();
  }
  return "Error: Unknown VFS action '" + action_name + "'";
}