"""
Cross-language parity for the keypair bundle KDF.

The bundle is the one thing a user carries between browsers, and the passphrase
is all that stands between it and whoever holds the disk of a node they joined
(finding C4). It moved from PBKDF2-SHA512 to Argon2id for that reason — PBKDF2 is
compute-only, which is what makes it cheap on a GPU.

Two implementations now have to agree byte for byte: the vendored WebAssembly the
browser runs, and `argon2-cffi` used by the QE harness. A disagreement would not
show up as an error — it would show up as a bundle nobody can open, which is
somebody's account gone.

Skipped when node or argon2-cffi is missing; that is a coverage gap, not a pass.
"""

import hashlib
import json
import shutil
import subprocess
from pathlib import Path

import pytest

STATIC = Path(__file__).resolve().parents[1] / "src" / "meshbay_hub" / "static"
VENDOR = STATIC / "vendor"

try:
    from argon2.low_level import Type, hash_secret_raw
    HAVE_ARGON2 = True
except ImportError:
    HAVE_ARGON2 = False

pytestmark = pytest.mark.skipif(
    shutil.which("node") is None
    or not (VENDOR / "argon2.min.js").exists()
    or not HAVE_ARGON2,
    reason="node, the vendored argon2, or argon2-cffi is unavailable",
)

# Parameters must match keyderive.js. If someone tunes them there and not here,
# this test fails — which is the point: changing them silently orphans every
# bundle already written.
MEM_KIB, TIME_COST, LANES = 131072, 3, 1

CASES = ["alice", "grenet", "utilisateur-é", ""]
PASSWORDS = ["correct horse battery staple", "p", "üñïçø∂é ✓ 🔐"]

_HARNESS = r"""
const fs = require('fs'), webcrypto = require('crypto').webcrypto;
global.self = global; global.crypto = webcrypto;
// The browser uses the copy inlined in the bundle; under node the emscripten
// loader looks for a file, so hand it the same bytes explicitly.
global.Module = { wasmBinary: fs.readFileSync(process.argv[2]) };
const argon2 = require(process.argv[3]);

(async () => {
  const input = JSON.parse(fs.readFileSync(process.argv[4], 'utf8'));
  const out = [];
  for (const v of input) {
    const salt = new Uint8Array(await webcrypto.subtle.digest(
      'SHA-256', new TextEncoder().encode(`meshbay:bundle:v2:${v.username}`)
    )).slice(0, 16);
    const r = await argon2.hash({
      pass: v.password, salt,
      time: v.time, mem: v.mem, parallelism: v.lanes,
      hashLen: 32, type: argon2.ArgonType.Argon2id,
    });
    out.push(Buffer.from(r.hash).toString('hex'));
  }
  process.stdout.write(JSON.stringify(out));
})();
"""


@pytest.fixture(scope="module")
def js_hashes(tmp_path_factory):
    d = tmp_path_factory.mktemp("kdf")
    harness = d / "harness.cjs"
    harness.write_text(_HARNESS, encoding="utf-8")
    vectors = [
        {"username": u, "password": p,
         "mem": MEM_KIB, "time": TIME_COST, "lanes": LANES}
        for u in CASES for p in PASSWORDS
    ]
    payload = d / "vectors.json"
    payload.write_text(json.dumps(vectors), encoding="utf-8")

    proc = subprocess.run(
        ["node", str(harness), str(VENDOR / "argon2.wasm"),
         str(VENDOR / "argon2.min.js"), str(payload)],
        capture_output=True, text=True, encoding="utf-8", timeout=300,
    )
    if proc.returncode != 0:
        pytest.fail(f"node harness failed:\n{proc.stderr[-2000:]}")
    return vectors, json.loads(proc.stdout)


def _python_hash(username: str, password: str) -> str:
    salt = hashlib.sha256(f"meshbay:bundle:v2:{username}".encode()).digest()[:16]
    return hash_secret_raw(
        password.encode(), salt, time_cost=TIME_COST, memory_cost=MEM_KIB,
        parallelism=LANES, hash_len=32, type=Type.ID,
    ).hex()


def test_bundle_key_matches_across_languages(js_hashes):
    vectors, js = js_hashes
    for i, v in enumerate(vectors):
        assert js[i] == _python_hash(v["username"], v["password"]), (
            f"argon2id disagrees for username={v['username']!r} — a bundle "
            f"written by one implementation would be unreadable by the other"
        )


def test_the_salt_separates_users(js_hashes):
    """Two accounts with the same passphrase must not share a bundle key."""
    assert _python_hash("alice", "same passphrase") != \
        _python_hash("bob", "same passphrase")


def test_parameters_still_match_the_client():
    """
    The numbers live in keyderive.js; this test is the second copy. Tuning one
    without the other orphans every bundle already written, so make it fail.
    """
    source = (STATIC / "keyderive.js").read_text(encoding="utf-8")
    assert f"ARGON2_MEM_KIB = {MEM_KIB}" in source
    assert f"ARGON2_TIME    = {TIME_COST}" in source
    assert f"ARGON2_LANES   = {LANES}" in source
    assert "meshbay:bundle:v2:" in source


# ── The whole chain: A, the pepper, M, the node key, the playlist key ─────────
#
# The real keyderive.js and playlist-crypto.js, over the real WebAssembly
# Argon2, against a reference written from the specification with nothing
# shared: argon2-cffi, `cryptography`'s HKDF and AES-GCM. Down to opening an
# MBK3 bundle, so the format and its associated data agree too.

