import { describe, it, expect } from "vitest";
import * as soliton from "../pkg/soliton_wasm.js";

const utf8 = (s) => new TextEncoder().encode(s);
const fromUtf8 = (buf) => new TextDecoder().decode(buf);
const randomBytes = (n) => crypto.getRandomValues(new Uint8Array(n));
const bytesEqual = (a, b) => {
  if (a.length !== b.length) return false;
  for (let i = 0; i < a.length; i++) if (a[i] !== b[i]) return false;
  return true;
};

// ── Version ─────────────────────────────────────────────────────────────

describe("version", () => {
  it("returns non-empty string", () => {
    const v = soliton.version();
    expect(v.length).toBeGreaterThan(0);
  });
});

// ── Primitives ──────────────────────────────────────────────────────────

describe("primitives", () => {
  it("sha3_256 known vector — empty", () => {
    const hash = soliton.sha3_256(new Uint8Array(0));
    expect(hash.length).toBe(32);
    expect(hash[0]).toBe(0xa7);
    expect(hash[1]).toBe(0xff);
  });

  it("sha3_256 known vector — abc", () => {
    const hash = soliton.sha3_256(utf8("abc"));
    expect(hash.length).toBe(32);
    expect(hash[0]).toBe(0x3a);
    expect(hash[1]).toBe(0x98);
  });

  it("fingerprintHex", () => {
    const hex = soliton.fingerprintHex(utf8("abc"));
    expect(hex.length).toBe(64);
    expect(typeof hex).toBe("string");
  });

  it("hmacSha3_256 and verify", () => {
    const key = new Uint8Array(32).fill(0x0b);
    const tag = soliton.hmacSha3_256(key, utf8("Hi There"));
    expect(tag.length).toBe(32);
    expect(soliton.hmacSha3_256Verify(tag, tag)).toBe(true);

    const bad = new Uint8Array(tag);
    bad[0] ^= 0xff;
    expect(soliton.hmacSha3_256Verify(tag, bad)).toBe(false);
  });

  it("hkdfSha3_256", () => {
    const salt = new Uint8Array(32);
    const ikm = new Uint8Array(32).fill(0x0b);
    const okm = soliton.hkdfSha3_256(salt, ikm, utf8("test"), 64);
    expect(okm.length).toBe(64);
  });

  it("xwingKeygen", () => {
    const kp = soliton.xwingKeygen();
    expect(kp.publicKey.length).toBe(1216);
    expect(kp.secretKey.length).toBe(2432);
  });

  it("argon2id", () => {
    const salt = randomBytes(16);
    const key1 = soliton.argon2id(utf8("password"), salt, 19456, 2, 1, 32);
    const key2 = soliton.argon2id(utf8("password"), salt, 19456, 2, 1, 32);
    expect(key1.length).toBe(32);
    expect(bytesEqual(key1, key2)).toBe(true);

    const key3 = soliton.argon2id(utf8("wrong"), salt, 19456, 2, 1, 32);
    expect(bytesEqual(key1, key3)).toBe(false);
  });
});

// ── Identity ────────────────────────────────────────────────────────────

