"""
One account, several nodes: one linked key per machine the person runs a node on.

The hub used to keep one node key per account, so linking a second machine
unlinked the first, which then failed its next sign-in with nothing said. These
pin the rules that replaced it: every linked key signs in, unlinking one cuts
that machine off at once and no other, a node token can neither link nor unlink,
and no other account learns which keys an account has.
"""

import base64
import json
import time
from pathlib import Path

import pytest
from alembic import command
from alembic.config import Config
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
from sqlalchemy import create_engine, text

HUB = Path(__file__).resolve().parents[1]


def _keypair():
    sk = Ed25519PrivateKey.generate()
    pk = sk.public_key().public_bytes(serialization.Encoding.Raw,
                                      serialization.PublicFormat.Raw)
    return sk, base64.b64encode(pk).decode()


def _h(token):
    return {"Authorization": f"Bearer {token}"}


async def _account(client, username):
    _, pk_ed = _keypair()
    r = await client.post("/v1/users/register", json={
        "username": username, "email": f"{username}@test.local",
        "password": "testpass99", "pk_user_ed25519": pk_ed, "pk_user_x25519": pk_ed,
    })
    assert r.status_code == 201, r.text
    r = await client.post("/v1/users/login", json={
        "username": username, "password": "testpass99"})
    assert r.status_code == 200, r.text
    return r.json()["access_token"]


async def _link(client, token, pk):
    return await client.put("/v1/users/me/node_key", json={"pk_node_ed25519": pk},
                            headers=_h(token))


async def _unlink(client, token, pk):
    return await client.request("DELETE", "/v1/users/me/node_key",
                                json={"pk_node_ed25519": pk}, headers=_h(token))


async def _node_auth(client, username, sk):
    ts = int(time.time())
    sig = sk.sign(f"meshbay:node_auth:{username}:{ts}".encode())
    return await client.post("/v1/nodes/auth", json={
        "username": username, "timestamp": ts,
        "signature": base64.b64encode(sig).decode()})


async def _announce(client, node_token, user_id, sk, pk):
    ts = int(time.time())
    sig = sk.sign(f"meshbay:node_announce:{user_id}:{pk}:{ts}".encode())
    r = await client.post("/v1/nodes/announce", json={
        "pk_node": pk, "timestamp": ts, "signature": base64.b64encode(sig).decode()},
        headers=_h(node_token))
    assert r.status_code == 201, r.text
    return r.json()["node_id"]


def _claims(token):
    body = token.split(".")[1]
    return json.loads(base64.urlsafe_b64decode(body + "=" * (-len(body) % 4)))


async def _two_nodes(client, username):
    token = await _account(client, username)
    sk_a, pk_a = _keypair()
    sk_b, pk_b = _keypair()
    assert (await _link(client, token, pk_a)).status_code == 200
    assert (await _link(client, token, pk_b)).status_code == 200
    return token, (sk_a, pk_a), (sk_b, pk_b)


class _FakeSocket:
    def __init__(self):
        self.sent, self.closed = [], None

    async def send_text(self, data):
        self.sent.append(json.loads(data))

    async def close(self, code=1000):
        self.closed = code


# ── Two machines, one account ────────────────────────────────────────────────

@pytest.mark.asyncio
async def test_two_nodes_of_one_account_both_sign_in(client):
    _, (sk_a, pk_a), (sk_b, pk_b) = await _two_nodes(client, "twonodes")

    ra = await _node_auth(client, "twonodes", sk_a)
    rb = await _node_auth(client, "twonodes", sk_b)
    assert ra.status_code == 200, ra.text
    assert rb.status_code == 200, rb.text
    # Each token names the key that signed in, which is what unlinking cuts.
    assert _claims(ra.json()["access_token"])["nk"] == pk_a
    assert _claims(rb.json()["access_token"])["nk"] == pk_b


