"""WebSocket transport with Noise_XX encryption + identity binding.

Wire protocol (per Noise message):
  <peer_id>\\x00<ed25519_pub_hex>\\x00<noise_message>

After the 3-message Noise handshake, both sides exchange an identity
binding proof: Ed25519.sign(handshake_hash) encrypted under the session.
This cryptographically binds the Ed25519 identity (PeerID) to the Noise
X25519 session.
"""

from __future__ import annotations

import asyncio
import hashlib
import logging
import struct
from collections.abc import AsyncIterator
from urllib.parse import urlparse

import websockets
import websockets.asyncio.client
import websockets.asyncio.server

from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey

from p2p.crypto import (
    NoiseInitiator,
    NoiseResponder,
    NoiseSession,
    derive_x25519_from_ed25519,
)
from p2p.errors import CryptoError, PeerUnreachable
from p2p.identity import PeerIdentity, _b58encode
from p2p.transport.base import Connection, Transport

logger = logging.getLogger(__name__)


class WSConnection(Connection):
    def __init__(self, ws, session: NoiseSession, peer_id: str, addr: str) -> None:
        self._ws = ws
        self._session = session
        self._peer_id = peer_id
        self._addr = addr
        self._closed = False

    @property
    def peer_id(self) -> str:
        return self._peer_id

    @property
    def remote_addr(self) -> str:
        return self._addr

    @property
    def is_open(self) -> bool:
        return not self._closed and self._ws.close_code is None

    async def send(self, data: bytes) -> None:
        if not self.is_open:
            raise ConnectionError("closed")
        ct = self._session.encrypt(data)
        await self._ws.send(struct.pack(">I", len(ct)) + ct)

    async def recv(self) -> bytes:
        if not self.is_open:
            raise ConnectionError("closed")
        try:
            raw = await self._ws.recv()
            if isinstance(raw, str):
                raw = raw.encode()
            if len(raw) < 4:
                raise CryptoError("frame too short")
            length = struct.unpack(">I", raw[:4])[0]
            return self._session.decrypt(raw[4:4 + length])
        except websockets.ConnectionClosed:
            self._closed = True
            raise ConnectionError("ws closed")

    async def close(self) -> None:
        self._closed = True
        try:
            await self._ws.close()
        except Exception:
            pass


# ── Helpers ─────────────────────────────────────────────────────────


def _x25519_key(identity: PeerIdentity):
    from cryptography.hazmat.primitives.serialization import Encoding, NoEncryption, PrivateFormat
    raw = identity.private_key.private_bytes(Encoding.Raw, PrivateFormat.Raw, NoEncryption())
    return derive_x25519_from_ed25519(raw)


def _verify_peer_id_binding(claimed_id: str, ed_pub_hex: str) -> None:
    """Verify PeerID is correctly derived from the Ed25519 public key."""
    ed_pub_bytes = bytes.fromhex(ed_pub_hex)
    if len(ed_pub_bytes) != 32:
        raise CryptoError("Ed25519 public key must be 32 bytes")
    expected = f"p2p:{_b58encode(hashlib.sha256(ed_pub_bytes).digest()[:20])}"
    if claimed_id != expected:
        raise CryptoError(f"PeerID mismatch: claimed {claimed_id}, derived {expected}")


def _verify_binding_proof(
    session: NoiseSession, peer_id: str, ed_pub_hex: str, sig: bytes,
) -> None:
    """Verify identity binding: Ed25519 signature over Noise handshake hash."""
    ed_pub = Ed25519PublicKey.from_public_bytes(bytes.fromhex(ed_pub_hex))
    try:
        ed_pub.verify(sig, session.handshake_hash)
    except Exception as e:
        raise CryptoError(f"Identity binding proof failed for {peer_id}: {e}") from e


def _pack_header(identity: PeerIdentity) -> bytes:
    """Build the wire header: peer_id + ed25519_pub_hex."""
    return (
        identity.peer_id.encode()
        + b"\x00"
        + identity.public_key_raw().hex().encode()
    )


def _parse_header(raw: bytes) -> tuple[str, str, bytes]:
    """Parse wire header → (peer_id, ed25519_pub_hex, noise_payload)."""
    parts = raw.split(b"\x00", 2)
    if len(parts) != 3:
        raise CryptoError("Bad wire format: expected <id>\\x00<ed_pub>\\x00<noise_msg>")
    return parts[0].decode(), parts[1].decode(), parts[2]


