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

#include "i_network_bridge.h"
#include "lua_api.h"
#include "session.h"
#include "machine.h"
#include "vfs.h"
#include "util.h"

namespace api {

vfs_node* get_current_dir(Session& context) {
  return context.get_current_dir();
}

std::string rm(Session& context, const std::string& filename) {
  vfs_node* current_dir = context.get_current_dir();
  auto it = current_dir->children.find(filename);

  if(it == current_dir->children.end()) {
    return "rm: cannot remove '" + filename + "': No such file or directory.";
  }
  if(it->second->type == DIR_NODE) {
    return "rm: cannot remove '" + filename + "': Is a directory.";
  }

  delete it->second; /* Free the memory for the node. */
  current_dir->children.erase(it); /* Remove from map. */
  return "";
}

std::string write_file(Session& context, const std::string& path,
                       const std::string& content) {
  vfs_node* parent_dir = nullptr;
  std::string filename;

  if(path[0] == '/') {
    Machine* session_machine = context.get_session_machine();
    vfs_node* root = session_machine->vfs_root;

    size_t last_slash = path.find_last_of('/');
    if(last_slash == 0) {
      /* File in root. */
      parent_dir = root;
      filename = path.substr(1);
    } else {
      std::string parent_path = path.substr(0, last_slash);
      parent_dir = find_node_by_path(root, parent_path);
      filename = path.substr(last_slash+1);
    }
  } else {
    /* Relative path. */
    parent_dir = context.get_current_dir();
    filename = path;
  }

  if(!parent_dir) {
    return "write: cannot create file '" + path + "': No such file or directory";
  }

  if(parent_dir->type != DIR_NODE) {
    return "write: cannot create file in '" + get_full_path(parent_dir) + "': Not a directory";
  }

  auto it = parent_dir->children.find(filename);
  if(it != parent_dir->children.end()) {
    if(it->second->type == DIR_NODE) {
      return "write: " + path + ": Is a directory";
    }
    it->second->content = content;
  } else {
    /* File does not exist, create it. */
    vfs_node* new_file = new_node(filename, FILE_NODE, parent_dir);
    new_file->content = content;
    parent_dir->children[filename] = new_file;
  }
  return "";
}

std::string create_executable(Session& context, const std::string& path,
                              const std::string& content) {
  vfs_node* parent_dir = nullptr;
  std::string filename;

  if(path[0] == '/') {
    Machine* session_machine = context.get_session_machine();
    vfs_node* root = session_machine->vfs_root;

    size_t last_slash = path.find_last_of('/');
    if(last_slash == 0) {
      parent_dir = root;
      filename = path.substr(1);
    } else {
      std::string parent_path = path.substr(0, last_slash);
      parent_dir = find_node_by_path(root, parent_path);
      filename = path.substr(last_slash+1);
    }
  } else {
    parent_dir = context.get_current_dir();
    filename = path;
  }

  if(!parent_dir) {
    return "create_executable: cannot create file '" + path + "': No such file or directory";
  }

  if(parent_dir->type != DIR_NODE) {
    return "create_executable: cannot create file in '" + get_full_path(parent_dir)
        + "': Not a directory";
  }

  /* Overwrite if exists. */
  auto it = parent_dir->children.find(filename);
  if(it != parent_dir->children.end()) {
    delete it->second;
    parent_dir->children.erase(it);
  }
  vfs_node* new_exec = new_node(filename, EXEC_NODE, parent_dir);
  new_exec->content = util::xor_string(content);
  parent_dir->children[filename] = new_exec;
  return "";
}

std::string cd(Session& context, const std::string& path) {
  vfs_node* current_dir = context.get_current_dir();

  if(path == "..") {
    if(current_dir->parent) {
      context.set_current_dir(current_dir->parent);
    }
  } else {
    auto it = current_dir->children.find(path);
    if(it != current_dir->children.end() && it->second->type == DIR_NODE) {
      context.set_current_dir(it->second);
    } else {
      return "cd: no such file or directory: " + path;
    }
  }
  return "";
}

std::string ls(Session& context) {
  vfs_node* dir = context.get_current_dir();
  if(dir->type != DIR_NODE) {
    return "ls: not a directory";
  }

  std::stringstream ss;
  for(auto const& [name, node] : dir->children) {
    ss << name;
    if(node->type == DIR_NODE) {
      ss << "/";
    }
    ss << "   ";
  }
  return ss.str();
}

std::string ssh(Session& context, const std::string& ip) {
  INetworkBridge* bridge = context.get_network_bridge();
  Machine* target_machine = bridge->get_machine_by_ip(ip);

  if(target_machine) {
    context.set_session_machine(target_machine);
    return "Connected to " + ip;
  } else {
    return "ssh: Could not resolve hostname " + ip + ": Name or service not found";
  }
}

std::string nmap(Session& context, const std::string& ip) {
  INetworkBridge* bridge = context.get_network_bridge();
  Machine* target_machine = bridge->get_machine_by_ip(ip);

  if(!target_machine) {
    return "nmap: Could not resolve host: " + ip;
  }

  /* Read services from the cached machine object, not the DB. */
  auto services = target_machine->services;

  /* Release the machine from the cache if it's a remote machine. */
  if(target_machine != context.get_session_machine()) {
    bridge->release_machine(target_machine->id);
  }

  if(services.empty()) {
    return "No open ports for " + ip;
  }

  std::stringstream ss;
  ss << "Host: " << ip << "\n";
  ss << "PORT\tSTATE\tSERVICE\n";
  for(auto const& [port, service_name] : services) {
    ss << port << "/tcp\t" << "open\t" << service_name << "\n";
  }
  return ss.str();
}

std::string disconnect(Session& context) {
  Machine* current_machine = context.get_session_machine();
  if(current_machine != context.get_home_machine()) {
    INetworkBridge* bridge = context.get_network_bridge();
    bridge->release_machine(current_machine->id);
  }
  context.set_session_machine(context.get_home_machine());
  return "Connection closed.";
}

std::string close_terminal(Session& context) {
  return "__CLOSE_CONNECTION__";
}

struct ScpPath {
  std::string host;
  std::string path;
};

ScpPath parse_scp_path(const std::string& arg) {
  size_t colon_pos = arg.find(':');
  if(colon_pos != std::string::npos) {
    return {arg.substr(0, colon_pos), arg.substr(colon_pos+1)};
  } else {
    return {"", arg};
  }
}

std::string scp(Session& context, const std::string& source_arg,
                const std::string& dest_arg) {
  ScpPath source_path = parse_scp_path(source_arg);
  ScpPath dest_path = parse_scp_path(dest_arg);
  Machine* session_machine = context.get_session_machine();
  INetworkBridge* bridge = context.get_network_bridge();

  Machine* source_machine = source_path.host.empty()
    ? session_machine : bridge->get_machine_by_ip(source_path.host);

  Machine* dest_machine = dest_path.host.empty()
    ? session_machine : bridge->get_machine_by_ip(dest_path.host);

  if(!source_machine) {
    return "scp: " + source_path.host + ": Name or service not known";
  }
  if(!dest_machine) {
    if(source_machine != session_machine) bridge->release_machine(source_machine->id);
    return "scp: " + dest_path.host + ": Name or service not known";
  }

  /* Simplified to use only absolute paths for now. */
  vfs_node* source_node = find_node_by_path(source_machine->vfs_root, source_path.path);

  if(!source_node || source_node->type != FILE_NODE) {
    if(source_machine != session_machine) bridge->release_machine(source_machine->id);
    if(dest_machine != session_machine) bridge->release_machine(dest_machine->id);
    return "scp: " + source_path.path + ": No such file or directory";
  }

  vfs_node* dest_node = find_node_by_path(dest_machine->vfs_root, dest_path.path);
  vfs_node* dest_dir = nullptr;
  std::string dest_filename;

  if(dest_node && dest_node->type == DIR_NODE) {
    dest_dir = dest_node;
    dest_filename = source_node->name;
  } else {
    size_t last_slash = dest_path.path.find_last_of('/');
    if(last_slash != std::string::npos) {
      std::string parent_path = dest_path.path.substr(0, last_slash);
      if(parent_path.empty()) parent_path = "/";
      dest_dir = find_node_by_path(dest_machine->vfs_root, parent_path);
      dest_filename = dest_path.path.substr(last_slash+1);
    } else {
      dest_dir = dest_machine->vfs_root;
      dest_filename = dest_path.path;
    }
  }

  if(!dest_dir) {
    if(source_machine != session_machine) bridge->release_machine(source_machine->id);
    if(dest_machine != session_machine) bridge->release_machine(dest_machine->id);
    return "scp: " + dest_path.path + ": No such file or directory";
  }

  if(dest_dir->children.count(dest_filename)) {
    vfs_node* existing_node = dest_dir->children[dest_filename];
    if(existing_node->type == DIR_NODE) {
      if(source_machine != session_machine) bridge->release_machine(source_machine->id);
      if(dest_machine != session_machine) bridge->release_machine(dest_machine->id);
      return "scp: " + dest_path.path + ": Is a directory";
    }
    existing_node->content = source_node->content;
  } else {
    vfs_node* new_file = new_node(dest_filename, FILE_NODE, dest_dir);
    new_file->content = source_node->content;
    dest_dir->children[dest_filename] = new_file;
  }

  if(source_machine != session_machine) bridge->release_machine(source_machine->id);
  if(dest_machine != session_machine) bridge->release_machine(dest_machine->id);

  return "";
}

} /* namespace api */