#include <array>
#include <atomic>
#include <chrono>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <iomanip>
#include <iostream>
#include <mutex>
#include <sstream>
#include <string>
#include <thread>
#include <vector>
#include <algorithm>

#include "wbag_test.h"
#include "crypto/hash.h"
#include "string_tools.h"
#include "cryptonote_basic/wbag_cache.h"

extern "C" {
  #include "crypto/wbag-hash.h"
}

// ---------------------- Anchors ----------------------
const WbagAnchorSample k_wbag_anchor_samples[] = {
  {770976, "731cabafc8626ce567be7d2ce161d0fefdaa0e6a4f88275513f199362b8d53fd"},
  {770832, "c8710d1ff2b98134c4777e1bbc8dc90ddb6d55a2be850dc904c20218cdb5a970"},
  {770688, "f54b797ed2c44e2c37a455118df2eef4656d159ac18ed8aa5de1affac2f5c20b"},
  {770544, "5197fb807e62e83a5e50e155ddb9975066aca574e22aa6e2f5ccc0a380f1618a"},
  {770400, "4b2981fd2b65aeeb28181ba37c1e49181db4123060f1fa205124aa78e77e4963"},
};
const int k_wbag_anchor_count =
    int(sizeof(k_wbag_anchor_samples) / sizeof(k_wbag_anchor_samples[0]));

// ---------------------- Helpers ----------------------
static inline std::string print_rate(double hs) {
  std::ostringstream o;
  if (hs >= 1000.0) o << std::fixed << std::setprecision(2) << (hs / 1000.0) << " kH/s";
  else              o << std::fixed << std::setprecision(0) << hs << " H/s";
  return o.str();
}

// ---------------------- Cache setup via real code ----------------------
static bool ensure_wbag_cache(int anchor_index,
                              WbagCacheView& out_cpu_view,
                              cryptonote::wbag::CacheView& out_gpu_view,
                              double& gen_seconds)
{
  using cryptonote::wbag::ensure_cache;
  using cryptonote::wbag::cache_view;
  using cryptonote::wbag::epoch;

  if (anchor_index < 0 || anchor_index >= k_wbag_anchor_count) return false;

  const auto& a = k_wbag_anchor_samples[anchor_index];
  crypto::hash anchor{};
  if (!epee::string_tools::hex_to_pod(a.block_id_hex, anchor)) return false;

  const uint64_t ep = epoch(a.height);

  const auto t0 = std::chrono::high_resolution_clock::now();
  if (!ensure_cache(ep, anchor)) return false;
  gen_seconds = std::chrono::duration<double>(std::chrono::high_resolution_clock::now() - t0).count();

  out_gpu_view = cache_view();
  if (!out_gpu_view || !out_gpu_view.data() || out_gpu_view.size() == 0) return false;

  const size_t resident_bytes = WBAG_scratchpad_resident();
  static std::vector<uint8_t> resident;
  resident.resize(resident_bytes);

  const uint64_t rseed = WBAG_resident_seed_from_cache(
                           out_gpu_view.data(), out_gpu_view.size());
  if (WBAG_resident_fill(resident.data(), resident.size(), rseed) != 0) return false;

  out_cpu_view.cache         = out_gpu_view.data();
  out_cpu_view.cache_size    = out_gpu_view.size();
  out_cpu_view.resident      = resident.data();
  out_cpu_view.resident_size = resident.size();
  return true;
}

// ---------------------- CPU bench ----------------------
static double bench_cpu_once(const WbagCacheView& view, int threads, int trials) {
  std::vector<std::thread> th;
  th.reserve(std::max(1, threads));

  std::atomic<int> ready{0};
  std::atomic<bool> go{false};
  const std::vector<uint8_t> base_msg(80, 0x42);

  auto worker = [&](int tid) {
    std::vector<uint8_t> msg = base_msg;
    if (!msg.empty()) msg[0] ^= uint8_t(tid * 31);
    std::array<uint8_t, 32> out{};
    ready.fetch_add(1, std::memory_order_acq_rel);
    while (!go.load(std::memory_order_acquire)) {}
    for (int i = 0; i < trials; ++i) {
      WBAG_hash(msg.data(), msg.size(),
                view.cache,
                view.resident, view.resident_size,
                out.data());
    }
  };

  for (int i = 0; i < threads; ++i) th.emplace_back(worker, i);
  while (ready.load(std::memory_order_acquire) != threads) {}
  const auto t1 = std::chrono::high_resolution_clock::now();
  go.store(true, std::memory_order_release);
  for (auto& x : th) x.join();

  const double sec = std::max(1e-12,
      std::chrono::duration<double>(std::chrono::high_resolution_clock::now() - t1).count());
  return (double(trials) * double(threads)) / sec;
}

