#include <chrono>
#include <cmath>
#include <cstdlib>
#include <filesystem>
#include <fstream>
#include <iostream>
#include <string>
#include <vector>

#include "compiler/jsonld_loader.h"
#include "compiler/pipeline_parser.h"
#include "compiler/plan_builder.h"
#include "cf_vector_f64.h"
#include "plugin_loader/plugin_loader.h"
#include "scheduler/scheduler.h"

namespace {

struct TestContext {
  int failures = 0;

  void expect(bool condition, const std::string& message) {
    if (!condition) {
      std::cerr << "FAIL: " << message << "\n";
      failures++;
    }
  }
};

bool load_pipeline(const std::string& path, cf::PipelineSpec& out, TestContext& ctx) {
  std::string error;
  if (!cf::load_pipeline_spec(path, out, error)) {
    ctx.expect(false, error);
    return false;
  }
  return true;
}

bool load_catalog(const std::string& path, cf::StepCatalog& out, TestContext& ctx) {
  std::string error;
  out = cf::load_step_catalog({path}, error);
  if (!error.empty()) {
    ctx.expect(false, error);
    return false;
  }
  return true;
}

bool load_plugins(const std::vector<std::string>& dirs, cf::PluginRegistry& out, TestContext& ctx) {
  std::string error;
  if (!out.load_from_directories(dirs, error, true)) {
    ctx.expect(false, error);
    return false;
  }
  return true;
}

std::string test_dir() {
  namespace fs = std::filesystem;
  fs::path base = fs::current_path() / "tests";
  return base.lexically_normal().string();
}

std::string basic_io_steps() {
  namespace fs = std::filesystem;
  fs::path steps =
      fs::current_path() / ".." / ".." / "cf_basic_steps" / "cf_basic_io" / "src" / "cf_basic_io" / "steps.jsonld";
  return steps.lexically_normal().string();
}

std::string basic_io_plugins() {
  namespace fs = std::filesystem;
  fs::path plugins =
      fs::current_path() / ".." / ".." / "cf_basic_steps" / "cf_basic_io" / "src" / "cf_basic_io" / "bin";
  return plugins.lexically_normal().string();
}

std::string basic_dev_steps() {
  namespace fs = std::filesystem;
  fs::path steps =
      fs::current_path() / ".." / ".." / "cf_basic_steps" / "cf_basic_dev" / "src" / "cf_basic_dev" / "steps.jsonld";
  return steps.lexically_normal().string();
}

std::string basic_dev_plugins() {
  namespace fs = std::filesystem;
  fs::path plugins =
      fs::current_path() / ".." / ".." / "cf_basic_steps" / "cf_basic_dev" / "src" / "cf_basic_dev" / "bin";
  return plugins.lexically_normal().string();
}

std::string basic_signal_steps() {
  namespace fs = std::filesystem;
  fs::path steps =
      fs::current_path() / ".." / ".." / "cf_basic_steps" / "cf_basic_signal" / "src" / "cf_basic_signal" / "steps.jsonld";
  return steps.lexically_normal().string();
}

std::string basic_signal_plugins() {
  namespace fs = std::filesystem;
  fs::path plugins =
      fs::current_path() / ".." / ".." / "cf_basic_steps" / "cf_basic_signal" / "src" / "cf_basic_signal" / "bin";
  return plugins.lexically_normal().string();
}

struct AllocatorState {
  std::vector<void*> allocations;
};

void* test_alloc(size_t size, void* user_data) {
  void* mem = std::malloc(size);
  if (mem && user_data) {
    static_cast<AllocatorState*>(user_data)->allocations.push_back(mem);
  }
  return mem;
}

void test_free(void* ptr, void*) {
  std::free(ptr);
}

void free_allocations(AllocatorState& state) {
  for (void* ptr : state.allocations) {
    std::free(ptr);
  }
  state.allocations.clear();
}

bool read_f64(const CfValueHandle& handle, double& out) {
  if (handle.storage_type != CF_STORAGE_F64 || !handle.data || handle.size < sizeof(double)) {
    return false;
  }
  out = *reinterpret_cast<const double*>(handle.data);
  return true;
}

std::string write_temp_step_file(const std::string& stem, const std::string& contents, TestContext& ctx) {
  namespace fs = std::filesystem;
  fs::path path = fs::temp_directory_path() /
                  ("cf_pipeline_vector_contract_" + stem + "_" +
                   std::to_string(std::chrono::steady_clock::now().time_since_epoch().count()) + ".jsonld");
  std::ofstream fh(path);
  if (!fh) {
    ctx.expect(false, "Failed to create temporary step manifest: " + path.string());
    return {};
  }
  fh << contents;
  return path.lexically_normal().string();
}

bool read_vector_f64_values(const CfValueHandle& handle, std::vector<double>& out) {
  if (handle.storage_type != CF_STORAGE_VECTOR_F64 || !handle.data || handle.size < sizeof(CfVectorF64)) {
    return false;
  }
  const CfVectorF64* vector = reinterpret_cast<const CfVectorF64*>(handle.data);
  if (!vector || !vector->data) {
    return false;
  }
  out.assign(vector->data, vector->data + vector->length);
  return true;
}

const cf::StepDef* find_step_def(const cf::StepCatalog& catalog, const std::string& step_iri) {
  auto it = catalog.steps.find(step_iri);
  if (it == catalog.steps.end()) {
    return nullptr;
  }
  return &it->second;
}

bool find_parameter_index(
  const cf::StepCatalog& catalog,
  const std::string& step_iri,
  const std::string& parameter_name,
  size_t& out_index) {
  const cf::StepDef* step = find_step_def(catalog, step_iri);
  if (!step) {
    return false;
  }
  for (size_t i = 0; i < step->parameters.size(); ++i) {
    if (step->parameters[i].name == parameter_name) {
      out_index = i;
      return true;
    }
  }
  return false;
}

const cf::EdgeSlot* find_edge(
  const cf::ExecutionPlan& plan,
  const std::string& from_node_id,
  const std::string& to_node_id) {
  int from_idx = -1;
  int to_idx = -1;
  for (size_t i = 0; i < plan.nodes.size(); ++i) {
    if (plan.nodes[i].node_id == from_node_id) from_idx = static_cast<int>(i);
    if (plan.nodes[i].node_id == to_node_id) to_idx = static_cast<int>(i);
  }
  if (from_idx < 0 || to_idx < 0) {
    return nullptr;
  }
  for (const auto& edge : plan.edges) {
    if (edge.from_node == from_idx && edge.to_node == to_idx) {
      return &edge;
    }
  }
  return nullptr;
}

}  // namespace