@pytest.mark.asyncio
async def test_linking_a_second_key_keeps_the_first(client):
    token, (_, pk_a), (_, pk_b) = await _two_nodes(client, "keepsfirst")
    r = await client.get("/v1/users/me/node_keys", headers=_h(token))
    assert r.status_code == 200
    assert [k["pk_node_ed25519"] for k in r.json()["keys"]] == [pk_a, pk_b]


@pytest.mark.asyncio
async def test_an_unlinked_key_is_told_to_link_not_to_register(client):
    """The daemon reads "No node key" to wait for its link (`waiting_for_node_key`)."""
    token = await _account(client, "notlinked")
    _, pk_a = _keypair()
    await _link(client, token, pk_a)
    sk_other, _ = _keypair()
    r = await _node_auth(client, "notlinked", sk_other)
    assert r.status_code == 401
    assert "No node key" in r.json()["detail"]


@pytest.mark.asyncio
async def test_linking_the_same_key_twice_changes_nothing(client):
    token = await _account(client, "twice_test")
    _, pk = _keypair()
    assert (await _link(client, token, pk)).status_code == 200
    assert (await _link(client, token, pk)).status_code == 200
    r = await client.get("/v1/users/me/node_keys", headers=_h(token))
    assert len(r.json()["keys"]) == 1


@pytest.mark.asyncio
async def test_a_key_is_stored_in_the_spelling_nodes_announce(client):
    """Pasted with spaces around it, it must still match the node's own announce."""
    token = await _account(client, "spelling")
    _, pk = _keypair()
    r = await _link(client, token, f"  {pk}\n")
    assert r.status_code == 200
    assert r.json()["pk_node_ed25519"] == pk


# ── Unlinking one machine ────────────────────────────────────────────────────

@pytest.mark.asyncio
async def test_unlinking_one_node_cuts_it_off_and_leaves_the_other(client):
    token, (sk_a, pk_a), (sk_b, pk_b) = await _two_nodes(client, "unlinkone")
    tok_a = (await _node_auth(client, "unlinkone", sk_a)).json()["access_token"]
    tok_b = (await _node_auth(client, "unlinkone", sk_b)).json()["access_token"]

    r = await _unlink(client, token, pk_a)
    assert r.status_code == 200, r.text

    # A's token stops working now, not in an hour…
    r = await client.get("/v1/users/unlinkone/pubkeys", headers=_h(tok_a))
    assert r.status_code == 401
    # …and A cannot sign in again until it is linked again.
    r = await _node_auth(client, "unlinkone", sk_a)
    assert r.status_code == 401
    assert "No node key" in r.json()["detail"]

    # B is untouched.
    r = await client.get("/v1/users/unlinkone/pubkeys", headers=_h(tok_b))
    assert r.status_code == 200
    assert (await _node_auth(client, "unlinkone", sk_b)).status_code == 200

    # Linking A again lets it back in.
    assert (await _link(client, token, pk_a)).status_code == 200
    assert (await _node_auth(client, "unlinkone", sk_a)).status_code == 200


@pytest.mark.asyncio
async def test_unlinking_closes_that_nodes_socket_only(client):
    from meshbay_hub.api import revocation

    token, (sk_a, pk_a), (sk_b, pk_b) = await _two_nodes(client, "sockets_test")
    me = (await client.get("/v1/users/sockets_test/pubkeys", headers=_h(token))).json()["user_id"]
    tok_a = (await _node_auth(client, "sockets_test", sk_a)).json()["access_token"]
    tok_b = (await _node_auth(client, "sockets_test", sk_b)).json()["access_token"]
    id_a = await _announce(client, tok_a, me, sk_a, pk_a)
    id_b = await _announce(client, tok_b, me, sk_b, pk_b)

    ws_a, ws_b = _FakeSocket(), _FakeSocket()
    revocation._connected_nodes[id_a] = ws_a
    revocation._connected_nodes[id_b] = ws_b
    try:
        listed = (await client.get("/v1/users/me/node_keys", headers=_h(token))).json()
        assert [k["connected"] for k in listed["keys"]] == [True, True]
        assert all(k["last_seen"] for k in listed["keys"])

        assert (await _unlink(client, token, pk_a)).status_code == 200
        assert ws_a.closed == 4003
        assert ws_b.closed is None
        assert id_a not in revocation._connected_nodes
        assert revocation._connected_nodes.get(id_b) is ws_b
    finally:
        revocation.forget_node(id_a)
        revocation.forget_node(id_b)


