import asyncio
import base64
import hashlib
import json
import time
from typing import cast

import pytest
from coincurve import PrivateKey, PublicKey
from Cryptodome import Random
from Cryptodome.Cipher import AES
from Cryptodome.Util.Padding import pad, unpad
from websockets import ServerConnection
from websockets import serve as ws_serve

from lnbits.wallets.nwc import NWCConnection, NWCWallet
from tests.wallets.helpers import (
    WalletTest,
    build_test_id,
    check_assertions,
    load_funding_source,
    wallet_fixtures_from_json,
)


def encrypt_content(priv_key, dest_pub_key, content):
    p = PublicKey(bytes.fromhex("02" + dest_pub_key))
    shared = p.multiply(bytes.fromhex(priv_key)).format()[1:]
    iv = Random.new().read(AES.block_size)
    aes = AES.new(shared, AES.MODE_CBC, iv)

    content_bytes = content.encode("utf-8")
    content_bytes = pad(content_bytes, AES.block_size)

    encrypted_b64 = base64.b64encode(aes.encrypt(content_bytes)).decode("ascii")
    iv_b64 = base64.b64encode(iv).decode("ascii")
    encrypted_content = encrypted_b64 + "?iv=" + iv_b64
    return encrypted_content


def decrypt_content(priv_key, source_pub_key, content):
    p = PublicKey(bytes.fromhex("02" + source_pub_key))
    shared = p.multiply(bytes.fromhex(priv_key)).format()[1:]
    encrypted_content_b64, iv_b64 = content.split("?iv=")
    encrypted_content = base64.b64decode(encrypted_content_b64.encode("ascii"))
    iv = base64.b64decode(iv_b64.encode("ascii"))
    aes = AES.new(shared, AES.MODE_CBC, iv)
    decrypted_bytes = aes.decrypt(encrypted_content)
    decrypted_bytes = unpad(decrypted_bytes, AES.block_size)
    return decrypted_bytes.decode("utf-8")


def json_dumps(data):
    if isinstance(data, dict):
        data = {k: v for k, v in data.items() if v is not None}
    return json.dumps(data, separators=(",", ":"), ensure_ascii=False)


def sign_event(pub_key, priv_key, event):
    signature_data = json_dumps(
        [
            0,
            pub_key,
            event["created_at"],
            event["kind"],
            event["tags"],
            event["content"],
        ]
    )
    event_id = hashlib.sha256(signature_data.encode()).hexdigest()
    event["id"] = event_id
    event["pubkey"] = pub_key
    s = PrivateKey(bytes.fromhex(priv_key))
    signature = s.sign_schnorr(bytes.fromhex(event_id)).hex()
    event["sig"] = signature
    return event


async def handle(  # noqa: C901
    wallet, mock_settings, data, websocket: ServerConnection
):
    async for message in websocket:
        if not wallet:
            continue
        msg = json.loads(message)
        if msg[0] == "REQ":
            sub_id = msg[1]
            sub_filter = msg[2]
            kinds = sub_filter["kinds"]
            if 13194 in kinds:  # Send info event
                event = {
                    "kind": 13194,
                    "content": " ".join(mock_settings["supported_methods"]),
                    "created_at": int(time.time()),
                    "tags": [],
                }
                sign_event(
                    mock_settings["service_public_key"],
                    mock_settings["service_private_key"],
                    event,
                )
                await websocket.send(json.dumps(["EVENT", sub_id, event]))
            elif 23195 in kinds:
                assert sub_filter["authors"] == [mock_settings["service_public_key"]]
        elif msg[0] == "EVENT":
            event = msg[1]
            decrypted_content = decrypt_content(
                mock_settings["service_private_key"],
                mock_settings["user_public_key"],
                event["content"],
            )
            content = json.loads(decrypted_content)
            mock = None
            for m in data.mocks:
                rb = m.request_body
                if rb and rb["method"] == content["method"]:
                    p1 = rb["params"]
                    p2 = content["params"]
                    p1 = json_dumps({k: v for k, v in p1.items() if v is not None})
                    p2 = json_dumps({k: v for k, v in p2.items() if v is not None})
                    if p1 == p2:
                        mock = m
                        break
            if mock:
                sub_id = None
                nwcwallet = cast(NWCWallet, wallet)
                for subscription in nwcwallet.conn.subscriptions.values():
                    if subscription["event_id"] == event["id"]:
                        sub_id = subscription["sub_id"]
                        break
                if sub_id:
                    response = mock.response
                    encrypted_content = encrypt_content(
                        mock_settings["service_private_key"],
                        mock_settings["user_public_key"],
                        json_dumps(response),
                    )
                    response_event = {
                        "kind": 23195,
                        "content": encrypted_content,
                        "created_at": int(time.time()),
                        "tags": [
                            ["e", event["id"]],
                            ["p", mock_settings["user_public_key"]],
                        ],
                    }
                    sign_event(
                        mock_settings["service_public_key"],
                        mock_settings["service_private_key"],
                        response_event,
                    )
                    await websocket.send(json.dumps(["EVENT", sub_id, response_event]))
            else:
                raise Exception(
                    "No mock found for "
                    + content["method"]
                    + " "
                    + json_dumps(content["params"])
                )


async def run(data: WalletTest):
    if data.skip:
        pytest.skip()

    wallet = None
    mock_settings = data.funding_source.mock_settings
    if mock_settings is None:
        return

    def handler(websocket):
        return handle(wallet, mock_settings, data, websocket)

    if mock_settings is not None:
        async with ws_serve(handler, "localhost", mock_settings["port"]) as server:
            await server.start_serving()
            wallet = load_funding_source(data.funding_source)
            await check_assertions(wallet, data)
            nwcwallet = cast(NWCWallet, wallet)
            await nwcwallet.cleanup()