int main() {
  TestContext ctx;
  const std::string base = test_dir();

  {
    std::string error;
    cf::StepCatalog catalog = cf::load_step_catalog({base + "/test_steps_invalid_type.jsonld"}, error);
    (void)catalog;
    ctx.expect(
      !error.empty() &&
      error.find("Unsupported parameter datatype") != std::string::npos &&
      error.find("http://www.w3.org/2001/XMLSchema#decimal") != std::string::npos,
      "Unknown step datatype should be rejected with the offending datatype IRI");
  }

  {
    cf::StepCatalog catalog;
    load_catalog(base + "/test_steps_vector_contract.jsonld", catalog, ctx);

    static constexpr const char* kVectorStepIri = "https://cogniflow.org/test#VectorBridgeStep";
    const cf::StepDef* vector_step = find_step_def(catalog, kVectorStepIri);
    ctx.expect(vector_step != nullptr, "VectorBridgeStep metadata should be present in the loaded StepCatalog");
    if (vector_step != nullptr) {
      ctx.expect(vector_step->inputs.size() == 1, "VectorBridgeStep should expose one input");
      ctx.expect(vector_step->outputs.size() == 1, "VectorBridgeStep should expose one output");
      if (vector_step->inputs.size() == 1) {
        ctx.expect(vector_step->inputs[0].storage_type == CF_STORAGE_VECTOR_F64,
                   "VectorBridgeStep input should bridge to VECTOR_F64 runtime storage");
        ctx.expect(vector_step->inputs[0].type_iri == "http://www.w3.org/2001/XMLSchema#double",
                   "VectorBridgeStep input semantic type should remain canonical xsd:double");
      }
      if (vector_step->outputs.size() == 1) {
        ctx.expect(vector_step->outputs[0].storage_type == CF_STORAGE_VECTOR_F64,
                   "VectorBridgeStep output should bridge to VECTOR_F64 runtime storage");
        ctx.expect(vector_step->outputs[0].type_iri == "http://www.w3.org/2001/XMLSchema#double",
                   "VectorBridgeStep output semantic type should remain canonical xsd:double");
      }
    }
  }

  {
    std::string error;
    cf::StepCatalog catalog = cf::load_step_catalog({base + "/test_steps_vector_contract_invalid.jsonld"}, error);
    (void)catalog;
    ctx.expect(
      !error.empty() &&
      error.find("Unsupported vector valueContract") != std::string::npos &&
      error.find("cf:valueContract/cf:elementType is required") != std::string::npos &&
      error.find("https://cogniflow.org/test#MissingElementVectorStep_input_data") != std::string::npos,
      "Vector contract without elementType should be rejected with a clear loader error");
  }

  {
    const std::string non_double_manifest = R"({
  "@context": {
    "cf": "https://cogniflow.org/ns#",
    "cftest": "https://cogniflow.org/test#",
    "xsd": "http://www.w3.org/2001/XMLSchema#"
  },
  "@graph": [
    {
      "@id": "cftest:NonDoubleVectorStep",
      "@type": "cf:ProcessingStep",
      "cf:hasInput": {"@id": "cftest:NonDoubleVectorStep_input_data"}
    },
    {
      "@id": "cftest:NonDoubleVectorStep_input_data",
      "@type": "cf:InputPort",
      "cf:portKey": {"@id": "cf:data"},
      "cf:portType": {"@id": "xsd:integer"},
      "cf:valueContract": {"@id": "cftest:NonDoubleVectorStep_input_data_contract"}
    },
    {
      "@id": "cftest:NonDoubleVectorStep_input_data_contract",
      "@type": "cf:PortValueContract",
      "cf:elementType": {"@id": "xsd:integer"},
      "cf:shapeKind": {"@id": "cf:shape_vector"}
    }
  ]
})";

    std::string error;
    const std::string manifest_path = write_temp_step_file("non_double", non_double_manifest, ctx);
    if (!manifest_path.empty()) {
      cf::StepCatalog catalog = cf::load_step_catalog({manifest_path}, error);
      (void)catalog;
      ctx.expect(
        !error.empty() &&
        error.find("must be xsd:double for cf:shape_vector") != std::string::npos &&
        error.find("http://www.w3.org/2001/XMLSchema#integer") != std::string::npos,
        "Vector contract with non-double elementType should be rejected");
    }
  }

  {
    const std::string mismatch_manifest = R"({
  "@context": {
    "cf": "https://cogniflow.org/ns#",
    "cftest": "https://cogniflow.org/test#",
    "xsd": "http://www.w3.org/2001/XMLSchema#"
  },
  "@graph": [
    {
      "@id": "cftest:MismatchVectorStep",
      "@type": "cf:ProcessingStep",
      "cf:hasOutput": {"@id": "cftest:MismatchVectorStep_output_data"}
    },
    {
      "@id": "cftest:MismatchVectorStep_output_data",
      "@type": "cf:OutputPort",
      "cf:portKey": {"@id": "cf:data"},
      "cf:portType": {"@id": "xsd:double"},
      "cf:valueContract": {"@id": "cftest:MismatchVectorStep_output_data_contract"}
    },
    {
      "@id": "cftest:MismatchVectorStep_output_data_contract",
      "@type": "cf:PortValueContract",
      "cf:elementType": {"@id": "xsd:integer"},
      "cf:shapeKind": {"@id": "cf:shape_vector"}
    }
  ]
})";

    std::string error;
    const std::string manifest_path = write_temp_step_file("mismatch", mismatch_manifest, ctx);
    if (!manifest_path.empty()) {
      cf::StepCatalog catalog = cf::load_step_catalog({manifest_path}, error);
      (void)catalog;
      ctx.expect(
        !error.empty() &&
        error.find("Contract mismatch") != std::string::npos &&
        error.find("cf:portType") != std::string::npos &&
        error.find("cf:valueContract/cf:elementType") != std::string::npos,
        "Vector contract with mismatched elementType should be rejected before bridge application");
    }
  }

  {
    cf::PipelineSpec pipeline;
    cf::StepCatalog catalog;
    cf::PluginRegistry plugins;
    std::string error;

    load_pipeline(base + "/pipeline_json_to_double_mismatch.jsonld", pipeline, ctx);
    catalog = cf::load_step_catalog({basic_io_steps(), basic_signal_steps()}, error);
    if (!error.empty()) {
      ctx.expect(false, error);
    }
    load_plugins({basic_io_plugins(), basic_signal_plugins()}, plugins, ctx);

    cf::ExecutionPlan plan;
    bool ok = cf::build_execution_plan(pipeline, catalog, plugins, plan, error);
    ctx.expect(
      !ok &&
      error.find("http://cogniflow.org/basic-io#JsonScalarToDoubleStep") != std::string::npos &&
      error.find("http://www.w3.org/1999/02/22-rdf-syntax-ns#JSON") != std::string::npos &&
      error.find("http://www.w3.org/2001/XMLSchema#double") != std::string::npos,
      "Direct json-to-double wiring should be rejected with guidance to insert the explicit converter step");
  }

  {
    cf::PipelineSpec pipeline;
    cf::StepCatalog catalog;
    cf::PluginRegistry plugins;
    std::string error;

    load_pipeline(base + "/pipeline_json_scalar_to_double.jsonld", pipeline, ctx);
    catalog = cf::load_step_catalog({basic_io_steps(), basic_signal_steps()}, error);
    if (!error.empty()) {
      ctx.expect(false, error);
    }
    load_plugins({basic_io_plugins(), basic_signal_plugins()}, plugins, ctx);

    cf::ExecutionPlan plan;
    cf::PlanBuildOptions options;
    options.allow_missing_inputs = true;
    bool ok = cf::build_execution_plan(pipeline, catalog, plugins, plan, error, options);
    ctx.expect(ok, "Explicit json-to-double converter should make the wiring graph buildable");
  }

  {
    cf::PipelineSpec pipeline;
    cf::StepCatalog catalog;
    cf::PluginRegistry plugins;
    std::string error;

    load_pipeline(base + "/pipeline_json_scalar_to_double_execute.jsonld", pipeline, ctx);
    catalog = cf::load_step_catalog({base + "/test_steps.jsonld", basic_io_steps()}, error);
    if (!error.empty()) {
      ctx.expect(false, error);
    }
    load_plugins({base, basic_io_plugins()}, plugins, ctx);

    cf::ExecutionPlan plan;
    bool ok = cf::build_execution_plan(pipeline, catalog, plugins, plan, error);
    ctx.expect(ok, "Execution fixture for explicit json-to-double converter should build");
    if (ok) {
      cf::ExecutionResult exec = cf::execute_plan(plan, 1);
      ctx.expect(exec.ok, "Execution fixture for explicit json-to-double converter should run successfully");
      const cf::EdgeSlot* edge = find_edge(plan, "node_convert", "node_sink");
      ctx.expect(edge != nullptr, "Execution fixture should contain converter-to-sink edge");
      if (edge != nullptr) {
        double actual = 0.0;
        ctx.expect(edge->filled, "Converter-to-sink edge should be filled after execution");
        ctx.expect(read_f64(edge->value, actual), "Converter-to-sink edge should carry an f64 payload");
        if (read_f64(edge->value, actual)) {
          ctx.expect(std::fabs(actual - 42.5) < 1e-9, "Converter-to-sink edge should carry the converted value");
        }
      }
    }
  }

  {
    cf::PluginRegistry plugins;
    std::string error;
    load_plugins({basic_io_plugins()}, plugins, ctx);

    const CfStepVTable* step = plugins.resolve_step("http://cogniflow.org/basic-io#JsonScalarToDoubleStep", error);
    ctx.expect(step != nullptr, "JsonScalarToDoubleStep should resolve from the basic_io plugin");

    if (step != nullptr) {
      struct ConversionCase {
        CfStorageType storage_type;
        std::string payload;
        double expected;
        const char* label;
      };
      const std::vector<ConversionCase> cases = {
        {CF_STORAGE_JSON, "42.5", 42.5, "JSON number payload"},
        {CF_STORAGE_JSON, "\"42.5\"", 42.5, "JSON quoted numeric string"},
        {CF_STORAGE_STRING, "42.5", 42.5, "UTF-8 string number"},
        {CF_STORAGE_BYTES, "42.5", 42.5, "bytes payload number"},
      };

      for (const auto& item : cases) {
        CfValueHandle input{};
        input.storage_type = item.storage_type;
        input.data = const_cast<char*>(item.payload.data());
        input.size = item.payload.size();

        CfValueHandle output{};
        AllocatorState allocator_state;
        CfAllocator allocator{test_alloc, test_free, &allocator_state};
        CfCallContext call_ctx{&allocator, nullptr};

        CfStatusCode code = step->call(&call_ctx, nullptr, 0, &input, 1, &output, 1);
        double actual = 0.0;
        ctx.expect(code == CF_STATUS_OK, std::string(item.label) + " should convert successfully");
        ctx.expect(read_f64(output, actual), std::string(item.label) + " should produce an f64 output");
        if (read_f64(output, actual)) {
          ctx.expect(std::fabs(actual - item.expected) < 1e-9, std::string(item.label) + " should match expected value");
        }

        free_allocations(allocator_state);
      }
    }
  }

  {
    cf::StepCatalog catalog;
    cf::PluginRegistry plugins;
    std::string error;
    catalog = cf::load_step_catalog({basic_signal_steps()}, error);
    if (!error.empty()) {
      ctx.expect(false, error);
    }
    load_plugins({basic_signal_plugins()}, plugins, ctx);

    static constexpr const char* kFifoStepIri = "http://cogniflow.org/basic-signal#FifoWindowBufferStep";
    const CfStepVTable* step = plugins.resolve_step("http://cogniflow.org/basic-signal#FifoWindowBufferStep", error);
    ctx.expect(step != nullptr, "FifoWindowBufferStep should resolve from the basic_signal plugin");
    const cf::StepDef* fifo_step_def = find_step_def(catalog, kFifoStepIri);
    ctx.expect(fifo_step_def != nullptr, "FifoWindowBufferStep metadata should be present in the loaded StepCatalog");

    size_t window_size_param_index = 0;
    bool have_window_size_param = find_parameter_index(catalog, kFifoStepIri, "windowSize", window_size_param_index);
    ctx.expect(have_window_size_param, "FifoWindowBufferStep windowSize parameter should be present in the loaded StepCatalog");

    if (step != nullptr && fifo_step_def != nullptr && have_window_size_param) {
      const size_t fifo_param_count = fifo_step_def->parameters.size();

      {
        double input_value = 42.5;
        int64_t window_size = 1;
        std::vector<CfValueHandle> params(fifo_param_count);
        params[window_size_param_index].storage_type = CF_STORAGE_I64;
        params[window_size_param_index].data = &window_size;
        params[window_size_param_index].size = sizeof(window_size);

        CfValueHandle input{};
        input.storage_type = CF_STORAGE_F64;
        input.data = &input_value;
        input.size = sizeof(input_value);

        CfValueHandle output{};
        AllocatorState allocator_state;
        CfAllocator allocator{test_alloc, test_free, &allocator_state};
        CfCallContext call_ctx{&allocator, nullptr};

        CfStatusCode code = step->call(&call_ctx, params.data(), params.size(), &input, 1, &output, 1);
        std::vector<double> values;
        ctx.expect(code == CF_STATUS_OK, "FifoWindowBufferStep should accept direct f64 input");
        ctx.expect(read_vector_f64_values(output, values), "FifoWindowBufferStep should emit a vector for direct f64 input");
        if (read_vector_f64_values(output, values)) {
          ctx.expect(values.size() == 1, "FifoWindowBufferStep f64 output should contain one sample");
          if (values.size() == 1) {
            ctx.expect(std::fabs(values[0] - 42.5) < 1e-9, "FifoWindowBufferStep f64 output should preserve the sample value");
          }
        }

        free_allocations(allocator_state);
      }

      {
        double vector_values[] = {1.0, 2.0, 3.0};
        CfVectorF64 vector{vector_values, 3};
        int64_t window_size = 3;
        std::vector<CfValueHandle> params(fifo_param_count);
        params[window_size_param_index].storage_type = CF_STORAGE_I64;
        params[window_size_param_index].data = &window_size;
        params[window_size_param_index].size = sizeof(window_size);

        CfValueHandle input{};
        input.storage_type = CF_STORAGE_VECTOR_F64;
        input.data = &vector;
        input.size = sizeof(vector);

        CfValueHandle output{};
        AllocatorState allocator_state;
        CfAllocator allocator{test_alloc, test_free, &allocator_state};
        CfCallContext call_ctx{&allocator, nullptr};

        CfStatusCode code = step->call(&call_ctx, params.data(), params.size(), &input, 1, &output, 1);
        std::vector<double> values;
        ctx.expect(code == CF_STATUS_OK, "FifoWindowBufferStep should accept direct vector input");
        ctx.expect(read_vector_f64_values(output, values), "FifoWindowBufferStep should emit a vector for vector input");
        if (read_vector_f64_values(output, values)) {
          ctx.expect(values.size() == 3, "FifoWindowBufferStep vector output should preserve the full window");
          if (values.size() == 3) {
            ctx.expect(std::fabs(values[0] - 1.0) < 1e-9, "FifoWindowBufferStep vector output should preserve the first sample");
            ctx.expect(std::fabs(values[1] - 2.0) < 1e-9, "FifoWindowBufferStep vector output should preserve the middle sample");
            ctx.expect(std::fabs(values[2] - 3.0) < 1e-9, "FifoWindowBufferStep vector output should preserve the last sample");
          }
        }

        free_allocations(allocator_state);
      }

      struct InvalidCase {
        CfStorageType storage_type;
        std::string payload;
        const char* label;
      };
      const std::vector<InvalidCase> invalid_cases = {
        {CF_STORAGE_JSON, "42.5", "FifoWindowBufferStep should reject JSON number payloads"},
        {CF_STORAGE_JSON, "\"42.5\"", "FifoWindowBufferStep should reject JSON quoted numeric string payloads"},
        {CF_STORAGE_STRING, "42.5", "FifoWindowBufferStep should reject UTF-8 numeric string payloads"},
        {CF_STORAGE_BYTES, "42.5", "FifoWindowBufferStep should reject bytes numeric payloads"},
      };

      for (const auto& item : invalid_cases) {
        int64_t window_size = 1;
        std::vector<CfValueHandle> params(fifo_param_count);
        params[window_size_param_index].storage_type = CF_STORAGE_I64;
        params[window_size_param_index].data = &window_size;
        params[window_size_param_index].size = sizeof(window_size);

        CfValueHandle input{};
        input.storage_type = item.storage_type;
        input.data = const_cast<char*>(item.payload.data());
        input.size = item.payload.size();

        CfValueHandle output{};
        AllocatorState allocator_state;
        CfAllocator allocator{test_alloc, test_free, &allocator_state};
        CfCallContext call_ctx{&allocator, nullptr};

        CfStatusCode code = step->call(&call_ctx, params.data(), params.size(), &input, 1, &output, 1);
        ctx.expect(code == CF_STATUS_INVALID, item.label);
        free_allocations(allocator_state);
      }
    }
  }

  {
    cf::PipelineSpec pipeline;
    cf::StepCatalog catalog;
    cf::PluginRegistry plugins;
    std::string error;

    load_pipeline(base + "/pipeline_type_mismatch.jsonld", pipeline, ctx);
    load_catalog(base + "/test_steps.jsonld", catalog, ctx);
    load_plugins({base}, plugins, ctx);

    cf::ExecutionPlan plan;
    bool ok = cf::build_execution_plan(pipeline, catalog, plugins, plan, error);
    ctx.expect(
      !ok &&
      error.find("Incompatible connection from node 'node_source'") != std::string::npos &&
      error.find("node_sink") != std::string::npos &&
      error.find("https://cogniflow.org/ns#data") != std::string::npos &&
      error.find("http://cogniflow.org/basic-signal#observed_thread_budget") != std::string::npos,
      "Plan build should reject incompatible port datatype connections");
  }

  {
    cf::PipelineSpec pipeline;
    cf::StepCatalog catalog;
    cf::PluginRegistry plugins;
    std::string error;

    load_pipeline(base + "/pipeline_mismatch.jsonld", pipeline, ctx);
    load_catalog(base + "/test_steps_mismatch.jsonld", catalog, ctx);
    load_plugins({base}, plugins, ctx);

    cf::ExecutionPlan plan;
    bool ok = cf::build_execution_plan(pipeline, catalog, plugins, plan, error);
    ctx.expect(!ok && error.find("Signature mismatch") != std::string::npos,
      "Signature mismatch should be detected");
  }

  {
    cf::PipelineSpec pipeline;
    cf::StepCatalog catalog;
    cf::PluginRegistry plugins;
    std::string error;

    load_pipeline(base + "/pipeline_edges.jsonld", pipeline, ctx);
    load_catalog(base + "/test_steps.jsonld", catalog, ctx);
    load_plugins({base}, plugins, ctx);

    cf::ExecutionPlan plan;
    bool ok = cf::build_execution_plan(pipeline, catalog, plugins, plan, error);
    ctx.expect(ok, "Plan build should succeed for edge test");
    if (ok) {
      ctx.expect(plan.edges.size() == 1, "Edge slot count should be 1");
      if (!plan.edges.empty()) {
        ctx.expect(plan.edges[0].id == 0, "Edge slot id should be 0");
      }
    }
  }

  {
    cf::PipelineSpec pipeline;
    cf::StepCatalog catalog;
    cf::PluginRegistry plugins;
    std::string error;

    load_pipeline(base + "/pipeline_basic_io_inline.jsonld", pipeline, ctx);
    catalog = cf::load_step_catalog({basic_io_steps(), basic_dev_steps()}, error);
    if (!error.empty()) {
      ctx.expect(false, error);
    }
    load_plugins({basic_io_plugins(), basic_dev_plugins()}, plugins, ctx);

    cf::ExecutionPlan plan;
    bool ok = cf::build_execution_plan(pipeline, catalog, plugins, plan, error);
    ctx.expect(ok, "Plan build should succeed for basic IO pipeline");
    if (ok) {
      cf::ExecutionResult exec = cf::execute_plan(plan, 1);
      ctx.expect(exec.ok, "Basic IO pipeline execution should succeed");
      if (exec.ok) {
        int handle_idx = -1;
        int stats_idx = -1;
        for (size_t i = 0; i < plan.nodes.size(); ++i) {
          if (plan.nodes[i].node_id == "node_handle") handle_idx = static_cast<int>(i);
          if (plan.nodes[i].node_id == "node_stats") stats_idx = static_cast<int>(i);
        }
        ctx.expect(handle_idx >= 0 && stats_idx >= 0, "Basic IO nodes should be present");
        const cf::EdgeSlot* status_edge = nullptr;
        for (const auto& edge : plan.edges) {
          if (edge.from_node == handle_idx && edge.to_node == stats_idx) {
            status_edge = &edge;
            break;
          }
        }
        ctx.expect(status_edge != nullptr, "Status edge should exist");
        if (status_edge) {
          ctx.expect(status_edge->filled, "Status edge should be filled");
          if (status_edge->filled && status_edge->value.data && status_edge->value.size > 0) {
            std::string status(reinterpret_cast<const char*>(status_edge->value.data), status_edge->value.size);
            ctx.expect(status == "handled via memory://test", "Status output should match handle uri");
          } else {
            ctx.expect(false, "Status edge should contain output payload");
          }
        }
      }
    }
  }

  {
    cf::PipelineSpec pipeline;
    cf::StepCatalog catalog;
    cf::PluginRegistry plugins;
    std::string error;

    load_pipeline(base + "/pipeline_levels.jsonld", pipeline, ctx);
    load_catalog(base + "/test_steps.jsonld", catalog, ctx);
    load_plugins({base}, plugins, ctx);

    cf::ExecutionPlan plan;
    bool ok = cf::build_execution_plan(pipeline, catalog, plugins, plan, error);
    ctx.expect(ok, "Plan build should succeed for levels test");
    if (ok) {
      ctx.expect(plan.levels.size() == 2, "Levels should have 2 stages");
      if (plan.levels.size() == 2) {
        bool root_in_level0 = false;
        for (int node_idx : plan.levels[0]) {
          if (node_idx == 0) root_in_level0 = true;
        }
        ctx.expect(root_in_level0, "Root node should be in level 0");
        ctx.expect(plan.levels[1].size() == 2, "Level 1 should contain two nodes");
      }
    }
  }

  {
    cf::PipelineSpec pipeline;
    cf::StepCatalog catalog;
    cf::PluginRegistry plugins;
    std::string error;

    load_pipeline(base + "/pipeline_parallel.jsonld", pipeline, ctx);
    load_catalog(base + "/test_steps.jsonld", catalog, ctx);
    load_plugins({base}, plugins, ctx);

    cf::ExecutionPlan plan;
    bool ok = cf::build_execution_plan(pipeline, catalog, plugins, plan, error);
    ctx.expect(ok, "Plan build should succeed for parallel test");
    if (ok) {
      auto start = std::chrono::steady_clock::now();
      cf::ExecutionResult exec = cf::execute_plan(plan, 2);
      auto end = std::chrono::steady_clock::now();
      auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();
      ctx.expect(exec.ok, "Parallel execution should succeed");
      ctx.expect(elapsed < 350, "Parallel sleep should finish under 350ms");
    }
  }

  {
    cf::PipelineSpec pipeline;
    cf::StepCatalog catalog;
    cf::PluginRegistry plugins;
    std::string error;

    load_pipeline(base + "/pipeline_budget_aware_hints.jsonld", pipeline, ctx);

    catalog = cf::load_step_catalog({base + "/test_steps.jsonld", basic_dev_steps()}, error);
    if (!error.empty()) {
      ctx.expect(false, error);
    }

    load_plugins({base, basic_dev_plugins()}, plugins, ctx);

    cf::ExecutionPlan plan;
    bool ok = cf::build_execution_plan(pipeline, catalog, plugins, plan, error);
    ctx.expect(ok, "Plan build should succeed for budget-aware hints test");
    if (ok) {
      cf::ConcurrencyOptions options;
      options.engine_threads = 1;
      options.max_total_threads = 8;
      options.default_step_threads = 4;

      cf::ExecutionResult exec = cf::execute_plan(plan, options);
      ctx.expect(exec.ok, "Budget-aware hints execution should succeed");
      ctx.expect(plan.edges.size() == 1, "Budget-aware hints plan should contain one edge");
      if (plan.edges.size() == 1) {
        const cf::EdgeSlot& edge = plan.edges[0];
        ctx.expect(edge.filled, "Budget-aware hints output edge should be filled");
        ctx.expect(edge.value.storage_type == CF_STORAGE_I64, "observed_thread_budget output should be I64");
        int64_t observed = 0;
        if (edge.value.data && edge.value.size >= sizeof(int64_t)) {
          observed = *reinterpret_cast<const int64_t*>(edge.value.data);
        }
        ctx.expect(observed == 4, "observed_thread_budget should match the assigned thread budget (4)");
      }
    }
  }

  {
    cf::PipelineSpec pipeline;
    cf::StepCatalog catalog;
    cf::PluginRegistry plugins;
    std::string error;

    load_pipeline(base + "/pipeline_budget_weighted.jsonld", pipeline, ctx);
    load_catalog(base + "/test_steps.jsonld", catalog, ctx);
    load_plugins({base}, plugins, ctx);

    cf::ExecutionPlan plan;
    bool ok = cf::build_execution_plan(pipeline, catalog, plugins, plan, error);
    ctx.expect(ok, "Plan build should succeed for budget-weighted test");
    if (ok) {
      cf::ConcurrencyOptions options;
      options.engine_threads = 2;
      options.max_total_threads = 4;

      auto start = std::chrono::steady_clock::now();
      cf::ExecutionResult exec = cf::execute_plan(plan, options);
      auto end = std::chrono::steady_clock::now();
      auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();
      ctx.expect(exec.ok, "Budget-weighted execution should succeed");
      ctx.expect(elapsed > 350, "Heavy steps should serialize under tight max_total_threads budget");
      ctx.expect(plan.nodes.size() == 2, "Budget-weighted plan should contain two nodes");
      if (plan.nodes.size() == 2) {
        ctx.expect(plan.nodes[0].thread_budget == 4, "BudgetCheckStep thread_budget should be 4");
        ctx.expect(plan.nodes[1].thread_budget == 4, "BudgetCheckStep thread_budget should be 4");
      }

      options.max_total_threads = 8;
      start = std::chrono::steady_clock::now();
      exec = cf::execute_plan(plan, options);
      end = std::chrono::steady_clock::now();
      elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();
      ctx.expect(exec.ok, "Budget-weighted execution should succeed under larger budget");
      ctx.expect(elapsed < 350, "Heavy steps should run concurrently under larger max_total_threads budget");
    }
  }

  {
    cf::PipelineSpec pipeline;
    cf::StepCatalog catalog;
    cf::PluginRegistry plugins;
    std::string error;

    load_pipeline(base + "/pipeline_budget_deterministic.jsonld", pipeline, ctx);
    load_catalog(base + "/test_steps.jsonld", catalog, ctx);
    load_plugins({base}, plugins, ctx);

    cf::ExecutionPlan plan;
    bool ok = cf::build_execution_plan(pipeline, catalog, plugins, plan, error);
    ctx.expect(ok, "Plan build should succeed for deterministic budget test");
    if (ok) {
      cf::ConcurrencyOptions options;
      options.engine_threads = 1;
      options.max_total_threads = 4;
      options.deterministic = true;

      cf::ExecutionResult exec = cf::execute_plan(plan, options);
      ctx.expect(exec.ok, "Deterministic budget execution should succeed");
      ctx.expect(plan.nodes.size() == 1, "Deterministic budget plan should contain one node");
      if (plan.nodes.size() == 1) {
        ctx.expect(plan.nodes[0].thread_budget == 1,
          "Deterministic mode should clamp thread_budget to 1 unless step declares deterministic=true");
      }
    }
  }

  {
    cf::PipelineSpec pipeline;
    cf::StepCatalog catalog;
    cf::PluginRegistry plugins;
    std::string error;

    load_pipeline(base + "/pipeline_invalid.jsonld", pipeline, ctx);
    load_catalog(base + "/test_steps.jsonld", catalog, ctx);
    load_plugins({base}, plugins, ctx);

    cf::ExecutionPlan plan;
    bool ok = cf::build_execution_plan(pipeline, catalog, plugins, plan, error);
    ctx.expect(ok, "Plan build should succeed for validation test");
    if (ok) {
      cf::ExecutionResult exec = cf::execute_plan(plan, 1);
      ctx.expect(!exec.ok && exec.error.find("Validation error") != std::string::npos,
        "Validation error should prevent execution");
    }
  }

  if (ctx.failures > 0) {
    std::cerr << ctx.failures << " tests failed\n";
    return 1;
  }
  std::cout << "All tests passed\n";
  return 0;
}

