#include "compiler/plan_builder.h"

#include <nlohmann/json.hpp>
#include <queue>
#include <unordered_map>
#include <cstring>
#include <sstream>

#include "common/type_registry.h"
#include "compiler/signature.h"

namespace cf {

namespace {

using json = nlohmann::json;

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);
}

int find_port_index(const std::vector<PortDef>& ports, const std::string& key) {
  for (size_t i = 0; i < ports.size(); ++i) {
    if (ports[i].key == key) return static_cast<int>(i);
  }
  std::string key_local = local_name(key);
  for (size_t i = 0; i < ports.size(); ++i) {
    if (local_name(ports[i].key) == key_local) return static_cast<int>(i);
  }
  return -1;
}

CfValueHandle make_value_handle(ExecutionPlan& plan, CfStorageType storage, const std::string& json_value, bool& ok) {
  ok = true;
  json val;
  try {
    val = json::parse(json_value);
  } catch (...) {
    ok = false;
    return {};
  }

  auto alloc_bytes = [&](size_t size) -> uint8_t* {
    auto buffer = std::make_unique<uint8_t[]>(size);
    uint8_t* ptr = buffer.get();
    plan.owned_buffers.push_back(std::move(buffer));
    return ptr;
  };

  CfValueHandle handle{};
  handle.storage_type = storage;
  switch (storage) {
    case CF_STORAGE_BOOL: {
      if (!val.is_boolean()) {
        ok = false;
        return {};
      }
      uint8_t* data = alloc_bytes(sizeof(uint8_t));
      data[0] = val.get<bool>() ? 1 : 0;
      handle.data = data;
      handle.size = sizeof(uint8_t);
      break;
    }
    case CF_STORAGE_I64: {
      if (!val.is_number_integer()) {
        ok = false;
        return {};
      }
      int64_t* data = reinterpret_cast<int64_t*>(alloc_bytes(sizeof(int64_t)));
      *data = val.get<int64_t>();
      handle.data = data;
      handle.size = sizeof(int64_t);
      break;
    }
    case CF_STORAGE_F64: {
      if (!val.is_number()) {
        ok = false;
        return {};
      }
      double* data = reinterpret_cast<double*>(alloc_bytes(sizeof(double)));
      *data = val.get<double>();
      handle.data = data;
      handle.size = sizeof(double);
      break;
    }
    case CF_STORAGE_STRING:
    case CF_STORAGE_JSON:
    case CF_STORAGE_BYTES:
    case CF_STORAGE_OPAQUE: {
      std::string payload;
      if (storage == CF_STORAGE_JSON) {
        payload = val.dump();
      } else if (val.is_string()) {
        payload = val.get<std::string>();
      } else {
        payload = val.dump();
      }
      uint8_t* data = alloc_bytes(payload.size());
      if (!payload.empty()) {
        std::memcpy(data, payload.data(), payload.size());
      }
      handle.data = data;
      handle.size = payload.size();
      break;
    }
    case CF_STORAGE_VECTOR_F64:
      ok = false;
      return {};
    default:
      ok = false;
      break;
  }
  return handle;
}

bool validate_parameter_handle(const ParameterDef& param,
                               const CfValueHandle& handle,
                               const std::string& node_id,
                               std::string& error) {
  if (param.required) {
    if (!handle.data || handle.size == 0) {
      error = "Missing required parameter '" + param.name + "' for node " + node_id;
      return false;
    }
  }

  if (handle.data && handle.size > 0) {
    if (!param.allowed_values.empty() && handle.storage_type == CF_STORAGE_STRING) {
      std::string value(reinterpret_cast<const char*>(handle.data), handle.size);
      bool ok = false;
      for (const auto& allowed : param.allowed_values) {
        if (value == allowed) {
          ok = true;
          break;
        }
      }
      if (!ok) {
        std::ostringstream msg;
        msg << "Invalid value for parameter '" << param.name << "' on node " << node_id << ". Allowed: ";
        for (size_t i = 0; i < param.allowed_values.size(); ++i) {
          if (i > 0) msg << ", ";
          msg << param.allowed_values[i];
        }
        error = msg.str();
        return false;
      }
    }

    if ((param.min_value.has_value() || param.max_value.has_value()) &&
        (handle.storage_type == CF_STORAGE_I64 || handle.storage_type == CF_STORAGE_F64)) {
      double value = 0.0;
      if (handle.storage_type == CF_STORAGE_I64) {
        if (handle.size < sizeof(int64_t)) {
          error = "Failed to read integer parameter '" + param.name + "' for node " + node_id;
          return false;
        }
        value = static_cast<double>(*reinterpret_cast<const int64_t*>(handle.data));
      } else {
        if (handle.size < sizeof(double)) {
          error = "Failed to read float parameter '" + param.name + "' for node " + node_id;
          return false;
        }
        value = *reinterpret_cast<const double*>(handle.data);
      }

      if (param.min_value.has_value() && value < param.min_value.value()) {
        error = "Parameter '" + param.name + "' on node " + node_id + " must be >= " +
                std::to_string(param.min_value.value());
        return false;
      }
      if (param.max_value.has_value() && value > param.max_value.value()) {
        error = "Parameter '" + param.name + "' on node " + node_id + " must be <= " +
                std::to_string(param.max_value.value());
        return false;
      }
    }
  }

  return true;
}