_CHAIN_HARNESS = r"""
const fs = require('fs'), path = require('path'), url = require('url');
const webcrypto = require('crypto').webcrypto;
global.self = global; global.window = global; global.crypto = webcrypto;
global.Module = { wasmBinary: fs.readFileSync(process.argv[2]) };
global.argon2 = require(process.argv[3]);
eval(fs.readFileSync(process.argv[4], 'utf8'));
const K = window.MeshBayKeys;
const fp = async (key) => Buffer.from(await webcrypto.subtle.encrypt(
  { name: 'AES-GCM', iv: new Uint8Array(12) }, key, new Uint8Array(16))).toString('hex');
(async () => {
  const { derivePlaylistKey } = await import(url.pathToFileURL(process.argv[5]).href);
  const out = [];
  for (const v of JSON.parse(fs.readFileSync(process.argv[6], 'utf8'))) {
    const sk = await K.deriveBundleSessionKey(v.password, v.username, v.user_id, v.pepper, 1);
    const kNode = await K.nodeBundleKey(sk, v.node_pk);
    const bundle = await K.encryptBundle(
      Buffer.from('ed-private'), Buffer.from('x-private'), kNode,
      { userId: v.user_id, nodePk: v.node_pk, pepperVersion: 1 });
    out.push({ node: await fp(kNode), playlists: await fp(await derivePlaylistKey(sk.v3)),
               bundle });
  }
  process.stdout.write(JSON.stringify(out));
})().catch((e) => { console.error(e); process.exit(1); });
"""

CHAIN = [
    {"username": "alice", "password": "correct horse battery staple",
     "user_id": "0b4f6f0e-5d7e-4e8a-9d2b-6a1c1b9e2f11", "node_pk": "Tm9kZUtleUE="},
    {"username": "utilisateur-é", "password": "üñïçø∂é ✓ 🔐",
     "user_id": "7d1e0c2a-3b4c-4d5e-8f60-718293a4b5c6", "node_pk": "Tm9kZUtleUI="},
]


def _hkdf(ikm: bytes, info: str) -> bytes:
    from cryptography.hazmat.primitives.hashes import SHA256
    from cryptography.hazmat.primitives.kdf.hkdf import HKDF
    return HKDF(algorithm=SHA256(), length=32, salt=None, info=info.encode()).derive(ikm)


def _fp(key: bytes) -> str:
    from cryptography.hazmat.primitives.ciphers.aead import AESGCM
    return AESGCM(key).encrypt(bytes(12), bytes(16), None).hex()


@pytest.fixture(scope="module")
def chain(tmp_path_factory):
    import base64
    d = tmp_path_factory.mktemp("chain")
    vectors = [{**v, "pepper": base64.b64encode(bytes([i + 1]) * 32).decode()}
               for i, v in enumerate(CHAIN)]
    (d / "harness.cjs").write_text(_CHAIN_HARNESS, encoding="utf-8")
    (d / "vectors.json").write_text(json.dumps(vectors), encoding="utf-8")
    proc = subprocess.run(
        ["node", str(d / "harness.cjs"), str(VENDOR / "argon2.wasm"),
         str(VENDOR / "argon2.min.js"), str(STATIC / "keyderive.js"),
         str(STATIC / "playlist-crypto.js"), str(d / "vectors.json")],
        capture_output=True, text=True, encoding="utf-8", timeout=300)
    if proc.returncode != 0:
        pytest.fail(f"node harness failed:\n{proc.stderr[-2000:]}")
    return vectors, json.loads(proc.stdout)


def _reference(v: dict) -> tuple[bytes, bytes]:
    import base64
    a = bytes.fromhex(_python_hash(v["username"], v["password"]))
    m = _hkdf(a + base64.b64decode(v["pepper"]), f"meshbay:bundle-master:v3|{v['user_id']}")
    return (_hkdf(m, f"meshbay:bundle:v3|node|{v['node_pk']}"),
            _hkdf(m, "meshbay:playlists:v2"))


def test_the_node_and_playlist_keys_match_across_languages(chain):
    vectors, js = chain
    for v, got in zip(vectors, js):
        k_node, k_pl = _reference(v)
        assert got["node"] == _fp(k_node), f"node key disagrees for {v['username']!r}"
        assert got["playlists"] == _fp(k_pl), f"playlist key disagrees for {v['username']!r}"


def test_an_mbk3_bundle_opens_from_the_specification(chain):
    """Magic, pepper version, nonce, AES-GCM with the account and node as
    associated data — read back by code that shares nothing with the writer."""
    import base64

    from cryptography.hazmat.primitives.ciphers.aead import AESGCM
    vectors, js = chain
    for v, got in zip(vectors, js):
        raw = base64.b64decode(got["bundle"])
        assert raw[:4] == b"MBK3" and raw[4] == 1
        aad = f"meshbay:bundle:v3|{v['user_id']}|{v['node_pk']}".encode()
        plain = json.loads(AESGCM(_reference(v)[0]).decrypt(raw[5:17], raw[17:], aad))
        assert base64.b64decode(plain["skEd"]) == b"ed-private"
        assert base64.b64decode(plain["skX"]) == b"x-private"