@pytest.mark.asyncio
async def test_a_node_token_registers_its_own_node_only(client):
    """A token names its key: it cannot register the socket of the account's other machine."""
    from meshbay_hub.api.revocation import _authorize_node_ws

    token, (sk_a, pk_a), (sk_b, pk_b) = await _two_nodes(client, "ownnode_test")
    me = (await client.get("/v1/users/ownnode_test/pubkeys", headers=_h(token))).json()["user_id"]
    tok_a = (await _node_auth(client, "ownnode_test", sk_a)).json()["access_token"]
    tok_b = (await _node_auth(client, "ownnode_test", sk_b)).json()["access_token"]
    id_a = await _announce(client, tok_a, me, sk_a, pk_a)
    id_b = await _announce(client, tok_b, me, sk_b, pk_b)

    assert (await _authorize_node_ws(tok_a, id_a, []))[0] == id_a
    assert (await _authorize_node_ws(tok_b, id_b, []))[0] == id_b
    assert (await _authorize_node_ws(tok_a, id_b, []))[0] is None

    await _unlink(client, token, pk_a)
    refused, why = await _authorize_node_ws(tok_a, id_a, [])
    assert refused is None and "no longer linked" in why
    assert (await _authorize_node_ws(tok_b, id_b, []))[0] == id_b


@pytest.mark.asyncio
async def test_a_node_token_from_before_the_claim_still_works(client):
    """A token without `nk` predates it; it expires within the hour on its own."""
    from meshbay_hub.auth import issue_access_token

    token = await _account(client, "legacy_test")
    me = (await client.get("/v1/users/legacy_test/pubkeys", headers=_h(token))).json()["user_id"]
    old = issue_access_token(me, scope="node")
    r = await client.get("/v1/users/legacy_test/pubkeys", headers=_h(old))
    assert r.status_code == 200


@pytest.mark.asyncio
async def test_a_token_from_before_the_claim_cannot_bring_an_unlinked_node_back(client):
    """Found in the real run: the node held its pre-upgrade token and re-registered at once."""
    from meshbay_hub.api.revocation import _authorize_node_ws
    from meshbay_hub.auth import issue_access_token

    token, (sk_a, pk_a), (sk_b, pk_b) = await _two_nodes(client, "legacyws_test")
    me = (await client.get("/v1/users/legacyws_test/pubkeys",
                           headers=_h(token))).json()["user_id"]
    tok_a = (await _node_auth(client, "legacyws_test", sk_a)).json()["access_token"]
    tok_b = (await _node_auth(client, "legacyws_test", sk_b)).json()["access_token"]
    id_a = await _announce(client, tok_a, me, sk_a, pk_a)
    id_b = await _announce(client, tok_b, me, sk_b, pk_b)
    old = issue_access_token(me, scope="node")

    assert (await _authorize_node_ws(old, id_a, []))[0] == id_a
    await _unlink(client, token, pk_a)
    refused, why = await _authorize_node_ws(old, id_a, [])
    assert refused is None and "no longer linked" in why
    assert (await _authorize_node_ws(old, id_b, []))[0] == id_b


# ── Who may link, unlink and see ─────────────────────────────────────────────

@pytest.mark.asyncio
async def test_a_node_token_cannot_link_list_or_unlink_a_sibling(client):
    token, (sk_a, pk_a), (sk_b, pk_b) = await _two_nodes(client, "nodescope")
    tok_a = (await _node_auth(client, "nodescope", sk_a)).json()["access_token"]
    _, pk_evil = _keypair()
    assert (await _link(client, tok_a, pk_evil)).status_code == 403
    assert (await _unlink(client, tok_a, pk_b)).status_code == 403
    r = await client.get("/v1/users/me/node_keys", headers=_h(tok_a))
    assert r.status_code == 403
    keys = (await client.get("/v1/users/me/node_keys", headers=_h(token))).json()["keys"]
    assert [k["pk_node_ed25519"] for k in keys] == [pk_a, pk_b]


