#include "db.h"
#include "parser.h"
#include "util.h"
#include <algorithm>
#include <filesystem>
#include <functional>
#include <span>
#include <thread>
#include <unordered_map>

constexpr static const std::string BUILD_PATH = "/tmp/shrap";

std::unordered_map<std::string, Expr> packages;

void load_packages() {
  for (const auto &e : std::filesystem::directory_iterator("./packages/")) {
    std::ifstream file(e.path());
    if (!file) {
      std::cerr << "Failed to open " << e.path() << "\n";
      std::exit(1);
    }

    Expr pkg = Expr::parse(file);
    packages[pkg.get_one("name").value] = pkg;
  }
}

std::vector<std::string>
resolve_dependencies(const std::vector<std::string> &roots, DB &db) {
  std::unordered_map<std::string, int> state;
  std::vector<std::string> order;

  std::function<void(const std::string &)> dfs = [&](const std::string &name) {
    if (db.is_installed(name)) {
      return;
    }

    int s = state[name];
    if (s == 2) {
      return;
    }
    if (s == 1) {
      throw std::runtime_error("dependency cycle detected");
    }

    if (!packages.contains(name)) {
      std::cerr << "Package not found: " << name << "\n";
      std::exit(1);
    }

    state[name] = 1;

    try {
      const Expr &pkg = packages.at(name);
      for (const auto &dep : pkg.get("dependencies").children) {
        dfs(dep.value);
      }
    } catch (std::out_of_range &) {
    }

    state[name] = 2;
    order.push_back(name);
  };

  for (const auto &r : roots) {
    dfs(r);
  }
  return order;
}

void install_package(const std::string &name, DB &db) {
  // TODO: track installed packages

  if (!packages.contains(name)) {
    std::cerr << "Package not found: " << name << "\n";
    std::exit(1);
  }

  Expr pkg = packages[name];

  std::cout << "\n\n\tInstalling " << pkg.get_one("name").value << " ("
            << pkg.get_one("version").value << ")...\n\n";

  Expr src = pkg.get_one("src");
  std::string src_type = src.children[0].value;
  std::string src_url = src.get_one("url").value;
  std::string src_path = BUILD_PATH + "/" + src.get_one("dir").value;

  if (src_type == "tar") {
    std::string archive_path = BUILD_PATH + "/" + Util::basename(src_url);
    std::string expected_hash = src.get_one("blake3").value;

    bool needs_download = true;

    // dont redownload if checksum matches
    if (std::filesystem::exists(archive_path)) {
      std::cout << "\tComputing checksum...\n";
      std::string current_hash = Util::hash_file(archive_path);
      if (current_hash == expected_hash) {
        needs_download = false;
      }
    }

    if (needs_download) {
      Util::shell_command("./curl -L -o " + archive_path + " " + src_url);

      std::cout << "\n\tComputing checksum...\n";
      std::string hash = Util::hash_file(archive_path);
      if (expected_hash != hash) {
        std::cerr << "Checksum check failed.\n";
        std::cerr << "Expected: " << expected_hash << "\n";
        std::cerr << "Got:      " << hash << "\n";
        std::exit(1);
      }
    }

    std::cout << "\n\tExtracting...\n";
    Util::shell_command("tar xf " + archive_path + " -C " + BUILD_PATH);
  } else {
    throw std::runtime_error("unrecognized src type");
  }

  std::string jobs = std::to_string(std::thread::hardware_concurrency());

  std::cout << "\n\tBuilding...\n\n";
  for (const auto &step : pkg.get("build").children) {
    std::string step_type = step.children[0].value;

    if (step_type == "configure_make") {
      std::string configure_flags;
      try {
        configure_flags = step.get_one("configure_flags").value;
      } catch (std::out_of_range &) {
      }

      bool do_install = true;
      try {
        if (step.get_one("do_install").value == "no") {
          do_install = false;
        }
      } catch (std::out_of_range &) {
      }

      Util::shell_command("./configure --prefix=/usr " + configure_flags,
                          src_path);
      Util::shell_command("make -j " + jobs, src_path);
      if (do_install) {
        Util::shell_command("make install", src_path);
      }
    } else if (step_type == "make") {
      bool do_install = true;
      try {
        if (step.get_one("do_install").value == "no") {
          do_install = false;
        }
      } catch (std::out_of_range &) {
      }

      std::string make_flags;
      try {
        make_flags = step.get_one("make_flags").value;
      } catch (std::out_of_range &) {
      }
      std::string install_flags;
      try {
        install_flags = step.get_one("install_flags").value;
      } catch (std::out_of_range &) {
      }

      Util::shell_command("make " + make_flags + " -j " + jobs, src_path);
      if (do_install) {
        Util::shell_command("make install " + install_flags, src_path);
      }
    } else if (step_type == "cmake") {
      std::string configure_flags;
      try {
        configure_flags = step.get_one("configure_flags").value;
      } catch (std::out_of_range &) {
      }

      Util::shell_command("mkdir -p build", src_path);
      Util::shell_command("cmake -DCMAKE_INSTALL_PREFIX=/usr " +
                              configure_flags + " ..",
                          src_path + "/build");
      Util::shell_command("make -j" + jobs, src_path + "/build");
      Util::shell_command("make install", src_path + "/build");
    } else if (step_type == "meson") {
      std::string configure_flags;
      try {
        configure_flags = step.get_one("configure_flags").value;
      } catch (std::out_of_range &) {
      }

      Util::shell_command(
          "meson setup --prefix=/usr builddir " + configure_flags, src_path);
      Util::shell_command("meson compile -C builddir -j" + jobs, src_path);
      Util::shell_command("meson install -C builddir", src_path);
    } else if (step_type == "shell") {
      Util::shell_command(step.children[1].value, src_path);
    } else {
      throw std::runtime_error("unrecognized step type");
    }
  }

  db.mark_installed(pkg.get_one("name").value, pkg.get_one("version").value);
}

