#include "ir_gen.h"
#include "parser.h"
#include "serializer.h"
#include <cctype>
#include <cstring>
#include <filesystem>
#include <iostream>
#include <stdexcept>
#include <string>
#include <vector>

namespace nix_irc {
namespace fs = std::filesystem;

void print_usage(const char* prog) {
  std::cout << "Usage: " << prog << " [options] <input.nix|flake#attr> [output.nixir]\n"
            << "\nOptions:\n"
            << "  -I <path>    Add search path for imports\n"
            << "  --no-imports Disable import resolution\n"
            << "  --help       Show this help\n";
}

static bool is_flake_reference(const std::string& input) {
  return input.find('#') != std::string::npos;
}

static std::string sanitize_output_stem(const std::string& input) {
  std::string stem;
  stem.reserve(input.size());

  for (char ch : input) {
    if (std::isalnum(static_cast<unsigned char>(ch))) {
      stem.push_back(ch);
    } else if (stem.empty() || stem.back() != '-') {
      stem.push_back('-');
    }
  }

  while (!stem.empty() && stem.back() == '-') {
    stem.pop_back();
  }

  return stem.empty() ? "bundle" : stem;
}

static std::string default_output_path_for(const std::string& input) {
  if (!is_flake_reference(input)) {
    return input + "ir";
  }

  return sanitize_output_stem(input) + ".nixir";
}

static std::string normalize_local_flake_path(const std::string& raw_path) {
  fs::path path = raw_path.empty() ? fs::current_path() : fs::path(raw_path);
  fs::path absolute = path.is_absolute() ? path : fs::absolute(path);
  fs::path normalized = absolute.lexically_normal();

  if (!fs::exists(normalized)) {
    throw std::runtime_error("Flake path does not exist: " + normalized.string());
  }

  if (fs::is_directory(normalized) && !fs::exists(normalized / "flake.nix")) {
    throw std::runtime_error("Flake directory does not contain flake.nix: " + normalized.string());
  }

  return normalized.string();
}

static std::string normalize_flake_ref_source(const std::string& ref) {
  if (ref.empty()) {
    return normalize_local_flake_path(".");
  }

  if (ref.rfind("path:", 0) == 0) {
    return "path:" + normalize_local_flake_path(ref.substr(5));
  }

  if (ref[0] == '.' || ref[0] == '/') {
    return normalize_local_flake_path(ref);
  }

  if (fs::exists(ref)) {
    return normalize_local_flake_path(ref);
  }

  return ref;
}

static std::vector<std::string> parse_flake_attr_path(const std::string& raw_attr_path) {
  if (raw_attr_path.empty()) {
    throw std::runtime_error("Flake reference is missing an attribute path after '#'");
  }

  std::vector<std::string> segments;
  std::string current;
  bool in_quotes = false;
  bool escaping = false;

  for (char ch : raw_attr_path) {
    if (escaping) {
      current.push_back(ch);
      escaping = false;
      continue;
    }

    if (in_quotes) {
      if (ch == '\\') {
        escaping = true;
      } else if (ch == '"') {
        in_quotes = false;
      } else {
        current.push_back(ch);
      }
      continue;
    }

    if (ch == '"') {
      in_quotes = true;
    } else if (ch == '.') {
      if (current.empty()) {
        throw std::runtime_error("Flake attribute path contains an empty segment");
      }
      segments.push_back(current);
      current.clear();
    } else {
      current.push_back(ch);
    }
  }

  if (escaping || in_quotes) {
    throw std::runtime_error("Unterminated quoted segment in flake attribute path");
  }

  if (current.empty()) {
    throw std::runtime_error("Flake attribute path contains an empty segment");
  }

  segments.push_back(current);
  return segments;
}

static std::shared_ptr<Node> build_flake_ref_ast(const std::string& input) {
  size_t hash_pos = input.find('#');
  if (hash_pos == std::string::npos) {
    throw std::runtime_error("Not a flake reference: " + input);
  }

  std::string flake_source = normalize_flake_ref_source(input.substr(0, hash_pos));
  auto attr_path = parse_flake_attr_path(input.substr(hash_pos + 1));

  auto expr = std::make_shared<Node>(BuiltinCallNode(
      "getFlake",
      std::vector<std::shared_ptr<Node>>{std::make_shared<Node>(ConstStringNode(flake_source))}));

  for (const auto& attr : attr_path) {
    expr = std::make_shared<Node>(SelectNode(expr, std::make_shared<Node>(ConstStringNode(attr))));
  }

  return expr;
}

int run_compile(int argc, char** argv) {
  std::string input_file;
  std::string output_file;
  std::vector<std::string> search_paths;
  bool resolve_imports = true;

  int i = 1;
  while (i < argc) {
    std::string arg = argv[i];
    if (arg == "-I") {
      if (i + 1 >= argc) {
        std::cerr << "Error: -I requires a path argument\n";
        return 1;
      }
      search_paths.push_back(argv[++i]);
    } else if (arg == "--no-imports") {
      resolve_imports = false;
    } else if (arg == "--help" || arg == "-h") {
      print_usage(argv[0]);
      return 0;
    } else if (arg[0] != '-') {
      input_file = arg;
      if (i + 1 < argc && argv[i + 1][0] != '-') {
        output_file = argv[++i];
      }
    } else {
      std::cerr << "Unknown option: " << arg << "\n";
      print_usage(argv[0]);
      return 1;
    }
    i++;
  }

  if (input_file.empty()) {
    std::cerr << "Error: No input file specified\n";
    print_usage(argv[0]);
    return 1;
  }

  if (output_file.empty()) {
    output_file = default_output_path_for(input_file);
  }

  try {
    Parser parser;
    (void) search_paths;
    (void) resolve_imports;

    std::shared_ptr<Node> ast;

    if (is_flake_reference(input_file)) {
      std::cout << "Compiling flake reference: " << input_file << "\n";
      ast = build_flake_ref_ast(input_file);
    } else {
      std::cout << "Parsing: " << input_file << "\n";
      ast = parser.parse_file(input_file);
    }

    if (!ast) {
      std::cerr << "Error: Failed to parse input\n";
      return 1;
    }

    std::cout << "Resolving imports...\n";

    IRGenerator ir_gen;

    std::cout << "Generating IR...\n";
    auto ir = ir_gen.generate(ast);

    IRModule module;
    module.version = IR_VERSION;
    module.entry = ir;

    std::cout << "Serializing to: " << output_file << "\n";
    Serializer serializer;
    serializer.serialize(module, output_file);

    std::cout << "Done!\n";
    return 0;

  } catch (const std::exception& e) {
    std::cerr << "Error: " << e.what() << "\n";
    return 1;
  }
}

void print_decompile_usage(const char* prog) {
  std::cout << "Usage: " << prog << " decompile <input.nixir>\n";
}

int run_decompile(int argc, char** argv) {
  if (argc < 3) {
    print_decompile_usage(argv[0]);
    return 1;
  }

  std::string input_file = argv[2];

  try {
    Deserializer deserializer;
    auto module = deserializer.deserialize(input_file);

    std::cout << "IR Version: " << module.version << "\n";
    std::cout << "Sources: " << module.sources.size() << "\n";
    std::cout << "Imports: " << module.imports.size() << "\n";

    return 0;
  } catch (const std::exception& e) {
    std::cerr << "Error: " << e.what() << "\n";
    return 1;
  }
}

} // namespace nix_irc

int main(int argc, char** argv) {
  if (argc < 2) {
    nix_irc::print_usage(argv[0]);
    return 1;
  }

  std::string cmd = argv[1];

  if (cmd == "compile" || cmd == "c") {
    return nix_irc::run_compile(argc - 1, argv + 1);
  } else if (cmd == "decompile" || cmd == "d") {
    return nix_irc::run_decompile(argc, argv);
  } else if (cmd == "help" || cmd == "--help" || cmd == "-h") {
    nix_irc::print_usage(argv[0]);
    return 0;
  } else {
    return nix_irc::run_compile(argc, argv);
  }
}