describe("identity", () => {
  it("keygen", () => {
    const id = new soliton.Identity();
    expect(id.publicKey().length).toBe(3200);
    expect(id.secretKey().length).toBe(2496);
    id.free();
  });

  it("fingerprint", () => {
    const id = new soliton.Identity();
    const fp = id.fingerprint();
    expect(fp.length).toBe(32);

    // Verify fingerprint matches sha3_256(pk).
    const expected = soliton.sha3_256(id.publicKey());
    expect(bytesEqual(fp, expected)).toBe(true);
    id.free();
  });

  it("sign and verify", () => {
    const id = new soliton.Identity();
    const msg = utf8("test message");
    const sig = id.sign(msg);
    expect(sig.length).toBe(3373);
    id.verify(msg, sig); // throws on failure
    id.free();
  });

  it("verify wrong message throws", () => {
    const id = new soliton.Identity();
    const sig = id.sign(utf8("correct"));
    expect(() => id.verify(utf8("wrong"), sig)).toThrow();
    id.free();
  });

  it("verify wrong key throws", () => {
    const alice = new soliton.Identity();
    const bob = new soliton.Identity();
    const sig = alice.sign(utf8("hello"));
    expect(() => bob.verify(utf8("hello"), sig)).toThrow();
    alice.free();
    bob.free();
  });

  it("fromBytes round trip", () => {
    const id = new soliton.Identity();
    const pk = id.publicKey();
    const sk = id.secretKey();
    id.free();

    const id2 = soliton.Identity.fromBytes(pk, sk);
    const sig = id2.sign(utf8("roundtrip"));
    id2.verify(utf8("roundtrip"), sig);
    id2.free();
  });

  it("fromPublicBytes cannot sign", () => {
    const id = new soliton.Identity();
    const pk = id.publicKey();
    id.free();

    const pub_only = soliton.Identity.fromPublicBytes(pk);
    expect(() => pub_only.sign(utf8("test"))).toThrow();
    pub_only.free();
  });
});

// ── Auth ────────────────────────────────────────────────────────────────

describe("auth", () => {
  it("challenge-respond-verify round trip", () => {
    const id = new soliton.Identity();
    const pk = id.publicKey();
    const sk = id.secretKey();

    const { ciphertext, token } = soliton.authChallenge(pk);
    const proof = soliton.authRespond(sk, ciphertext);
    expect(soliton.authVerify(token, proof)).toBe(true);
    id.free();
  });

  it("wrong proof fails", () => {
    const id = new soliton.Identity();
    const { token } = soliton.authChallenge(id.publicKey());
    expect(soliton.authVerify(token, new Uint8Array(32))).toBe(false);
    id.free();
  });
});

// ── Verification ────────────────────────────────────────────────────────

describe("verification", () => {
  it("phrase is symmetric", () => {
    const alice = new soliton.Identity();
    const bob = new soliton.Identity();

    const p1 = soliton.verificationPhrase(alice.publicKey(), bob.publicKey());
    const p2 = soliton.verificationPhrase(bob.publicKey(), alice.publicKey());
    expect(p1.length).toBeGreaterThan(0);
    expect(p1).toBe(p2);

    alice.free();
    bob.free();
  });
});

// ── Storage ─────────────────────────────────────────────────────────────

describe("storage", () => {
  it("encrypt/decrypt round trip", () => {
    const key = randomBytes(32);
    const ring = new soliton.StorageKeyRing(1, key);
    const pt = utf8("encrypted storage data");
    const blob = ring.encryptBlob("channel-1", "segment-0", pt);
    const result = ring.decryptBlob("channel-1", "segment-0", blob);
    expect(bytesEqual(result, pt)).toBe(true);
    ring.free();
  });

  it("wrong channel fails", () => {
    const ring = new soliton.StorageKeyRing(1, randomBytes(32));
    const blob = ring.encryptBlob("ch-1", "seg", utf8("data"));
    expect(() => ring.decryptBlob("ch-2", "seg", blob)).toThrow();
    ring.free();
  });

  it("key rotation", () => {
    const ring = new soliton.StorageKeyRing(1, randomBytes(32));
    const blob_v1 = ring.encryptBlob("ch", "seg", utf8("v1"));
    ring.addKey(2, randomBytes(32), true);
    const blob_v2 = ring.encryptBlob("ch", "seg", utf8("v2"));

    expect(fromUtf8(ring.decryptBlob("ch", "seg", blob_v1))).toBe("v1");
    expect(fromUtf8(ring.decryptBlob("ch", "seg", blob_v2))).toBe("v2");
    ring.free();
  });

  it("dm queue round trip", () => {
    const ring = new soliton.StorageKeyRing(1, randomBytes(32));
    const fp = randomBytes(32);
    const pt = utf8("queued DM");
    const blob = ring.encryptDmQueue(fp, "batch-1", pt);
    const result = ring.decryptDmQueue(fp, "batch-1", blob);
    expect(bytesEqual(result, pt)).toBe(true);
    ring.free();
  });

  it("compress round trip", () => {
    const ring = new soliton.StorageKeyRing(1, randomBytes(32));
    const data = utf8("compressible ".repeat(100));
    const blob = ring.encryptBlob("ch", "seg", data, true);
    const result = ring.decryptBlob("ch", "seg", blob);
    expect(bytesEqual(result, data)).toBe(true);
    ring.free();
  });
});