bool build_levels(size_t node_count, const std::vector<EdgeSlot>& edges, std::vector<std::vector<int>>& levels) {
  std::vector<int> indegree(node_count, 0);
  std::vector<std::vector<int>> adj(node_count);
  for (const auto& edge : edges) {
    if (edge.from_node >= 0 && edge.to_node >= 0) {
      indegree[edge.to_node] += 1;
      adj[edge.from_node].push_back(edge.to_node);
    }
  }
  std::queue<int> q;
  for (size_t i = 0; i < node_count; ++i) {
    if (indegree[i] == 0) q.push(static_cast<int>(i));
  }
  size_t visited = 0;
  while (!q.empty()) {
    size_t level_size = q.size();
    std::vector<int> level;
    level.reserve(level_size);
    for (size_t i = 0; i < level_size; ++i) {
      int node = q.front();
      q.pop();
      level.push_back(node);
      visited++;
      for (int next : adj[node]) {
        indegree[next] -= 1;
        if (indegree[next] == 0) {
          q.push(next);
        }
      }
    }
    levels.push_back(level);
  }
  return visited == node_count;
}

}  // namespace

bool build_execution_plan(
  const PipelineSpec& pipeline,
  const StepCatalog& catalog,
  PluginRegistry& plugins,
  ExecutionPlan& out_plan,
  std::string& error,
  const PlanBuildOptions& options) {

  out_plan = ExecutionPlan{};

  TypeRegistry type_registry;
  if (!load_default_type_registry(type_registry, error)) {
    return false;
  }

  std::unordered_map<std::string, int> node_index;
  for (size_t i = 0; i < pipeline.nodes.size(); ++i) {
    node_index[pipeline.nodes[i].id] = static_cast<int>(i);
  }

  out_plan.nodes.reserve(pipeline.nodes.size());
  for (const auto& node : pipeline.nodes) {
    auto step_it = catalog.steps.find(node.step_iri);
    if (step_it == catalog.steps.end()) {
      error = "Step not found in catalog: " + node.step_iri;
      return false;
    }
    PlanNode plan_node;
    plan_node.node_id = node.id;
    plan_node.step_iri = node.step_iri;
    plan_node.step_def = step_it->second;

    std::string resolve_error;
    plan_node.vtable = plugins.resolve_step(node.step_iri, resolve_error);
    if (!plan_node.vtable) {
      error = resolve_error;
      return false;
    }
    auto spec = build_signature_spec(plan_node.step_def);
    std::string expected = signature_hash(spec);
    if (plan_node.vtable->signature_hash == nullptr || expected != plan_node.vtable->signature_hash) {
      error = "Signature mismatch for step: " + node.step_iri;
      return false;
    }

    for (const auto& param : plan_node.step_def.parameters) {
      auto provided = node.parameters_json.find(param.name);
      std::string value_json;
      if (provided != node.parameters_json.end()) {
        value_json = provided->second;
      } else {
        auto by_iri = node.parameters_json_by_iri.find(param.iri);
        if (by_iri != node.parameters_json_by_iri.end()) {
          value_json = by_iri->second;
        } else if (param.has_default) {
          value_json = param.default_json;
        } else if (param.required) {
          error = "Missing required parameter '" + param.name + "' for node " + node.id;
          return false;
        } else {
          value_json = "null";
        }
      }
      bool ok = false;
      CfValueHandle handle = make_value_handle(out_plan, param.storage_type, value_json, ok);
      if (!ok) {
        error = "Failed to parse parameter '" + param.name + "' for node " + node.id;
        return false;
      }
      if (!validate_parameter_handle(param, handle, node.id, error)) {
        return false;
      }
      plan_node.params.push_back(handle);
    }
    out_plan.nodes.push_back(std::move(plan_node));
  }

  out_plan.edges.reserve(pipeline.connections.size());
  for (size_t i = 0; i < pipeline.connections.size(); ++i) {
    const auto& conn = pipeline.connections[i];
    auto from_it = node_index.find(conn.from_node);
    auto to_it = node_index.find(conn.to_node);
    if (from_it == node_index.end() || to_it == node_index.end()) {
      error = "Connection references unknown node";
      return false;
    }
    PlanNode& from_node = out_plan.nodes[from_it->second];
    PlanNode& to_node = out_plan.nodes[to_it->second];

    int from_idx = find_port_index(from_node.step_def.outputs, conn.from_port_key);
    int to_idx = find_port_index(to_node.step_def.inputs, conn.to_port_key);
    if (from_idx < 0) {
      error = "Unknown output port key '" + conn.from_port_key + "' on node '" + conn.from_node + "'";
      return false;
    }
    if (to_idx < 0) {
      error = "Unknown input port key '" + conn.to_port_key + "' on node '" + conn.to_node + "'";
      return false;
    }

    const PortDef& from_port = from_node.step_def.outputs[from_idx];
    const PortDef& to_port = to_node.step_def.inputs[to_idx];
    if (!from_port.type_iri.empty() && !to_port.type_iri.empty() &&
        !type_registry.are_port_types_compatible(from_port.type_iri, to_port.type_iri)) {
      std::ostringstream msg;
      msg << "Incompatible connection from node '" << conn.from_node << "' port '" << conn.from_port_key
          << "' to node '" << conn.to_node << "' port '" << conn.to_port_key << "'"
          << " (output type '" << from_port.type_iri << "', input type '" << to_port.type_iri << "')";
      if (from_port.type_iri == "http://www.w3.org/1999/02/22-rdf-syntax-ns#JSON" &&
          to_port.type_iri == "http://www.w3.org/2001/XMLSchema#double") {
        msg << ". Insert step 'http://cogniflow.org/basic-io#JsonScalarToDoubleStep' between these ports.";
      }
      error = msg.str();
      return false;
    }

    EdgeSlot edge;
    edge.id = out_plan.edges.size();
    edge.from_node = from_it->second;
    edge.from_output = from_idx;
    edge.to_node = to_it->second;
    edge.to_input = to_idx;
    edge.label = conn.to_port_label.empty() ? local_name(conn.from_port_key) : conn.to_port_label;
    edge.from_port_key = conn.from_port_key;
    out_plan.edges.push_back(edge);
  }

  for (size_t i = 0; i < out_plan.edges.size(); ++i) {
    const auto& edge = out_plan.edges[i];
    out_plan.nodes[edge.from_node].output_edge_ids.push_back(static_cast<int>(i));
    out_plan.nodes[edge.to_node].input_edge_ids.push_back(static_cast<int>(i));
  }

  if (!build_levels(out_plan.nodes.size(), out_plan.edges, out_plan.levels)) {
    error = "Pipeline has cycles (DAG check failed)";
    return false;
  }

  if (!options.allow_missing_inputs) {
    for (const auto& node : out_plan.nodes) {
      std::vector<bool> has_input(node.step_def.inputs.size(), false);
      for (int edge_id : node.input_edge_ids) {
        if (edge_id >= 0 && static_cast<size_t>(edge_id) < out_plan.edges.size()) {
          const auto& edge = out_plan.edges[edge_id];
          if (edge.to_input >= 0 && static_cast<size_t>(edge.to_input) < has_input.size()) {
            has_input[edge.to_input] = true;
          }
        }
      }
      for (size_t idx = 0; idx < has_input.size(); ++idx) {
        if (!has_input[idx]) {
          error = "Missing required input connection on node " + node.node_id;
          return false;
        }
      }
    }
  }

  return true;
}

}  // namespace cf

