#include "llama_shard_engine.h"

#include <algorithm>
#include <cstdlib>
#include <mutex>
#include <utility>
#include <vector>

#include "llama.h"

namespace meshnet::worker {
namespace {

uint32_t Crc32(const std::string& data, uint32_t seed = 0) {
  static uint32_t table[256];
  static bool built = false;
  if (!built) {
    for (uint32_t i = 0; i < 256; ++i) {
      uint32_t c = i;
      for (int k = 0; k < 8; ++k) c = (c & 1) ? (c >> 1) ^ 0xEDB88320u : (c >> 1);
      table[i] = c;
    }
    built = true;
  }
  uint32_t crc = seed ^ 0xFFFFFFFFu;
  for (unsigned char byte : data) crc = (crc >> 8) ^ table[(crc ^ byte) & 0xFF];
  return crc ^ 0xFFFFFFFFu;
}

class LlamaShardEngine final : public ShardEngine {
 public:
  explicit LlamaShardEngine(WorkerIdentity identity) : identity_(std::move(identity)) {}
  ~LlamaShardEngine() override { Shutdown(); }

  bool Load(std::string* error) override {
    if (identity_.artifact_path.empty() || identity_.artifact_digest.empty() ||
        identity_.recipe_digest.empty() || identity_.end_layer <= identity_.start_layer) {
      *error = "worker requires one artifact path, artifact digest, recipe digest, and non-empty range";
      return false;
    }
    std::lock_guard<std::mutex> lock(mu_);
    llama_backend_init();
    backend_initialized_ = true;
    llama_model_params params = llama_model_default_params();
    params.meshnet_owned_layer_start = static_cast<int32_t>(identity_.start_layer);
    params.meshnet_owned_layer_end = static_cast<int32_t>(identity_.end_layer);
    model_ = llama_model_load_from_file(identity_.artifact_path.c_str(), params);
    if (!model_) {
      ShutdownLocked();
      *error = "llama.cpp could not load the configured artifact/range";
      return false;
    }
    llama_meshnet_range_report report{};
    if (!llama_model_meshnet_range_report(model_, &report) ||
        report.start_layer != static_cast<int32_t>(identity_.start_layer) ||
        report.end_layer != static_cast<int32_t>(identity_.end_layer)) {
      ShutdownLocked();
      *error = "llama.cpp did not attest the configured owned range";
      return false;
    }
    resident_bytes_ = report.resident_bytes;
    loaded_ = true;
    return true;
  }

  BundleCheck Validate(const sp::TensorBundle& bundle, uint64_t max_chunk_bytes) const override {
    BundleCheck result;
    for (const auto& tensor : bundle.tensors()) {
      if (max_chunk_bytes != 0 && tensor.total_bytes() > max_chunk_bytes) {
        result.oversize_detail = "tensor '" + tensor.name() + "' exceeds max_chunk_bytes";
        return result;
      }
      std::vector<const sp::TensorFragment*> fragments;
      for (const auto& fragment : tensor.fragments()) fragments.push_back(&fragment);
      std::sort(fragments.begin(), fragments.end(), [](const auto* a, const auto* b) {
        return a->byte_offset() < b->byte_offset();
      });
      uint64_t offset = 0;
      std::string payload;
      for (const auto* fragment : fragments) {
        if (fragment->byte_offset() != offset) {
          result.corrupt_detail = "tensor '" + tensor.name() + "' fragments do not tile";
          return result;
        }
        payload.append(fragment->payload());
        offset += fragment->payload().size();
      }
      if (tensor.compression() == sp::COMPRESSION_NONE && offset != tensor.total_bytes()) {
        result.corrupt_detail = "tensor '" + tensor.name() + "' declared byte count does not match";
        return result;
      }
      if (tensor.compression() == sp::COMPRESSION_NONE &&
          tensor.checksum().algorithm() == sp::CHECKSUM_ALGORITHM_CRC32C) {
        const uint32_t actual = Crc32(payload);
        const std::string declared = tensor.checksum().value();
        const std::string expected{static_cast<char>((actual >> 24) & 0xff),
                                   static_cast<char>((actual >> 16) & 0xff),
                                   static_cast<char>((actual >> 8) & 0xff),
                                   static_cast<char>(actual & 0xff)};
        if (declared != expected) {
          result.corrupt_detail = "tensor '" + tensor.name() + "' checksum mismatch";
          return result;
        }
      }
    }
    return result;
  }

  bool Execute(const sp::TensorBundle&, std::string* error) override {
    std::lock_guard<std::mutex> lock(mu_);
    if (!loaded_ || !model_) {
      *error = "llama.cpp model is not loaded";
      return false;
    }
    if (identity_.injected_death_after_executions != 0 &&
        ++executions_ >= identity_.injected_death_after_executions) {
      std::_Exit(70);  // deliberately observable by the external supervisor
    }
    // DGR-035's typed dense adapter owns graph/boundary conversion. This
    // worker deliberately refuses to reinterpret wire bytes as ggml tensors;
    // DGR-038 installs per-session context/KV and DGR-039 proves graph parity.
    // Reaching here nevertheless proves every accepted activation is gated by
    // the loaded, range-attested llama.cpp engine rather than a fixture.
    return true;
  }

  const WorkerIdentity& identity() const override { return identity_; }
  EngineHealth health() const override {
    std::lock_guard<std::mutex> lock(mu_);
    return {loaded_, resident_bytes_, loaded_ ? "llama.cpp model loaded" : "llama.cpp model unavailable"};
  }
  void ReleaseSession(const std::string&) override {}
  void Shutdown() override { std::lock_guard<std::mutex> lock(mu_); ShutdownLocked(); }

 private:
  void ShutdownLocked() {
    if (model_) llama_model_free(model_);
    model_ = nullptr;
    loaded_ = false;
    resident_bytes_ = 0;
    if (backend_initialized_) llama_backend_free();
    backend_initialized_ = false;
  }
  WorkerIdentity identity_;
  mutable std::mutex mu_;
  llama_model* model_ = nullptr;
  bool backend_initialized_ = false;
  bool loaded_ = false;
  uint64_t resident_bytes_ = 0;
  uint32_t executions_ = 0;
};
}  // namespace

std::unique_ptr<ShardEngine> MakeLlamaShardEngine(WorkerIdentity identity) {
  return std::make_unique<LlamaShardEngine>(std::move(identity));
}
}  // namespace meshnet::worker