// ── Streaming AEAD ──────────────────────────────────────────────────────

describe("streaming", () => {
  it("single chunk round trip", () => {
    const key = randomBytes(32);
    const enc = new soliton.StreamEncryptor(key);
    const header = enc.header();
    expect(header.length).toBe(26);
    const ct = enc.encryptChunk(utf8("hello stream"), true);
    enc.free();

    const dec = new soliton.StreamDecryptor(key, header);
    const { plaintext, isLast } = dec.decryptChunk(ct);
    expect(fromUtf8(plaintext)).toBe("hello stream");
    expect(isLast).toBe(true);
    dec.free();
  });

  it("encrypt_at/decrypt_at", () => {
    const key = randomBytes(32);
    const enc = new soliton.StreamEncryptor(key);
    const header = enc.header();
    const ct = enc.encryptChunkAt(5n, utf8("at index five"), true);
    enc.free();

    const dec = new soliton.StreamDecryptor(key, header);
    const { plaintext } = dec.decryptChunkAt(5n, ct);
    expect(fromUtf8(plaintext)).toBe("at index five");
    dec.free();
  });

  it("wrong key fails", () => {
    const enc = new soliton.StreamEncryptor(randomBytes(32));
    const header = enc.header();
    const ct = enc.encryptChunk(utf8("data"), true);
    enc.free();

    const dec = new soliton.StreamDecryptor(randomBytes(32), header);
    expect(() => dec.decryptChunk(ct)).toThrow();
    dec.free();
  });
});

// ── KEX + Ratchet ───────────────────────────────────────────────────────

describe("kex", () => {
  it("sign and verify bundle", () => {
    const bob = new soliton.Identity();
    const { publicKey: spkPub } = soliton.xwingKeygen();
    const sig = soliton.kexSignPrekey(bob.secretKey(), spkPub);
    expect(sig.length).toBe(3373);

    soliton.kexVerifyBundle(
      bob.publicKey(), bob.publicKey(),
      spkPub, 1, sig, "lo-crypto-v1",
    );
    bob.free();
  });

  it("verify bundle wrong key throws", () => {
    const bob = new soliton.Identity();
    const eve = new soliton.Identity();
    const { publicKey: spkPub } = soliton.xwingKeygen();
    const sig = soliton.kexSignPrekey(bob.secretKey(), spkPub);

    expect(() => soliton.kexVerifyBundle(
      bob.publicKey(), eve.publicKey(),
      spkPub, 1, sig, "lo-crypto-v1",
    )).toThrow();
    bob.free();
    eve.free();
  });
});

describe("first message", () => {
  it("encrypt/decrypt round trip", () => {
    const ck = randomBytes(32);
    const aad = utf8("test-aad");
    const pt = utf8("first application message");

    const { encryptedPayload, ratchetInitKey: rikA } =
      soliton.Ratchet.encryptFirstMessage(ck, pt, aad);
    const { plaintext, ratchetInitKey: rikB } =
      soliton.Ratchet.decryptFirstMessage(ck, encryptedPayload, aad);

    expect(fromUtf8(plaintext)).toBe("first application message");
    expect(bytesEqual(rikA, rikB)).toBe(true);
  });

  it("wrong key throws", () => {
    const ck = randomBytes(32);
    const { encryptedPayload } = soliton.Ratchet.encryptFirstMessage(
      ck, utf8("hello"), utf8("aad"),
    );
    expect(() =>
      soliton.Ratchet.decryptFirstMessage(randomBytes(32), encryptedPayload, utf8("aad")),
    ).toThrow();
  });
});