@pytest.mark.anyio
async def test_nwc_rejects_event_from_unexpected_pubkey(mocker):
    async def _noop(*args, **kwargs):
        return None

    mocker.patch("lnbits.wallets.nwc.NWCConnection._connect_to_relay", new=_noop)
    mocker.patch("lnbits.wallets.nwc.NWCConnection._handle_timeouts", new=_noop)

    service_private_key = PrivateKey()
    service_public_key = service_private_key.public_key.format().hex()[2:]
    attacker_private_key = PrivateKey()
    attacker_public_key = attacker_private_key.public_key.format().hex()[2:]
    account_private_key = PrivateKey()

    conn = NWCConnection(
        service_public_key,
        account_private_key.secret.hex(),
        "ws://127.0.0.1:8555",
    )
    try:
        event = {
            "kind": 23195,
            "content": "{}",
            "created_at": int(time.time()),
            "tags": [["e", "request-event-id"]],
        }
        sign_event(attacker_public_key, attacker_private_key.secret.hex(), event)

        with pytest.raises(Exception, match="Invalid event signature"):
            await conn._on_event_message(["EVENT", "subid", event])
    finally:
        await conn.close()


@pytest.mark.anyio
async def test_nwc_marks_pending_invoice_settled_only_once():
    wallet = NWCWallet.__new__(NWCWallet)
    wallet.pending_invoice_details = {"checking-id": {"checking_id": "checking-id"}}
    wallet.pending_invoices = ["checking-id"]
    wallet.paid_invoices_queue = asyncio.Queue(0)

    wallet._mark_invoice_settled("checking-id", source="notification")
    wallet._mark_invoice_settled("checking-id", source="notification")

    assert wallet.paid_invoices_queue.qsize() == 1
    assert await wallet.paid_invoices_queue.get() == "checking-id"


@pytest.mark.anyio
async def test_nwc_registers_notification_subscriptions(mocker):
    async def _noop(*args, **kwargs):
        return None

    mocker.patch("lnbits.wallets.nwc.NWCConnection._connect_to_relay", new=_noop)
    mocker.patch("lnbits.wallets.nwc.NWCConnection._handle_timeouts", new=_noop)

    service_private_key = PrivateKey()
    service_public_key = service_private_key.public_key.format().hex()[2:]
    account_private_key = PrivateKey()

    conn = NWCConnection(
        service_public_key,
        account_private_key.secret.hex(),
        "ws://127.0.0.1:8555",
    )
    send_mock = mocker.patch.object(conn, "_send", mocker.AsyncMock())

    try:
        await conn._subscribe_to_notifications()

        assert len(conn.notification_subscription_ids) == 2
        assert len(conn.subscriptions) == 2
        assert set(conn.subscriptions.keys()) == conn.notification_subscription_ids
        assert all(
            subscription["method"] == "notification_sub"
            and subscription["event_id"] == subscription["sub_id"]
            for subscription in conn.subscriptions.values()
        )
        assert send_mock.await_count == 2
    finally:
        await conn.close()


@pytest.mark.anyio
async def test_nwc_spreads_fallback_lookups_with_cooldown(mocker):
    def _schedule_next_lookup(
        invoice: dict[str, object], now: float | None = None
    ) -> None:
        assert now is not None
        invoice["next_lookup_at"] = now + 1

    wallet = NWCWallet.__new__(NWCWallet)
    wallet.shutdown = False
    wallet.pending_invoices = ["checking-1", "checking-2"]
    wallet.pending_invoice_details = {
        "checking-1": {
            "checking_id": "checking-1",
            "next_lookup_at": 0.0,
            "lookup_attempts": 0,
        },
        "checking-2": {
            "checking_id": "checking-2",
            "next_lookup_at": 0.0,
            "lookup_attempts": 0,
        },
    }
    wallet.pending_invoices_lookup_cooldown = 1.0
    wallet._is_shutting_down = lambda: False
    wallet._payment_data_is_settled = lambda payment_data: False
    wallet._cache_payment_data = lambda *args, **kwargs: None
    wallet._schedule_next_lookup = _schedule_next_lookup
    wallet.conn = mocker.Mock()
    wallet.conn.get_info = mocker.AsyncMock()
    wallet.conn.supports_method = mocker.Mock(return_value=True)
    wallet.conn.call = mocker.AsyncMock(return_value={"settled_at": None})
    sleep_mock = mocker.patch("lnbits.wallets.nwc.asyncio.sleep", mocker.AsyncMock())

    await wallet._run_fallback_lookups(100.0)

    assert wallet.conn.call.await_count == 2
    sleep_mock.assert_awaited_once_with(1.0)


@pytest.mark.anyio
@pytest.mark.parametrize(
    "test_data",
    wallet_fixtures_from_json("tests/wallets/fixtures/json/fixtures_nwc.json"),
    ids=build_test_id,
)
async def test_nwc_wallet(test_data: WalletTest):
    await run(test_data)


@pytest.mark.anyio
@pytest.mark.parametrize(
    "test_data",
    wallet_fixtures_from_json("tests/wallets/fixtures/json/fixtures_nwc_bad.json"),
    ids=build_test_id,
)
async def test_nwc_wallet_bad(test_data: WalletTest):
    await run(test_data)