// ---------------------- Public C API wrappers ----------------------
extern "C" void wbag_hash(const void* data, size_t length, char* hash) {
  static std::once_flag once;
  static bool ok = false;
  static cryptonote::wbag::CacheView g_cache;
  static std::vector<uint8_t> g_resident;
  static size_t g_resident_size = 0;

  std::call_once(once, [&]{
    const auto& a = k_wbag_anchor_samples[0];
    crypto::hash anchor{};
    if (!epee::string_tools::hex_to_pod(a.block_id_hex, anchor)) return;
    const uint64_t ep = cryptonote::wbag::epoch(a.height);
    if (!cryptonote::wbag::ensure_cache(ep, anchor)) return;
    g_cache = cryptonote::wbag::cache_view();
    if (!g_cache || !g_cache.data() || g_cache.size() == 0) return;

    g_resident_size = WBAG_scratchpad_resident();
    g_resident.resize(g_resident_size);

    const uint64_t rseed = WBAG_resident_seed_from_cache(g_cache.data(), g_cache.size());
    if (WBAG_resident_fill(g_resident.data(), g_resident.size(), rseed) != 0) return;

    ok = true;
  });

  uint8_t out[32]{};
  if (!ok) { std::memset(hash, 0, 32); return; }

  if (WBAG_hash(data, length, g_cache.data(), g_resident.data(), g_resident_size, out) != 0) {
    std::memset(hash, 0, 32);
    return;
  }
  std::memcpy(hash, out, 32);
}

// ---------------------- CLI ----------------------
extern "C" int wbag_hashrate_cli(int argc, char** argv) {
  bool want_cpu = false, want_gpu = false;
  int threads = std::max(1u, std::thread::hardware_concurrency());
  int trials  = 2048;
  int runs    = 1;
  int anchor_index = 0;

  for (int i = 1; i < argc; ++i) {
    const std::string a = argv[i];
    if (a == "--hashrate") continue;
    else if (a == "--cpu") want_cpu = true;
    else if (a == "--gpu") want_gpu = true;
    else if (a == "--threads" && i + 1 < argc) { threads = std::max(1, std::atoi(argv[++i])); }
    else if (a == "--trials"  && i + 1 < argc) { trials  = std::max(1, std::atoi(argv[++i])); }
    else if (a == "--runs"    && i + 1 < argc) { runs    = std::max(1, std::atoi(argv[++i])); }
    else if (a == "--anchor"  && i + 1 < argc) { anchor_index = std::atoi(argv[++i]); }
  }
  if (!(want_cpu || want_gpu)) want_cpu = true;
  if (anchor_index < 0) anchor_index = 0;
  if (anchor_index >= k_wbag_anchor_count) anchor_index = k_wbag_anchor_count - 1;

  WbagCacheView cpu_view{};
  cryptonote::wbag::CacheView gpu_view{};
  double cache_sec = 0.0;
  if (!ensure_wbag_cache(anchor_index, cpu_view, gpu_view, cache_sec)) {
    std::cerr << "Cache build failed\n";
    return 1;
  }

  std::cout << "Mode: " << (want_cpu ? "CPU" : "GPU") << "\n";
  std::cout << "Anchor: [" << anchor_index << "] height=" << k_wbag_anchor_samples[anchor_index].height << "\n";
  std::cout << "Cache Gen Time: " << std::fixed << std::setprecision(3) << cache_sec << " s\n";
  std::cout << "Threads: " << threads << "\n";
  std::cout << "Trials: "  << trials  << "\n\n";
  std::cout << "Hashing...\n";

  if (want_cpu) {
    for (int r = 1; r <= runs; ++r) {
      const double hs = bench_cpu_once(cpu_view, threads, trials);
      std::cout << "[" << r << "/" << runs << "] HR: " << print_rate(hs) << "\n";
    }
    return 0;
  }

#if defined(WBAG_HAVE_WBAG_GPU_BENCH)
  for (int r = 1; r <= runs; ++r) {
    double hs = 0.0;
    const int rc = wbag_gpu_bench_once(gpu_view, threads, trials, hs);
    if (rc != 0) {
      std::cerr << "GPU bench failed (rc=" << rc << ")\n";
      return 1;
    }
    std::cout << "[" << r << "/" << runs << "] HR: " << print_rate(hs) << "\n";
  }
  return 0;
#else
  std::cerr << "GPU support not compiled. Define WBAG_HAVE_WBAG_GPU_BENCH and provide wbag_gpu_bench_once().\n";
  return 2;
#endif
}