// whatsapp_test.cpp: validates detz's native WhatsApp implementation against
// golden vectors captured from the real Baileys client (tests/whatsapp_data,
// copied from jfelipesjc/wa-go, MIT). No network; these are offline byte-exact
// checks of the wire codec, crypto, Noise handshake and Signal session. The
// live pairing/send paths need a phone and are not covered here. See
// docs/whatsapp.md.
#include <cstdio>
#include <fstream>
#include <string>

#include "json.h"
#include "wa_binary.h"
#include "wa_crypto.h"
#include "wa_noise.h"
#include "wa_proto.h"
#include "wa_signal.h"

using namespace detz::wa;
using detz::json::Value;

static int g_fail = 0;
static void check(bool ok, const std::string& name) {
  std::printf("  %s %s\n", ok ? "ok  " : "FAIL", name.c_str());
  if (!ok) g_fail++;
}

static Bytes fromhex(const std::string& h) {
  Bytes b;
  for (size_t i = 0; i + 1 < h.size(); i += 2) b.push_back((uint8_t)std::stoi(h.substr(i, 2), nullptr, 16));
  return b;
}
static std::string tohex(const Bytes& b) {
  std::string s;
  char t[3];
  for (auto c : b) {
    std::snprintf(t, 3, "%02x", c);
    s += t;
  }
  return s;
}
static Bytes b64(const std::string& in) {
  static const std::string T = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
  Bytes out;
  int val = 0, bits = -8;
  for (char c : in) {
    if (c == '=') break;
    auto p = T.find(c);
    if (p == std::string::npos) continue;
    val = (val << 6) | (int)p;
    bits += 6;
    if (bits >= 0) {
      out.push_back((uint8_t)((val >> bits) & 0xff));
      bits -= 8;
    }
  }
  return out;
}
static std::string slurp(const std::string& p) {
  std::ifstream f(p);
  return std::string((std::istreambuf_iterator<char>(f)), std::istreambuf_iterator<char>());
}
static std::array<uint8_t, 32> a32(const Bytes& b) {
  std::array<uint8_t, 32> a{};
  for (size_t i = 0; i < 32 && i < b.size(); ++i) a[i] = b[i];
  return a;
}
static SignalKey sk33(const Bytes& b) {
  SignalKey k{};
  for (size_t i = 0; i < 33 && i < b.size(); ++i) k[i] = b[i];
  return k;
}
static WaCrypto::KeyPair kpair(const Value& o) {
  WaCrypto::KeyPair k;
  k.priv = a32(fromhex(o.get("priv").str()));
  Bytes pb = fromhex(o.get("pub").str());
  k.pub = a32(Bytes(pb.begin() + (pb.size() == 33 ? 1 : 0), pb.end()));
  return k;
}

static const char* DIR = nullptr;
static std::string path(const std::string& f) { return std::string(DIR) + "/" + f; }

// ---- crypto (RFC vectors + XEdDSA) ----
static void test_crypto() {
  std::printf("[crypto]\n");
  Bytes ikm(22, 0x0b), salt;
  for (int i = 0; i < 13; i++) salt.push_back(i);
  Bytes info;
  for (int i = 0xf0; i <= 0xf9; i++) info.push_back(i);
  check(tohex(WaCrypto::hkdf_sha256(ikm, salt, info, 42)) ==
            "3cb25f25faacd57a90434f64d0362f2a2d2d0a90cf1a5a4c5db02d56ecc4c5bf34007208d5b887185865",
        "HKDF-SHA256 RFC 5869 case 1");
  check(tohex(WaCrypto::pbkdf2_sha256(Bytes{'p', 'a', 's', 's', 'w', 'o', 'r', 'd'}, Bytes{'s', 'a', 'l', 't'}, 1, 32)) ==
            "120fb6cffcf8b32c43e7225256c4f837a86548c92ccc35480805987cb70be17b",
        "PBKDF2-SHA256");
  auto a = WaCrypto::generate_x25519(), b = WaCrypto::generate_x25519();
  check(WaCrypto::x25519(a.priv, b.pub) == WaCrypto::x25519(b.priv, a.pub), "X25519 DH agreement");
  bool xok = true;
  for (int i = 0; i < 64; ++i) {
    auto kp = WaCrypto::generate_x25519();
    Bytes m = WaCrypto::random_bytes(1 + i);
    auto sig = WaCrypto::xeddsa_sign(kp.priv, m);
    if (!WaCrypto::xeddsa_verify(kp.pub, m, sig)) xok = false;
    sig[i % 64] ^= 1;
    if (WaCrypto::xeddsa_verify(kp.pub, m, sig)) xok = false;
  }
  check(xok, "XEdDSA sign/verify roundtrip + tamper (64 keys)");
}