describe("full kex + ratchet", () => {
  it("complete session lifecycle", () => {
    const alice = new soliton.Identity();
    const bob = new soliton.Identity();
    const { publicKey: spkPub, secretKey: spkSk } = soliton.xwingKeygen();
    const spkSig = soliton.kexSignPrekey(bob.secretKey(), spkPub);

    // Alice initiates.
    const initiated = soliton.kexInitiate(
      alice.publicKey(), alice.secretKey(),
      bob.publicKey(), spkPub, 1, spkSig, "lo-crypto-v1",
    );

    // Bob receives.
    const siEncoded = initiated.sessionInitEncoded();
    const received = soliton.kexReceive(
      bob.publicKey(), bob.secretKey(), alice.publicKey(),
      siEncoded, initiated.senderSig(), spkSk,
    );

    // First message.
    const aad = soliton.kexBuildFirstMessageAad(
      initiated.senderFingerprint(),
      initiated.recipientFingerprint(),
      siEncoded,
    );
    const { encryptedPayload, ratchetInitKey: rikA } =
      soliton.Ratchet.encryptFirstMessage(initiated.takeInitialChainKey(), utf8("hello bob"), aad);
    const { plaintext: firstPt, ratchetInitKey: rikB } =
      soliton.Ratchet.decryptFirstMessage(received.takeInitialChainKey(), encryptedPayload, aad);

    expect(fromUtf8(firstPt)).toBe("hello bob");

    // Init ratchets.
    const aliceR = soliton.Ratchet.initAlice(
      initiated.takeRootKey(), rikA,
      alice.fingerprint(), bob.fingerprint(),
      received.peerEk(), initiated.ekSk(),
    );
    const bobR = soliton.Ratchet.initBob(
      received.takeRootKey(), rikB,
      bob.fingerprint(), alice.fingerprint(),
      received.peerEk(),
    );

    // Alice sends to Bob.
    const { header, ciphertext } = aliceR.encrypt(utf8("message 1"));
    const pt1 = bobR.decrypt(header, ciphertext);
    expect(fromUtf8(pt1)).toBe("message 1");

    // Bob replies (direction change).
    const { header: hdr2, ciphertext: ct2 } = bobR.encrypt(utf8("reply 1"));
    const pt2 = aliceR.decrypt(hdr2, ct2);
    expect(fromUtf8(pt2)).toBe("reply 1");

    // Cleanup.
    aliceR.free();
    bobR.free();
    initiated.free();
    received.free();
    alice.free();
    bob.free();
  });

  it("ratchet serialize/deserialize", () => {
    const alice = new soliton.Identity();
    const bob = new soliton.Identity();
    const { publicKey: spkPub, secretKey: spkSk } = soliton.xwingKeygen();
    const spkSig = soliton.kexSignPrekey(bob.secretKey(), spkPub);

    const initiated = soliton.kexInitiate(
      alice.publicKey(), alice.secretKey(),
      bob.publicKey(), spkPub, 1, spkSig, "lo-crypto-v1",
    );
    const siEncoded = initiated.sessionInitEncoded();
    const received = soliton.kexReceive(
      bob.publicKey(), bob.secretKey(), alice.publicKey(),
      siEncoded, initiated.senderSig(), spkSk,
    );

    const aad = soliton.kexBuildFirstMessageAad(
      initiated.senderFingerprint(), initiated.recipientFingerprint(), siEncoded,
    );
    const { ratchetInitKey: rik } =
      soliton.Ratchet.encryptFirstMessage(initiated.takeInitialChainKey(), utf8("x"), aad);

    const aliceR = soliton.Ratchet.initAlice(
      initiated.takeRootKey(), rik,
      alice.fingerprint(), bob.fingerprint(),
      received.peerEk(), initiated.ekSk(),
    );

    aliceR.encrypt(utf8("advance state"));
    const { blob, epoch } = aliceR.toBytes();
    expect(blob.length).toBeGreaterThan(0);
    expect(epoch).toBeGreaterThanOrEqual(1n);

    const restored = soliton.Ratchet.fromBytes(blob, 0n);
    const { ciphertext } = restored.encrypt(utf8("after restore"));
    expect(ciphertext.length).toBeGreaterThan(0);

    restored.free();
    initiated.free();
    received.free();
    alice.free();
    bob.free();
  });

  it("derive call keys", () => {
    const alice = new soliton.Identity();
    const bob = new soliton.Identity();
    const { publicKey: spkPub, secretKey: spkSk } = soliton.xwingKeygen();
    const spkSig = soliton.kexSignPrekey(bob.secretKey(), spkPub);

    const initiated = soliton.kexInitiate(
      alice.publicKey(), alice.secretKey(),
      bob.publicKey(), spkPub, 1, spkSig, "lo-crypto-v1",
    );
    const siEncoded = initiated.sessionInitEncoded();
    const received = soliton.kexReceive(
      bob.publicKey(), bob.secretKey(), alice.publicKey(),
      siEncoded, initiated.senderSig(), spkSk,
    );

    const aad = soliton.kexBuildFirstMessageAad(
      initiated.senderFingerprint(), initiated.recipientFingerprint(), siEncoded,
    );
    const { ratchetInitKey: rik } =
      soliton.Ratchet.encryptFirstMessage(initiated.takeInitialChainKey(), utf8("x"), aad);

    const aliceR = soliton.Ratchet.initAlice(
      initiated.takeRootKey(), rik,
      alice.fingerprint(), bob.fingerprint(),
      received.peerEk(), initiated.ekSk(),
    );

    const kemSs = randomBytes(32);
    const callId = randomBytes(16);
    const keys = aliceR.deriveCallKeys(kemSs, callId);
    const send = keys.sendKey();
    const recv = keys.recvKey();
    expect(send.length).toBe(32);
    expect(recv.length).toBe(32);
    expect(bytesEqual(send, recv)).toBe(false);

    keys.advance();
    const send2 = keys.sendKey();
    expect(bytesEqual(send, send2)).toBe(false);

    keys.free();
    aliceR.free();
    initiated.free();
    received.free();
    alice.free();
    bob.free();
  });
});

// ── Additional Coverage ─────────────────────────────────────────────────

describe("additional", () => {
  it("hybridVerify standalone", () => {
    const id = new soliton.Identity();
    const msg = utf8("test");
    const sig = id.sign(msg);
    soliton.hybridVerify(id.publicKey(), msg, sig);
    expect(() => soliton.hybridVerify(id.publicKey(), utf8("wrong"), sig)).toThrow();
    id.free();
  });

  it("storage removeKey", () => {
    const ring = new soliton.StorageKeyRing(1, randomBytes(32));
    ring.addKey(2, randomBytes(32), true);
    ring.removeKey(1);
    const blob = ring.encryptBlob("ch", "seg", utf8("data"));
    expect(fromUtf8(ring.decryptBlob("ch", "seg", blob))).toBe("data");
    ring.free();
  });

  it("errors are Error instances", () => {
    try {
      soliton.authVerify(new Uint8Array(10), new Uint8Array(32));
    } catch (e) {
      expect(e instanceof Error).toBe(true);
      expect(typeof e.message).toBe("string");
      return;
    }
    expect.unreachable("should have thrown");
  });
});