#include <algorithm>
#include <cstddef>
#include <fstream>
#include <iostream>
#include <string>
#include <unordered_map>
#include <vector>

#include "compiler/jsonld_loader.h"
#include "compiler/signature.h"

namespace {

struct Options {
  std::vector<std::string> step_files;
  std::string out_path;
  std::string only_prefix;
};

bool parse_args(int argc, char** argv, Options& opts) {
  for (int i = 1; i < argc; ++i) {
    std::string arg = argv[i];
    if (arg == "--steps" && i + 1 < argc) {
      opts.step_files.push_back(argv[++i]);
    } else if (arg == "--out" && i + 1 < argc) {
      opts.out_path = argv[++i];
    } else if (arg == "--only" && i + 1 < argc) {
      opts.only_prefix = argv[++i];
    } else if (arg == "--help" || arg == "-h") {
      return false;
    }
  }
  return !opts.step_files.empty() && !opts.out_path.empty();
}

std::string sanitize_name(const std::string& iri) {
  size_t pos = iri.find_last_of("#/");
  std::string base = (pos == std::string::npos) ? iri : iri.substr(pos + 1);
  std::string out;
  out.reserve(base.size());
  for (char c : base) {
    if ((c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9')) {
      out.push_back(c);
    } else {
      out.push_back('_');
    }
  }
  if (!out.empty() && (out[0] >= '0' && out[0] <= '9')) {
    out = "Step_" + out;
  }
  return out;
}

std::string local_key_name(const std::string& key) {
  size_t pos = key.find_last_of("#/");
  if (pos != std::string::npos) return key.substr(pos + 1);
  pos = key.find_last_of(':');
  if (pos != std::string::npos) return key.substr(pos + 1);
  return key;
}

std::string sanitize_key(const std::string& key) {
  std::string base = local_key_name(key);
  std::string out;
  out.reserve(base.size());
  for (char c : base) {
    if ((c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9')) {
      out.push_back(c);
    } else {
      out.push_back('_');
    }
  }
  if (!out.empty() && (out[0] >= '0' && out[0] <= '9')) {
    out = "Key_" + out;
  }
  return out;
}

std::vector<std::string> unique_names(const std::vector<std::string>& keys) {
  std::unordered_map<std::string, int> counts;
  std::vector<std::string> out;
  out.reserve(keys.size());
  for (const auto& key : keys) {
    std::string base = sanitize_key(key);
    int count = counts[base]++;
    if (count > 0) {
      base += "_" + std::to_string(count);
    }
    out.push_back(base);
  }
  return out;
}

}  // namespace

int main(int argc, char** argv) {
  Options opts;
  if (!parse_args(argc, argv, opts)) {
    std::cerr << "Usage: cf_siggen --steps <path> [--steps <path> ...] --out <path> [--only <iri_prefix>]\n";
    return 2;
  }

  std::string error;
  cf::StepCatalog catalog = cf::load_step_catalog(opts.step_files, error);
  if (!error.empty()) {
    std::cerr << error << "\n";
    return 1;
  }

  std::vector<std::pair<std::string, std::string>> entries;
  for (const auto& it : catalog.steps) {
    const std::string& iri = it.first;
    if (!opts.only_prefix.empty()) {
      if (iri.rfind(opts.only_prefix, 0) != 0) {
        continue;
      }
    }
    auto spec = cf::build_signature_spec(it.second);
    entries.emplace_back(iri, cf::signature_hash(spec));
  }

  std::sort(entries.begin(), entries.end(), [](const auto& a, const auto& b) {
    return a.first < b.first;
  });

  std::unordered_map<std::string, int> name_counts;
  std::ofstream out(opts.out_path, std::ios::trunc);
  if (!out) {
    std::cerr << "Failed to open output: " << opts.out_path << "\n";
    return 1;
  }

  out << "#pragma once\n";
  out << "#include <cstddef>\n";
  out << "namespace cf_generated {\n";

  std::unordered_map<std::string, std::string> step_names;
  for (const auto& entry : entries) {
    std::string name = sanitize_name(entry.first);
    int count = name_counts[name]++;
    if (count > 0) {
      name += "_" + std::to_string(count);
    }
    step_names[entry.first] = name;
    out << "inline constexpr const char* k" << name << "Iri = \"" << entry.first << "\";\n";
    out << "inline constexpr const char* k" << name << "SigHash = \"" << entry.second << "\";\n";
  }

  out << "}  // namespace cf_generated\n";
  out << "namespace cf_sig {\n";
  for (const auto& entry : entries) {
    const std::string& iri = entry.first;
    auto step_it = catalog.steps.find(iri);
    if (step_it == catalog.steps.end()) {
      continue;
    }
    auto name_it = step_names.find(iri);
    if (name_it == step_names.end()) {
      continue;
    }

    std::vector<std::string> param_keys;
    param_keys.reserve(step_it->second.parameters.size());
    for (const auto& param : step_it->second.parameters) {
      param_keys.push_back(param.name);
    }
    std::sort(param_keys.begin(), param_keys.end());

    std::vector<std::string> input_keys;
    input_keys.reserve(step_it->second.inputs.size());
    for (const auto& input : step_it->second.inputs) {
      input_keys.push_back(input.key);
    }
    std::sort(input_keys.begin(), input_keys.end());

    std::vector<std::string> output_keys;
    output_keys.reserve(step_it->second.outputs.size());
    for (const auto& output : step_it->second.outputs) {
      output_keys.push_back(output.key);
    }
    std::sort(output_keys.begin(), output_keys.end());

    std::vector<std::string> param_names = unique_names(param_keys);
    std::vector<std::string> input_names = unique_names(input_keys);
    std::vector<std::string> output_names = unique_names(output_keys);

    out << "struct " << name_it->second << " {\n";
    out << "  static constexpr size_t kParamCount = " << param_names.size() << ";\n";
    for (size_t i = 0; i < param_names.size(); ++i) {
      out << "  static constexpr size_t P_" << param_names[i] << " = " << i << ";\n";
    }
    out << "  static constexpr size_t kInputCount = " << input_names.size() << ";\n";
    for (size_t i = 0; i < input_names.size(); ++i) {
      out << "  static constexpr size_t I_" << input_names[i] << " = " << i << ";\n";
    }
    out << "  static constexpr size_t kOutputCount = " << output_names.size() << ";\n";
    for (size_t i = 0; i < output_names.size(); ++i) {
      out << "  static constexpr size_t O_" << output_names[i] << " = " << i << ";\n";
    }
    out << "};\n";
  }
  out << "}  // namespace cf_sig\n";
  return 0;
}