// ---- binary codec ----
static void test_codec() {
  std::printf("[codec]\n");
  std::ifstream f(path("codec_battery.jsonl"));
  std::string line;
  int n = 0, ok = 0;
  std::function<WaNode(const Value&)> build = [&](const Value& t) -> WaNode {
    WaNode node(t.get("tag").str());
    for (const auto& m : t.get("attrs").members()) node.attrs[m.first] = m.second.str();
    const Value& c = t.get("content");
    if (c.is_string())
      node.content = fromhex(c.str());
    else if (c.is_array()) {
      WaNode::Children ch;
      for (const auto& e : c.items()) ch.push_back(build(e));
      node.content = ch;
    }
    return node;
  };
  std::function<bool(const WaNode&, const WaNode&)> eq = [&](const WaNode& x, const WaNode& y) -> bool {
    if (x.tag != y.tag || x.attrs != y.attrs || x.content.index() != y.content.index()) return false;
    if (auto* a = x.bytes()) return *a == *y.bytes();
    if (auto* a = x.children()) {
      auto* b = y.children();
      if (a->size() != b->size()) return false;
      for (size_t i = 0; i < a->size(); ++i)
        if (!eq((*a)[i], (*b)[i])) return false;
    }
    return true;
  };
  while (std::getline(f, line)) {
    if (line.empty()) continue;
    ++n;
    Value v = Value::loads(line);
    WaNode want = build(v.get("decoded_tree"));
    WaNode got = WaBinary::decode(WaBinary::unpack(fromhex(v.get("encoded_hex").str())));
    Bytes re = WaBinary::encode(got);
    Bytes framed;
    framed.push_back(0);
    framed.insert(framed.end(), re.begin(), re.end());
    WaNode got2 = WaBinary::decode(WaBinary::unpack(framed));
    if (eq(got, want) && eq(got2, want)) ++ok;
  }
  check(ok == n && n == 19, "codec battery " + std::to_string(ok) + "/" + std::to_string(n) + " (decode + re-encode)");
}

// ---- ClientPayload byte-exact ----
static void test_client_payload() {
  std::printf("[proto]\n");
  Value d = Value::loads(slurp(path("client_payload.json")));
  const Value& fl = d.get("fields");
  const Value& ua = fl.get("userAgent");
  const Value& dp = fl.get("devicePairingData");
  WaClientPayload p;
  p.has_passive = true;
  p.has_user_agent = true;
  auto& u = p.user_agent;
  u.app_version.has = true;
  u.app_version.primary = ua.get("appVersion").get("primary").as_int();
  u.app_version.secondary = ua.get("appVersion").get("secondary").as_int();
  u.app_version.tertiary = ua.get("appVersion").get("tertiary").as_int();
  p.has_web_info = true;
  p.has_connect = true;
  WaDevicePairingData pd;
  pd.e_regid = b64(dp.get("eRegid").str());
  pd.e_keytype = b64(dp.get("eKeytype").str());
  pd.e_ident = b64(dp.get("eIdent").str());
  pd.e_skey_id = b64(dp.get("eSkeyId").str());
  pd.e_skey_val = b64(dp.get("eSkeyVal").str());
  pd.e_skey_sig = b64(dp.get("eSkeySig").str());
  pd.build_hash = b64(dp.get("buildHash").str());
  pd.device_props = b64(dp.get("deviceProps").str());
  p.pairing = pd;
  p.has_pull = true;
  check(tohex(p.encode()) == d.get("payloadHex").str(), "ClientPayload byte-exact (registration)");
}