# ── Noise + binding proof ───────────────────────────────────────────


async def _noise_initiator(ws, identity: PeerIdentity):
    init = NoiseInitiator(_x25519_key(identity))

    # Msg 1: → e (with identity header)
    msg1 = init.write_message_1()
    await ws.send(_pack_header(identity) + b"\x00" + msg1)

    # Msg 2: ← e, ee, s, es (with remote identity header)
    raw2 = await ws.recv()
    if isinstance(raw2, str):
        raw2 = raw2.encode()
    remote_id, remote_ed_pub_hex, noise_msg2 = _parse_header(raw2)

    msg3 = init.read_message_2(noise_msg2)

    # Msg 3: → s, se
    await ws.send(msg3)

    session = init.finalize()

    # Verify PeerID ↔ Ed25519 binding
    _verify_peer_id_binding(remote_id, remote_ed_pub_hex)

    # Identity binding proof exchange (encrypted under session)
    # Initiator sends first, responder second
    our_proof = identity.sign(session.handshake_hash)
    await ws.send(session.encrypt(our_proof))

    raw_proof = await ws.recv()
    if isinstance(raw_proof, str):
        raw_proof = raw_proof.encode()
    remote_proof = session.decrypt(raw_proof)
    _verify_binding_proof(session, remote_id, remote_ed_pub_hex, remote_proof)

    logger.debug("Identity binding verified for %s (initiator)", remote_id)
    return session, remote_id


async def _noise_responder(ws, identity: PeerIdentity):
    # Msg 1: ← e (with remote identity header)
    raw1 = await ws.recv()
    if isinstance(raw1, str):
        raw1 = raw1.encode()
    remote_id, remote_ed_pub_hex, noise_msg1 = _parse_header(raw1)

    resp = NoiseResponder(_x25519_key(identity))

    # Msg 2: → e, ee, s, es (with our identity header)
    msg2 = resp.read_message_1(noise_msg1)
    await ws.send(_pack_header(identity) + b"\x00" + msg2)

    # Msg 3: ← s, se
    raw3 = await ws.recv()
    if isinstance(raw3, str):
        raw3 = raw3.encode()
    resp.read_message_3(raw3)

    session = resp.finalize()

    # Verify PeerID ↔ Ed25519 binding
    _verify_peer_id_binding(remote_id, remote_ed_pub_hex)

    # Identity binding proof exchange (encrypted under session)
    # Receive initiator's proof first, then send ours
    raw_proof = await ws.recv()
    if isinstance(raw_proof, str):
        raw_proof = raw_proof.encode()
    remote_proof = session.decrypt(raw_proof)
    _verify_binding_proof(session, remote_id, remote_ed_pub_hex, remote_proof)

    our_proof = identity.sign(session.handshake_hash)
    await ws.send(session.encrypt(our_proof))

    logger.debug("Identity binding verified for %s (responder)", remote_id)
    return session, remote_id


# ── Raw Noise handshake (works over any send/recv interface) ────────


async def _noise_initiator_raw(conn, identity: PeerIdentity):
    """Noise_XX + identity binding over any connection with send()/recv()."""
    init = NoiseInitiator(_x25519_key(identity))
    msg1 = init.write_message_1()
    await conn.send(_pack_header(identity) + b"\x00" + msg1)

    raw2 = await conn.recv()
    if isinstance(raw2, str):
        raw2 = raw2.encode()
    remote_id, remote_ed_pub_hex, noise_msg2 = _parse_header(raw2)
    msg3 = init.read_message_2(noise_msg2)
    await conn.send(msg3)
    session = init.finalize()
    _verify_peer_id_binding(remote_id, remote_ed_pub_hex)

    our_proof = identity.sign(session.handshake_hash)
    await conn.send(session.encrypt(our_proof))
    raw_proof = await conn.recv()
    if isinstance(raw_proof, str):
        raw_proof = raw_proof.encode()
    _verify_binding_proof(session, remote_id, remote_ed_pub_hex, session.decrypt(raw_proof))

    return session, remote_id