void generate_graph(const std::vector<std::string> &to_highlight) {
  std::ofstream out(BUILD_PATH + "/graph.dot");

  out << "digraph dependencies {\n";
  if (to_highlight.empty()) {
    out << "    rankdir=LR;\n";
  }

  for (const auto &[pkg_name, pkg] : packages) {
    if (to_highlight.empty() ||
        std::find(to_highlight.begin(), to_highlight.end(), pkg_name) !=
            to_highlight.end()) {
      out << "    \"" << pkg_name << "\";\n";
      try {
        for (const auto &dep : pkg.get("dependencies").children) {
          out << "    \"" << pkg_name << "\" -> \"" << dep.value << "\";\n";
        }
      } catch (std::out_of_range &) {
      }
    }
  }
  out << "}\n";
  out.close();

  Util::shell_command("dot -Tpng graph.dot > graph.png", BUILD_PATH);
  std::cout << "Graph saved to " << BUILD_PATH << "/graph.png\n";
}

int main(int argc, char **argv) {
  auto args = std::span(argv, static_cast<size_t>(argc));

  if (geteuid() != 0) {
    std::cerr << "This program needs to be ran as root.\n";
    return 1;
  }

  std::filesystem::create_directories(BUILD_PATH);
  std::filesystem::create_directories(DB::DB_PATH);

  DB db;
  db.init_db();

  load_packages();

  bool flag_raw = false;
  bool flag_graph = false;
  std::vector<std::string> to_install;

  for (size_t i = 1; i < args.size(); ++i) {
    std::string arg = args[i];
    if (arg == "-r") {
      flag_raw = true;
    } else if (arg == "-g") {
      flag_graph = true;
    } else {
      to_install.push_back(arg);
    }
  }

  if (flag_graph) {
    generate_graph(to_install);
    return 0;
  }

  if (to_install.empty()) {
    std::cerr << "Usage: " << args[0] << " [-g] [-r] package1 [package2 ...]\n";
    return 1;
  }

  if (!flag_raw) {
    to_install = resolve_dependencies(to_install, db);
  }

  std::cout << "\nFollowing packages will be installed:\n";
  for (const std::string &pkg : to_install) {
    if (!packages.contains(pkg)) {
      std::cerr << "Package not found: " << pkg << "\n";
      std::exit(1);
    }
    std::cout << " - " << pkg << " (" << packages[pkg].get_one("version").value
              << ")" << std::endl;
  }
  std::cin.get();

  for (const std::string &pkg : to_install) {
    try {
      install_package(pkg, db);
    } catch (std::exception &e) {
      std::cerr << "ERROR: " << e.what() << std::endl;
      return 1;
    }
  }
}