// ---- Noise handshake ----
static void test_noise() {
  std::printf("[noise]\n");
  Value d = Value::loads(slurp(path("noise.json")));
  auto ephPriv = a32(fromhex(d.get("ephemeralKeyPriv").str()));
  auto ephPub = a32(fromhex(d.get("ephemeralKeyPub").str()));
  auto stPriv = a32(fromhex(d.get("authStaticKeyPriv").str()));
  auto stPub = a32(fromhex(d.get("authStaticKeyPub").str()));
  Bytes shFrame = fromhex(d.get("serverHelloFrameHex").str());
  size_t pos = 0;
  auto sh = WaFraming::unframe(shFrame, pos);
  WaNoise noise(ephPub);
  auto res = noise.client_handshake(ephPriv, ephPub, stPriv, stPub, *sh, Bytes{1, 2, 3, 4});
  check(res.has_value(), "handshake completed (server static+payload decrypted)");
  if (res) {
    check(tohex(WaFraming::frame(res->client_hello, true)) == d.get("clientHelloFrameHex").str(),
          "clientHello frame byte-exact");
    check(res->server_static.size() == 32, "server static key is 32 bytes");
    check(parse_cert_chain(res->server_payload).has_value(), "server payload parses as cert chain");
    check(noise.in_transport(), "transport keys derived");
  }
}

// ---- Signal session ----
static void test_signal() {
  std::printf("[signal]\n");
  Value d = Value::loads(slurp(path("session_ab.json")));
  auto& al = d.get("alice");
  auto& bo = d.get("bob");
  auto& eph = d.get("ephemeralsGenerated");
  WaSession::InitiatorParams p;
  p.local_identity = kpair(al.get("identityKeyPair"));
  p.local_base = kpair(al.get("baseKey"));
  p.remote_identity = sk33(fromhex(bo.get("identityKeyPair").get("pub").str()));
  p.remote_signed_pre = sk33(fromhex(bo.get("signedPreKey").get("keyPair").get("pub").str()));
  p.remote_pre_key = sk33(fromhex(bo.get("preKey").get("keyPair").get("pub").str()));
  p.has_pre_key = true;
  WaSession s = WaSession::initiate(p, kpair(eph.at(1)));
  s.set_pending(al.get("registrationId").as_int(), signal_pub(p.local_base.pub),
                bo.get("signedPreKey").get("id").as_int(), true, bo.get("preKey").get("id").as_int());
  WaCrypto::KeyPair nextR = kpair(eph.at(3));
  s.set_ratchet_generator([&]() { return nextR; });
  const Value& ex = d.get("exchanges");
  auto pt = [&](int i) {
    std::string t = ex.at(i).get("plaintext").str();
    return Bytes(t.begin(), t.end());
  };
  auto want = [&](int i) { return ex.at(i).get("ciphertextHex").str(); };
  auto o1 = s.encrypt(pt(0));
  check(o1.type == "pkmsg" && tohex(o1.serialized) == want(0), "exchange #1 pkmsg byte-exact");
  auto p2 = s.decrypt(fromhex(want(1)));
  check(p2 && std::string(p2->begin(), p2->end()) == ex.at(1).get("plaintext").str(), "exchange #2 decrypt");
  auto o3 = s.encrypt(pt(2));
  check(o3.type == "msg" && tohex(o3.serialized) == want(2), "exchange #3 msg byte-exact");
  auto o4 = s.encrypt(pt(3));
  check(o4.type == "msg" && tohex(o4.serialized) == want(3), "exchange #4 msg byte-exact");
  // session survives a serialize/deserialize round trip
  Bytes blob = s.serialize();
  check(WaSession::deserialize(blob).has_value(), "session serialize/deserialize");
}

int main(int argc, char** argv) {
  DIR = argc > 1 ? argv[1] : "tests/whatsapp_data";
  test_crypto();
  test_codec();
  test_client_payload();
  test_noise();
  test_signal();
  std::printf("%s\n", g_fail == 0 ? "whatsapp: ALL OK" : "whatsapp: FAILURES");
  return g_fail == 0 ? 0 : 1;
}