async def _noise_responder_raw(conn, identity: PeerIdentity):
    """Noise_XX + identity binding over any connection with send()/recv()."""
    raw1 = await conn.recv()
    if isinstance(raw1, str):
        raw1 = raw1.encode()
    remote_id, remote_ed_pub_hex, noise_msg1 = _parse_header(raw1)

    resp = NoiseResponder(_x25519_key(identity))
    msg2 = resp.read_message_1(noise_msg1)
    await conn.send(_pack_header(identity) + b"\x00" + msg2)

    raw3 = await conn.recv()
    if isinstance(raw3, str):
        raw3 = raw3.encode()
    resp.read_message_3(raw3)
    session = resp.finalize()
    _verify_peer_id_binding(remote_id, remote_ed_pub_hex)

    raw_proof = await conn.recv()
    if isinstance(raw_proof, str):
        raw_proof = raw_proof.encode()
    _verify_binding_proof(session, remote_id, remote_ed_pub_hex, session.decrypt(raw_proof))

    our_proof = identity.sign(session.handshake_hash)
    await conn.send(session.encrypt(our_proof))

    return session, remote_id


class EncryptedRelayConnection:
    """Wraps a RelayConnection with Noise session encryption.

    After the Noise handshake completes over the relay, this class
    encrypts/decrypts all subsequent messages using the session keys.
    Implements the same interface as WSConnection.
    """

    def __init__(self, relay_conn, session: NoiseSession, peer_id: str) -> None:
        self._relay = relay_conn
        self._session = session
        self._peer_id = peer_id

    @property
    def peer_id(self) -> str:
        return self._peer_id

    @property
    def remote_addr(self) -> str:
        return self._relay.remote_addr

    @property
    def is_open(self) -> bool:
        return self._relay.is_open

    async def send(self, data: bytes) -> None:
        import struct
        ct = self._session.encrypt(data)
        await self._relay.send(struct.pack(">I", len(ct)) + ct)

    async def recv(self) -> bytes:
        import struct
        raw = await self._relay.recv()
        if isinstance(raw, str):
            raw = raw.encode()
        if len(raw) < 4:
            raise CryptoError("frame too short")
        length = struct.unpack(">I", raw[:4])[0]
        ct = raw[4:4 + length]
        if len(ct) != length:
            raise CryptoError(
                f"Frame length mismatch: header={length}, got={len(ct)}, total={len(raw)}"
            )
        try:
            return self._session.decrypt(ct)
        except CryptoError:
            logger.debug(
                "Decrypt failed: raw_len=%d, frame_len=%d, nonce=%d",
                len(raw), length, self._session._recv.n - 1,
            )
            raise

    async def close(self) -> None:
        await self._relay.close()


# ── Transport ───────────────────────────────────────────────────────


class WSTransport(Transport):
    def __init__(self, identity: PeerIdentity) -> None:
        self._identity = identity
        self._server = None
        self._conns: list[WSConnection] = []

    async def listen(self, host: str, port: int) -> AsyncIterator[Connection]:
        q: asyncio.Queue[WSConnection] = asyncio.Queue()

        async def handler(ws):
            addr = f"{ws.remote_address[0]}:{ws.remote_address[1]}" if ws.remote_address else "?"
            try:
                session, rid = await asyncio.wait_for(
                    _noise_responder(ws, self._identity), timeout=30,
                )
                conn = WSConnection(ws, session, rid, addr)
                self._conns.append(conn)
                await q.put(conn)
                while conn.is_open:
                    await asyncio.sleep(1)
            except Exception:
                logger.debug("Handshake failed from %s", addr, exc_info=True)

        self._server = await websockets.asyncio.server.serve(
            handler, host, port, max_size=10 * 1024 * 1024,
        )
        logger.info("P2P WS server on %s:%d", host, port)
        try:
            while True:
                yield await q.get()
        except asyncio.CancelledError:
            pass

    async def connect(self, addr: str) -> Connection:
        url = addr if urlparse(addr).scheme in ("ws", "wss") else f"ws://{addr}"
        try:
            ws = await websockets.asyncio.client.connect(url, max_size=10*1024*1024, open_timeout=10)
        except Exception as e:
            raise PeerUnreachable("?", str(e)) from e
        try:
            session, rid = await asyncio.wait_for(
                _noise_initiator(ws, self._identity), timeout=30,
            )
        except Exception as e:
            await ws.close()
            raise CryptoError(f"Noise failed: {e}") from e
        p = urlparse(url)
        conn = WSConnection(ws, session, rid, f"{p.hostname}:{p.port}")
        self._conns.append(conn)
        return conn

    async def shutdown(self) -> None:
        for c in self._conns:
            await c.close()
        self._conns.clear()
        if self._server:
            self._server.close()
            await self._server.wait_closed()