@pytest.mark.asyncio
async def test_a_node_may_unlink_itself(client):
    """`meshbay-node setup --reset` and the Node page: the node lets go of its own key."""
    token, (sk_a, pk_a), (sk_b, pk_b) = await _two_nodes(client, "selfunlink")
    tok_a = (await _node_auth(client, "selfunlink", sk_a)).json()["access_token"]
    tok_b = (await _node_auth(client, "selfunlink", sk_b)).json()["access_token"]

    assert (await _unlink(client, tok_a, pk_a)).status_code == 200
    assert (await _node_auth(client, "selfunlink", sk_a)).status_code == 401

    # A node from before the key was named in the request asks with no body.
    r = await client.request("DELETE", "/v1/users/me/node_key", headers=_h(tok_b))
    assert r.status_code == 200, r.text
    assert r.json()["pk_node_ed25519"] == pk_b
    keys = (await client.get("/v1/users/me/node_keys", headers=_h(token))).json()["keys"]
    assert keys == []


@pytest.mark.asyncio
async def test_a_node_token_without_its_key_cannot_unlink(client):
    from meshbay_hub.auth import issue_access_token

    token, (_, pk_a), _ = await _two_nodes(client, "legacyunl")
    me = (await client.get("/v1/users/legacyunl/pubkeys", headers=_h(token))).json()["user_id"]
    old = issue_access_token(me, scope="node")
    assert (await _unlink(client, old, pk_a)).status_code == 403
    r = await client.request("DELETE", "/v1/users/me/node_key", headers=_h(old))
    assert r.status_code == 403


@pytest.mark.asyncio
async def test_another_account_cannot_unlink_or_see_my_keys(client):
    token, (sk_a, pk_a), _ = await _two_nodes(client, "victim_test")
    mallory = await _account(client, "mallory_test")

    r = await _unlink(client, mallory, pk_a)
    assert r.status_code == 404
    assert (await _node_auth(client, "victim_test", sk_a)).status_code == 200

    r = await client.get("/v1/users/victim_test/pubkeys", headers=_h(mallory))
    assert r.status_code == 200
    assert "pk_node_ed25519" not in r.json()
    assert "keys" not in r.json()
    mine = (await client.get("/v1/users/me/node_keys", headers=_h(mallory))).json()
    assert mine["keys"] == []


@pytest.mark.asyncio
async def test_installed_desktop_apps_still_see_their_own_latest_key(client):
    """Their interface shows the Node entries from `/pubkeys`; only for yourself."""
    token, (_, pk_a), (_, pk_b) = await _two_nodes(client, "compatapp_test")
    r = await client.get("/v1/users/compatapp_test/pubkeys", headers=_h(token))
    assert r.json()["pk_node_ed25519"] == pk_b
    await _unlink(client, token, pk_b)
    r = await client.get("/v1/users/compatapp_test/pubkeys", headers=_h(token))
    assert r.json()["pk_node_ed25519"] == pk_a
    await _unlink(client, token, pk_a)
    r = await client.get("/v1/users/compatapp_test/pubkeys", headers=_h(token))
    assert "pk_node_ed25519" not in r.json()


@pytest.mark.asyncio
async def test_the_same_key_on_another_account_signs_in_only_as_that_account(client):
    """Linking someone's public key gains nothing: the signature names the username."""
    alice = await _account(client, "keyowner")
    bob = await _account(client, "keycopier")
    sk, pk = _keypair()
    await _link(client, alice, pk)
    await _link(client, bob, pk)
    await _unlink(client, bob, pk)
    assert (await _node_auth(client, "keyowner", sk)).status_code == 200


