#include "compiler/jsonld_loader.h"

#include <algorithm>
#include <cctype>
#include <fstream>
#include <nlohmann/json.hpp>
#include <sstream>

#include "common/type_registry.h"

namespace cf {

namespace {

using json = nlohmann::json;

bool load_json_file(const std::string& path, json& out, std::string& error) {
  std::ifstream fh(path);
  if (!fh) {
    error = "Failed to open JSON-LD file: " + path;
    return false;
  }
  try {
    fh >> out;
  } catch (const std::exception& ex) {
    error = std::string("Failed to parse JSON-LD file: ") + ex.what();
    return false;
  }
  return true;
}

JsonLdContext parse_context(const json& ctx_node, const std::string& base_dir, std::string& error) {
  JsonLdContext ctx;
  if (ctx_node.is_object()) {
    for (auto it = ctx_node.begin(); it != ctx_node.end(); ++it) {
      if (it.value().is_string()) {
        ctx.prefixes[it.key()] = it.value().get<std::string>();
      }
    }
    return ctx;
  }
  if (!ctx_node.is_array()) {
    return ctx;
  }
  for (const auto& entry : ctx_node) {
    if (entry.is_object()) {
      for (auto it = entry.begin(); it != entry.end(); ++it) {
        if (it.value().is_string()) {
          ctx.prefixes[it.key()] = it.value().get<std::string>();
        }
      }
    } else if (entry.is_string()) {
      std::string ref = entry.get<std::string>();
      std::string path = ref;
      if (!base_dir.empty() && ref.rfind("http", 0) != 0 && ref.rfind("https", 0) != 0) {
        path = base_dir + "/" + ref;
      }
      json ref_json;
      if (!load_json_file(path, ref_json, error)) {
        return ctx;
      }
      if (ref_json.contains("@context")) {
        JsonLdContext nested = parse_context(ref_json["@context"], base_dir, error);
        ctx.prefixes.insert(nested.prefixes.begin(), nested.prefixes.end());
      }
    }
  }
  return ctx;
}

std::string dirname_of(const std::string& path) {
  size_t pos = path.find_last_of("/\\");
  if (pos == std::string::npos) return "";
  return path.substr(0, pos);
}

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

std::string to_lower_copy(std::string value) {
  std::transform(value.begin(), value.end(), value.begin(), [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
  return value;
}

ParallelismModel parse_parallelism_model_iri(const std::string& iri) {
  std::string name = to_lower_copy(local_name(iri));
  if (name == "none" || name == "serial") return ParallelismModel::kNone;
  if (name == "omp" || name == "openmp") return ParallelismModel::kOmp;
  if (name == "tbb") return ParallelismModel::kTbb;
  if (name == "blas") return ParallelismModel::kBlas;
  if (name == "mkl") return ParallelismModel::kMkl;
  if (name == "cuda") return ParallelismModel::kCuda;
  if (name == "external") return ParallelismModel::kExternal;
  return ParallelismModel::kUnknown;
}

std::optional<int> read_int_opt(const json& node) {
  if (node.is_number_integer()) {
    return node.get<int>();
  }
  if (node.is_object() && node.contains("@value") && node["@value"].is_number_integer()) {
    return node["@value"].get<int>();
  }
  return std::nullopt;
}

std::optional<bool> read_bool_opt(const json& node) {
  if (node.is_boolean()) {
    return node.get<bool>();
  }
  if (node.is_object() && node.contains("@value") && node["@value"].is_boolean()) {
    return node["@value"].get<bool>();
  }
  return std::nullopt;
}

std::optional<double> read_double_opt(const json& node) {
  if (node.is_number()) {
    return node.get<double>();
  }
  if (node.is_object() && node.contains("@value") && node["@value"].is_number()) {
    return node["@value"].get<double>();
  }
  if (node.is_string()) {
    try {
      return std::stod(node.get<std::string>());
    } catch (...) {
      return std::nullopt;
    }
  }
  if (node.is_object() && node.contains("@value") && node["@value"].is_string()) {
    try {
      return std::stod(node["@value"].get<std::string>());
    } catch (...) {
      return std::nullopt;
    }
  }
  return std::nullopt;
}

std::vector<std::string> read_string_list(const json& node) {
  std::vector<std::string> out;
  if (node.is_string()) {
    out.push_back(node.get<std::string>());
    return out;
  }
  if (node.is_object() && node.contains("@value") && node["@value"].is_string()) {
    out.push_back(node["@value"].get<std::string>());
    return out;
  }
  if (!node.is_array()) {
    return out;
  }
  for (const auto& entry : node) {
    if (entry.is_string()) {
      out.push_back(entry.get<std::string>());
    } else if (entry.is_object() && entry.contains("@value") && entry["@value"].is_string()) {
      out.push_back(entry["@value"].get<std::string>());
    }
  }
  return out;
}

std::vector<std::string> to_id_list(const json& node) {
  std::vector<std::string> ids;
  if (node.is_array()) {
    for (const auto& entry : node) {
      if (entry.is_object() && entry.contains("@id")) {
        ids.push_back(entry["@id"].get<std::string>());
      }
    }
  } else if (node.is_object()) {
    if (node.contains("@id")) {
      ids.push_back(node["@id"].get<std::string>());
    }
  }
  return ids;
}

bool type_matches(const JsonLdContext& ctx, const json& type_node, const std::string& expected) {
  auto matches = [&](const std::string& t) {
    return expand_iri(ctx, t) == expand_iri(ctx, expected);
  };
  if (type_node.is_string()) {
    return matches(type_node.get<std::string>());
  }
  if (type_node.is_array()) {
    for (const auto& entry : type_node) {
      if (entry.is_string() && matches(entry.get<std::string>())) {
        return true;
      }
    }
  }
  return false;
}

std::string extract_id_or_string(const json& node) {
  if (node.is_string()) {
    return node.get<std::string>();
  }
  if (node.is_object() && node.contains("@id") && node["@id"].is_string()) {
    return node["@id"].get<std::string>();
  }
  return {};
}

const json* resolve_referenced_node(const json& node, const std::unordered_map<std::string, json>& by_id) {
  if (!node.is_object()) {
    return nullptr;
  }
  if (node.contains("@id") && node["@id"].is_string()) {
    auto it = by_id.find(node["@id"].get<std::string>());
    if (it != by_id.end()) {
      return &it->second;
    }
  }
  if (node.contains("@type") || node.size() > 1) {
    return &node;
  }
  return nullptr;
}

bool apply_vector_contract_bridge(
  const JsonLdContext& ctx,
  const json& port,
  const std::unordered_map<std::string, json>& by_id,
  const TypeRegistry& type_registry,
  const std::string& step_iri,
  const std::string& port_iri,
  const std::string& semantic_type_iri,
  CfStorageType& storage_type,
  std::string& error) {

  auto contract_it = port.find("cf:valueContract");
  if (contract_it == port.end()) {
    return true;
  }

  const std::string contract_ref = extract_id_or_string(*contract_it);
  const json* contract = resolve_referenced_node(*contract_it, by_id);
  if (!contract) {
    error = "Could not resolve cf:valueContract '" +
            (contract_ref.empty() ? std::string("<inline>") : contract_ref) +
            "' on step '" + step_iri + "' port '" + port_iri + "'";
    return false;
  }

  const std::string contract_ref_label =
    contract_ref.empty() ? (contract->contains("@id") ? (*contract)["@id"].get<std::string>() : "<inline>") : contract_ref;

  const std::string shape_kind = expand_iri(ctx, extract_id_or_string(contract->value("cf:shapeKind", json::object())));
  const std::string vector_shape_iri = expand_iri(ctx, "cf:shape_vector");
  const std::string raw_contract_element = extract_id_or_string(contract->value("cf:elementType", json::object()));

  ResolvedDatatype resolved_contract_type;
  bool has_contract_type = false;
  if (!raw_contract_element.empty()) {
    const std::string contract_element_iri = expand_iri(ctx, raw_contract_element);
    if (!type_registry.resolve_port_type(contract_element_iri, resolved_contract_type)) {
      error = "Unsupported cf:valueContract/cf:elementType '" + contract_element_iri +
              "' on step '" + step_iri + "' port '" + port_iri + "'";
      return false;
    }
    has_contract_type = true;
    if (!semantic_type_iri.empty() && resolved_contract_type.canonical_port_type_iri != semantic_type_iri) {
      error = "Contract mismatch on step '" + step_iri + "' port '" + port_iri +
              "': cf:portType '" + semantic_type_iri +
              "' does not match cf:valueContract/cf:elementType '" +
              resolved_contract_type.canonical_port_type_iri +
              "' in contract '" + contract_ref_label + "'";
      return false;
    }
  }

  if (shape_kind != vector_shape_iri) {
    return true;
  }

  if (semantic_type_iri.empty()) {
    error = "Unsupported vector valueContract on step '" + step_iri + "' port '" + port_iri +
            "': cf:portType is required when cf:shapeKind is cf:shape_vector";
    return false;
  }
  if (!has_contract_type) {
    error = "Unsupported vector valueContract on step '" + step_iri + "' port '" + port_iri +
            "': cf:valueContract/cf:elementType is required for contract '" + contract_ref_label + "'";
    return false;
  }
  if (resolved_contract_type.canonical_port_type_iri != "http://www.w3.org/2001/XMLSchema#double") {
    error = "Unsupported vector valueContract on step '" + step_iri + "' port '" + port_iri +
            "': cf:valueContract/cf:elementType '" + resolved_contract_type.canonical_port_type_iri +
            "' must be xsd:double for cf:shape_vector";
    return false;
  }

  storage_type = CF_STORAGE_VECTOR_F64;
  return true;
}

}  // namespace

std::string expand_iri(const JsonLdContext& ctx, const std::string& value) {
  if (value.rfind("http://", 0) == 0 || value.rfind("https://", 0) == 0) {
    return value;
  }
  auto pos = value.find(':');
  if (pos == std::string::npos) {
    return value;
  }
  std::string prefix = value.substr(0, pos);
  std::string suffix = value.substr(pos + 1);
  auto it = ctx.prefixes.find(prefix);
  if (it == ctx.prefixes.end()) {
    return value;
  }
  return it->second + suffix;
}

StepCatalog load_step_catalog(const std::vector<std::string>& step_files, std::string& error) {
  StepCatalog catalog;
  TypeRegistry type_registry;
  if (!load_default_type_registry(type_registry, error)) {
    return catalog;
  }

  for (const auto& path : step_files) {
    json doc;
    if (!load_json_file(path, doc, error)) {
      return catalog;
    }
    if (!doc.contains("@context") || !doc.contains("@graph")) {
      error = "JSON-LD file missing @context or @graph: " + path;
      return catalog;
    }
    std::string base_dir = dirname_of(path);
    JsonLdContext ctx = parse_context(doc["@context"], base_dir, error);
    if (!error.empty()) {
      return catalog;
    }

    std::unordered_map<std::string, json> by_id;
    for (const auto& entry : doc["@graph"]) {
      if (entry.contains("@id")) {
        std::string id = entry["@id"].get<std::string>();
        by_id[id] = entry;
      }
    }

    for (const auto& entry : doc["@graph"]) {
      if (!entry.contains("@type") || !entry.contains("@id")) continue;
      if (!type_matches(ctx, entry["@type"], "cf:ProcessingStep")) continue;

      StepDef step;
      step.iri = expand_iri(ctx, entry["@id"].get<std::string>());

      if (entry.contains("cf:parallelismModel")) {
        const auto& model_node = entry["cf:parallelismModel"];
        std::string model_raw;
        if (model_node.is_object() && model_node.contains("@id") && model_node["@id"].is_string()) {
          model_raw = model_node["@id"].get<std::string>();
        } else if (model_node.is_string()) {
          model_raw = model_node.get<std::string>();
        }
        if (!model_raw.empty()) {
          step.parallelism.model = parse_parallelism_model_iri(expand_iri(ctx, model_raw));
        }
      }
      if (entry.contains("cf:parallelismPreferredThreads")) {
        step.parallelism.preferred_threads = read_int_opt(entry["cf:parallelismPreferredThreads"]);
      }
      if (entry.contains("cf:parallelismMaxThreads")) {
        step.parallelism.max_threads = read_int_opt(entry["cf:parallelismMaxThreads"]);
      }
      if (entry.contains("cf:parallelismDeterministic")) {
        step.parallelism.deterministic = read_bool_opt(entry["cf:parallelismDeterministic"]);
      }
      if (entry.contains("cf:parallelismSchemaVersion")) {
        step.parallelism.schema_version = read_int_opt(entry["cf:parallelismSchemaVersion"]);
      }

      std::vector<std::string> param_ids;
      if (entry.contains("cf:hasParameter")) {
        auto ids = to_id_list(entry["cf:hasParameter"]);
        param_ids.insert(param_ids.end(), ids.begin(), ids.end());
      }
      std::vector<std::string> input_ids;
      if (entry.contains("cf:hasInput")) {
        auto ids = to_id_list(entry["cf:hasInput"]);
        input_ids.insert(input_ids.end(), ids.begin(), ids.end());
      }
      std::vector<std::string> output_ids;
      if (entry.contains("cf:hasOutput")) {
        auto ids = to_id_list(entry["cf:hasOutput"]);
        output_ids.insert(output_ids.end(), ids.begin(), ids.end());
      }

      std::unordered_map<std::string, bool> required_params;
      if (entry.contains("cf:requiresParameter")) {
        for (const auto& req_id : to_id_list(entry["cf:requiresParameter"])) {
          required_params[req_id] = true;
        }
      }

      for (const auto& pid : param_ids) {
        auto it = by_id.find(pid);
        if (it == by_id.end()) continue;
        const json& param = it->second;
        ParameterDef def;
        def.iri = expand_iri(ctx, pid);
        if (param.contains("cf:paramKey")) {
          const auto& key_node = param["cf:paramKey"];
          if (key_node.is_object() && key_node.contains("@id")) {
            def.name = expand_iri(ctx, key_node["@id"].get<std::string>());
          } else if (key_node.is_string()) {
            def.name = key_node.get<std::string>();
          }
        }
        if (def.name.empty() && param.contains("cf:parameterName")) {
          def.name = param["cf:parameterName"].get<std::string>();
        }
        std::string raw_type = extract_id_or_string(param.value("cf:parameterType", json::object()));
        if (raw_type.empty()) {
          def.storage_type = CF_STORAGE_OPAQUE;
          def.type_iri.clear();
        } else {
          raw_type = expand_iri(ctx, raw_type);
          ResolvedDatatype resolved_type;
          if (!type_registry.resolve_parameter_type(raw_type, resolved_type)) {
            error = "Unsupported parameter datatype '" + raw_type + "' on step '" + step.iri +
                    "' parameter '" + (def.name.empty() ? def.iri : def.name) + "'";
            return catalog;
          }
          def.storage_type = resolved_type.storage_type;
          def.type_iri = resolved_type.canonical_parameter_type_iri;
        }
        if (param.contains("cf:defaultValue")) {
          def.default_json = param["cf:defaultValue"].dump();
          def.has_default = true;
        }
        if (param.contains("cf:minValue")) {
          def.min_value = read_double_opt(param["cf:minValue"]);
        }
        if (param.contains("cf:maxValue")) {
          def.max_value = read_double_opt(param["cf:maxValue"]);
        }
        if (param.contains("cf:allowedValue")) {
          def.allowed_values = read_string_list(param["cf:allowedValue"]);
        }
        def.required = required_params.find(pid) != required_params.end();
        step.parameters.push_back(def);
      }

      for (const auto& iid : input_ids) {
        auto it = by_id.find(iid);
        if (it == by_id.end()) continue;
        const json& port = it->second;
        PortDef def;
        def.iri = expand_iri(ctx, iid);
        if (port.contains("cf:portKey") && port["cf:portKey"].contains("@id")) {
          def.key = expand_iri(ctx, port["cf:portKey"]["@id"].get<std::string>());
        }
        std::string raw_type = extract_id_or_string(port.value("cf:portType", json::object()));
        if (raw_type.empty()) {
          def.storage_type = CF_STORAGE_OPAQUE;
          def.type_iri.clear();
        } else {
          raw_type = expand_iri(ctx, raw_type);
          ResolvedDatatype resolved_type;
          if (!type_registry.resolve_port_type(raw_type, resolved_type)) {
            error = "Unsupported input port datatype '" + raw_type + "' on step '" + step.iri +
                    "' port '" + def.iri + "'";
            return catalog;
          }
          def.storage_type = resolved_type.storage_type;
          def.type_iri = resolved_type.canonical_port_type_iri;
          if (!apply_vector_contract_bridge(ctx, port, by_id, type_registry, step.iri, def.iri, def.type_iri,
                                            def.storage_type, error)) {
            return catalog;
          }
        }
        if (port.contains("cf:isMultiPort") && port["cf:isMultiPort"].is_boolean()) {
          def.is_multi = port["cf:isMultiPort"].get<bool>();
        }
        step.inputs.push_back(def);
      }

      for (const auto& oid : output_ids) {
        auto it = by_id.find(oid);
        if (it == by_id.end()) continue;
        const json& port = it->second;
        PortDef def;
        def.iri = expand_iri(ctx, oid);
        if (port.contains("cf:portKey") && port["cf:portKey"].contains("@id")) {
          def.key = expand_iri(ctx, port["cf:portKey"]["@id"].get<std::string>());
        }
        std::string raw_type = extract_id_or_string(port.value("cf:portType", json::object()));
        if (raw_type.empty()) {
          def.storage_type = CF_STORAGE_OPAQUE;
          def.type_iri.clear();
        } else {
          raw_type = expand_iri(ctx, raw_type);
          ResolvedDatatype resolved_type;
          if (!type_registry.resolve_port_type(raw_type, resolved_type)) {
            error = "Unsupported output port datatype '" + raw_type + "' on step '" + step.iri +
                    "' port '" + def.iri + "'";
            return catalog;
          }
          def.storage_type = resolved_type.storage_type;
          def.type_iri = resolved_type.canonical_port_type_iri;
          if (!apply_vector_contract_bridge(ctx, port, by_id, type_registry, step.iri, def.iri, def.type_iri,
                                            def.storage_type, error)) {
            return catalog;
          }
        }
        if (port.contains("cf:isMultiPort") && port["cf:isMultiPort"].is_boolean()) {
          def.is_multi = port["cf:isMultiPort"].get<bool>();
        }
        step.outputs.push_back(def);
      }

      std::sort(step.parameters.begin(), step.parameters.end(),
                [](const ParameterDef& a, const ParameterDef& b) { return a.name < b.name; });
      std::sort(step.inputs.begin(), step.inputs.end(),
                [](const PortDef& a, const PortDef& b) { return a.key < b.key; });
      std::sort(step.outputs.begin(), step.outputs.end(),
                [](const PortDef& a, const PortDef& b) { return a.key < b.key; });

      catalog.steps[step.iri] = step;
    }
  }
  return catalog;
}

}  // namespace cf

