#include "scheduler/scheduler.h"

#include <atomic>
#include <chrono>
#include <condition_variable>
#include <cstring>
#include <deque>
#include <iomanip>
#include <mutex>
#include <sstream>
#include <thread>
#include <unordered_map>
#include <utility>
#include <vector>

#include "scheduler/thread_pool.h"

namespace cf {

namespace {

struct ExecutionAllocator {
  std::mutex mutex;
  std::vector<std::unique_ptr<uint8_t[]>>* buffers = nullptr;

  static void* alloc_fn(size_t size, void* user_data) {
    auto* self = reinterpret_cast<ExecutionAllocator*>(user_data);
    if (!self || !self->buffers || size == 0) return nullptr;
    std::unique_ptr<uint8_t[]> buffer(new (std::nothrow) uint8_t[size]);
    if (!buffer) return nullptr;
    void* ptr = buffer.get();
    {
      std::lock_guard<std::mutex> lock(self->mutex);
      self->buffers->push_back(std::move(buffer));
    }
    return ptr;
  }

  static void free_fn(void*, void*) {}
};

void add_issue(CfValidationReport* report, CfValidationSeverity severity, const std::string& message) {
  if (!report || !report->allocator) return;
  if (report->issue_count >= report->issue_capacity) {
    size_t new_cap = report->issue_capacity == 0 ? 4 : report->issue_capacity * 2;
    size_t bytes = sizeof(CfValidationIssue) * new_cap;
    auto* mem = reinterpret_cast<CfValidationIssue*>(report->allocator->alloc(bytes, report->allocator->user_data));
    if (!mem) return;
    for (size_t i = 0; i < report->issue_count; ++i) {
      mem[i] = report->issues[i];
    }
    report->issues = mem;
    report->issue_capacity = new_cap;
  }
  auto* issue = &report->issues[report->issue_count++];
  issue->severity = severity;
  char* msg = reinterpret_cast<char*>(report->allocator->alloc(message.size() + 1, report->allocator->user_data));
  if (msg) {
    std::memcpy(msg, message.data(), message.size());
    msg[message.size()] = '\0';
    issue->message = msg;
  } else {
    issue->message = nullptr;
  }
  issue->parameter_key = nullptr;
  issue->port_key = nullptr;
}

std::string json_escape(const std::string& value) {
  std::ostringstream out;
  for (char c : value) {
    switch (c) {
      case '"': out << "\\\""; break;
      case '\\': out << "\\\\"; break;
      case '\b': out << "\\b"; break;
      case '\f': out << "\\f"; break;
      case '\n': out << "\\n"; break;
      case '\r': out << "\\r"; break;
      case '\t': out << "\\t"; break;
      default:
        if (static_cast<unsigned char>(c) < 0x20) {
          out << "\\u" << std::hex << std::uppercase << static_cast<int>(c);
        } else {
          out << c;
        }
        break;
    }
  }
  return out.str();
}

std::string handle_to_json(const CfValueHandle& handle) {
  if (!handle.data || handle.size == 0) {
    return "null";
  }
  switch (handle.storage_type) {
    case CF_STORAGE_BOOL: {
      const uint8_t v = *reinterpret_cast<const uint8_t*>(handle.data);
      return v ? "true" : "false";
    }
    case CF_STORAGE_I64: {
      if (handle.size < sizeof(int64_t)) return "null";
      const int64_t v = *reinterpret_cast<const int64_t*>(handle.data);
      return std::to_string(v);
    }
    case CF_STORAGE_F64: {
      if (handle.size < sizeof(double)) return "null";
      const double v = *reinterpret_cast<const double*>(handle.data);
      std::ostringstream out;
      out << std::setprecision(15) << v;
      return out.str();
    }
    case CF_STORAGE_STRING: {
      std::string value(reinterpret_cast<const char*>(handle.data), handle.size);
      return std::string("\"") + json_escape(value) + "\"";
    }
    case CF_STORAGE_JSON: {
      return std::string(reinterpret_cast<const char*>(handle.data), handle.size);
    }
    case CF_STORAGE_BYTES:
    case CF_STORAGE_OPAQUE: {
      std::string value(reinterpret_cast<const char*>(handle.data), handle.size);
      return std::string("\"") + json_escape(value) + "\"";
    }
    default:
      break;
  }
  return "null";
}

std::string build_multi_input_json(const std::vector<std::pair<std::string, CfValueHandle>>& items) {
  std::ostringstream out;
  out << '{';
  for (size_t i = 0; i < items.size(); ++i) {
    if (i > 0) out << ',';
    out << '"' << json_escape(items[i].first) << "\":" << handle_to_json(items[i].second);
  }
  out << '}';
  return out.str();
}

}  // namespace

ExecutionResult execute_plan(ExecutionPlan& plan, const ConcurrencyOptions& options) {
  ExecutionResult result;

  for (auto& edge : plan.edges) {
    edge.filled = false;
    edge.value = {};
  }
  plan.execution_owned_buffers.clear();

  const ResolvedConcurrencySettings resolved =
    ConcurrencyController::resolve(options, std::thread::hardware_concurrency());
  ConcurrencyController concurrency(resolved);
  auto process_scope = concurrency.enter_process_scope();
  concurrency.assign_plan_budgets(plan);

  ExecutionAllocator exec_alloc;
  exec_alloc.buffers = &plan.execution_owned_buffers;
  CfAllocator allocator{&ExecutionAllocator::alloc_fn, &ExecutionAllocator::free_fn, &exec_alloc};
  CfCallContext validate_ctx{&allocator, nullptr};

  for (auto& node : plan.nodes) {
    if (node.vtable && node.vtable->validate) {
      CfValidationReport report{};
      report.allocator = &allocator;
      std::vector<CfValueHandle> empty_inputs(node.step_def.inputs.size());
      CfStatusCode code = node.vtable->validate(
        &validate_ctx,
        node.params.data(),
        node.params.size(),
        empty_inputs.data(),
        empty_inputs.size(),
        &report
      );
      if (code != CF_STATUS_OK) {
        result.ok = false;
        result.error = "Validation call failed for node " + node.node_id;
        return result;
      }
      for (size_t i = 0; i < report.issue_count; ++i) {
        if (report.issues[i].severity == CF_VALIDATION_ERROR) {
          result.ok = false;
          result.error = "Validation error on node " + node.node_id + ": " +
            (report.issues[i].message ? report.issues[i].message : "unknown");
          return result;
        }
      }
    }
  }

  std::vector<std::atomic<int>> remaining_inputs(plan.nodes.size());
  std::vector<std::vector<CfValueHandle>> input_values(plan.nodes.size());
  std::vector<std::vector<std::vector<std::pair<std::string, CfValueHandle>>>> multi_inputs(plan.nodes.size());
  std::vector<std::mutex> node_input_mutex(plan.nodes.size());
  for (size_t i = 0; i < plan.nodes.size(); ++i) {
    const auto& node = plan.nodes[i];
    input_values[i].resize(node.step_def.inputs.size());
    multi_inputs[i].resize(node.step_def.inputs.size());
    std::vector<int> edge_counts(node.step_def.inputs.size(), 0);
    for (int edge_id : node.input_edge_ids) {
      if (edge_id >= 0 && static_cast<size_t>(edge_id) < plan.edges.size()) {
        const auto& edge = plan.edges[edge_id];
        if (edge.to_input >= 0 && static_cast<size_t>(edge.to_input) < edge_counts.size()) {
          edge_counts[edge.to_input] += 1;
        }
      }
    }
    int pending = 0;
    for (int count : edge_counts) {
      if (count > 0) pending += count;
    }
    remaining_inputs[i].store(pending);
  }

  std::mutex sched_mutex;
  std::deque<int> ready;
  size_t available_budget = resolved.max_total_threads;
  for (size_t i = 0; i < plan.nodes.size(); ++i) {
    if (remaining_inputs[i].load() == 0) {
      ready.push_back(static_cast<int>(i));
    }
  }

  std::atomic<size_t> completed{0};
  std::atomic<bool> failed{false};
  std::atomic<bool> halted{false};
  std::mutex error_mutex;
  std::string error_message;

  std::condition_variable cv;
  const size_t pool_threads = std::min(resolved.engine_threads, resolved.max_total_threads);
  ThreadPool pool(pool_threads);

  struct BudgetRelease {
    std::mutex& mutex;
    std::condition_variable& cv;
    size_t& available_budget;
    size_t cost;

    ~BudgetRelease() {
      std::lock_guard<std::mutex> lock(mutex);
      available_budget += cost;
      cv.notify_one();
    }
  };

  auto schedule_node = [&](int node_idx, size_t cost) {
    pool.enqueue([&, node_idx, cost]() {
      BudgetRelease release{sched_mutex, cv, available_budget, cost};
      if (failed.load()) return;
      PlanNode& node = plan.nodes[node_idx];

      std::vector<CfValueHandle> outputs(node.step_def.outputs.size());
      std::vector<CfValueHandle> inputs = input_values[node_idx];
      for (size_t i = 0; i < node.step_def.inputs.size(); ++i) {
        if (node.step_def.inputs[i].is_multi) {
          std::string json_payload = build_multi_input_json(multi_inputs[node_idx][i]);
          CfCallContext call_ctx{&allocator, nullptr};
          CfStatusCode code = cf_out_bytes(&call_ctx, &inputs[i], CF_STORAGE_JSON,
                                           json_payload.data(), json_payload.size());
          if (code != CF_STATUS_OK) {
            failed.store(true);
            std::lock_guard<std::mutex> lock(error_mutex);
            error_message = "Failed to allocate multi-input payload for node " + node.node_id;
            return;
          }
        }
      }

      CfCallContext call_ctx{&allocator, nullptr};
      auto step_scope = concurrency.enter_step_scope(node, call_ctx);
      CfStatusCode code = node.vtable->call(
        &call_ctx,
        node.params.data(),
        node.params.size(),
        inputs.data(),
        inputs.size(),
        outputs.data(),
        outputs.size()
      );
      if (code != CF_STATUS_OK) {
        failed.store(true);
        if (code == CF_STATUS_INVALID &&
            node.step_iri == "http://cogniflow.org/basic-signal#FifoWindowBufferStep") {
          halted.store(true);
          return;
        }
        std::lock_guard<std::mutex> lock(error_mutex);
        error_message = "Execution failed for node " + node.node_id;
        return;
      }

      for (int edge_id : node.output_edge_ids) {
        if (edge_id < 0 || static_cast<size_t>(edge_id) >= plan.edges.size()) continue;
        EdgeSlot& edge = plan.edges[edge_id];
        if (edge.from_output < 0 || static_cast<size_t>(edge.from_output) >= outputs.size()) continue;
        edge.value = outputs[edge.from_output];
        edge.filled = true;

        int downstream = edge.to_node;
        int input_idx = edge.to_input;
        if (downstream >= 0 && static_cast<size_t>(downstream) < input_values.size()) {
          if (input_idx >= 0 && static_cast<size_t>(input_idx) < input_values[downstream].size()) {
            {
              std::lock_guard<std::mutex> lock(node_input_mutex[downstream]);
              if (plan.nodes[downstream].step_def.inputs[input_idx].is_multi) {
                std::string label = edge.label.empty() ? edge.from_port_key : edge.label;
                multi_inputs[downstream][input_idx].push_back({label, edge.value});
              } else {
                input_values[downstream][input_idx] = edge.value;
              }
            }

            int remaining = --remaining_inputs[downstream];
            if (remaining == 0) {
              std::lock_guard<std::mutex> lock(sched_mutex);
              ready.push_back(downstream);
              cv.notify_one();
            }
          }
        }
      }

      completed.fetch_add(1);
    });
  };

  while (completed.load() < plan.nodes.size()) {
    if (failed.load()) break;

    int node_idx = -1;
    size_t cost = 0;

    {
      std::unique_lock<std::mutex> lock(sched_mutex);

      auto can_schedule = [&]() -> bool {
        for (int idx : ready) {
          if (idx < 0 || static_cast<size_t>(idx) >= plan.nodes.size()) continue;
          if (plan.nodes[idx].scheduler_cost <= available_budget) {
            return true;
          }
        }
        return false;
      };

      if (ready.empty() || !can_schedule()) {
        cv.wait_for(lock, std::chrono::milliseconds(10));
      }

      if (failed.load()) break;

      for (auto it = ready.begin(); it != ready.end(); ++it) {
        const int idx = *it;
        if (idx < 0 || static_cast<size_t>(idx) >= plan.nodes.size()) continue;
        const size_t node_cost = plan.nodes[idx].scheduler_cost;
        if (node_cost <= available_budget) {
          node_idx = idx;
          cost = node_cost;
          ready.erase(it);
          available_budget -= node_cost;
          break;
        }
      }
    }

    if (node_idx >= 0) {
      schedule_node(node_idx, cost);
      continue;
    }

    std::this_thread::sleep_for(std::chrono::milliseconds(1));
  }

  pool.wait_for_idle();

  if (failed.load()) {
    if (halted.load()) {
      result.ok = true;
      return result;
    }
    result.ok = false;
    result.error = error_message.empty() ? "Execution failed" : error_message;
    return result;
  }

  if (completed.load() != plan.nodes.size()) {
    result.ok = false;
    result.error = "Execution stalled";
    return result;
  }

  result.ok = true;
  return result;
}

ExecutionResult execute_plan(ExecutionPlan& plan, size_t threads) {
  ConcurrencyOptions options;
  options.engine_threads = threads;
  return execute_plan(plan, options);
}

}  // namespace cf