@pytest.mark.asyncio
async def test_unlinking_a_key_that_is_not_linked_says_so(client):
    token = await _account(client, "nothing_test")
    _, pk = _keypair()
    assert (await _unlink(client, token, pk)).status_code == 404
    r = await client.request("DELETE", "/v1/users/me/node_key", headers=_h(token))
    assert r.status_code == 422


@pytest.mark.asyncio
async def test_the_limit_is_the_announce_limit(client):
    from meshbay_hub.api.nodes import MAX_NODES_PER_ACCOUNT

    token = await _account(client, "manynodes")
    keys = [_keypair()[1] for _ in range(MAX_NODES_PER_ACCOUNT)]
    for pk in keys:
        assert (await _link(client, token, pk)).status_code == 200
    r = await _link(client, token, _keypair()[1])
    assert r.status_code == 409
    assert "Unlink one" in r.json()["detail"]
    # Linking one already there still answers, at the limit.
    assert (await _link(client, token, keys[0])).status_code == 200
    listed = (await client.get("/v1/users/me/node_keys", headers=_h(token))).json()
    assert listed["limit"] == MAX_NODES_PER_ACCOUNT
    assert len(listed["keys"]) == MAX_NODES_PER_ACCOUNT


@pytest.mark.asyncio
async def test_deleting_the_account_unlinks_every_key(client, db_session):
    import hashlib

    from meshbay_hub.db.models import NodeKey
    from sqlalchemy import select

    name, passphrase = "goneacct_test", "a-long-enough-passphrase"
    salt = hashlib.sha256(f"meshbay:auth:v1:{name}".encode()).digest()
    auth_key = base64.b64encode(
        hashlib.pbkdf2_hmac("sha512", passphrase.encode(), salt, 600_000, 32)).decode()
    r = await client.post("/v1/users/register", json={
        "username": name, "email": f"{name}@example.com", "auth_key": auth_key})
    assert r.status_code in (200, 201), r.text
    token = (await client.post("/v1/users/login", json={
        "username": name, "auth_key": auth_key})).json()["access_token"]
    sk_a, pk_a = _keypair()
    await _link(client, token, pk_a)
    await _link(client, token, _keypair()[1])

    r = await client.request("DELETE", "/v1/users/me", headers=_h(token),
                             json={"auth_key": auth_key})
    assert r.status_code == 200, r.text
    left = (await db_session.execute(select(NodeKey))).scalars().all()
    assert left == []
    assert (await _node_auth(client, name, sk_a)).status_code == 401


# ── The upgrade ──────────────────────────────────────────────────────────────

def test_the_upgrade_keeps_the_key_each_account_had_linked(tmp_path, monkeypatch):
    db = tmp_path / "hub.db"
    monkeypatch.setenv("MESHBAY_DATABASE_URL", f"sqlite+aiosqlite:///{db}")
    cfg = Config(str(HUB / "alembic.ini"))
    command.upgrade(cfg, "e7f8a9b0c1d2")

    engine = create_engine(f"sqlite:///{db}")
    with engine.begin() as conn:
        for uid, pk in (("u-linked", "PKLINKED"), ("u-none", None)):
            conn.execute(text(
                "INSERT INTO users (id, username, email, pw_hash, pw_salt, hub_id, "
                "pk_node_ed25519) VALUES (:id, :id, '', x'00', x'00', 'h', :pk)"),
                {"id": uid, "pk": pk})

    command.upgrade(cfg, "f8a9b0c1d2e3")
    with engine.begin() as conn:
        rows = conn.execute(text(
            "SELECT user_id, pk_node_ed25519 FROM node_keys")).fetchall()
    assert rows == [("u-linked", "PKLINKED")]

    command.downgrade(cfg, "e7f8a9b0c1d2")
    with engine.begin() as conn:
        back = dict(conn.execute(text("SELECT id, pk_node_ed25519 FROM users")).fetchall())
    assert back == {"u-linked": "PKLINKED", "u-none": None}