"""WebSocket transport with Noise_XX encryption."""

from __future__ import annotations

import asyncio
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 a2a_p2p.crypto import (
    NoiseInitiator,
    NoiseResponder,
    NoiseSession,
    generate_x25519_from_ed25519_seed,
)
from a2a_p2p.errors import CryptoError, PeerUnreachable
from a2a_p2p.identity import PeerIdentity
from a2a_p2p.transport.base import Connection, Transport

logger = logging.getLogger(__name__)


class WSConnection(Connection):
    """A Noise-encrypted WebSocket connection."""

    def __init__(
        self,
        ws: websockets.asyncio.client.ClientConnection
        | websockets.asyncio.server.ServerConnection,
        session: NoiseSession,
        peer_id: str,
        remote_addr: str,
    ) -> None:
        self._ws = ws
        self._session = session
        self._peer_id = peer_id
        self._remote_addr = remote_addr
        self._closed = False

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

    @property
    def remote_addr(self) -> str:
        return self._remote_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:
        """Encrypt and send data over the WebSocket."""
        if not self.is_open:
            raise ConnectionError("Connection closed")
        ct = self._session.encrypt(data)
        length_prefix = struct.pack(">I", len(ct))
        await self._ws.send(length_prefix + ct)

    async def recv(self) -> bytes:
        """Receive and decrypt data from the WebSocket."""
        if not self.is_open:
            raise ConnectionError("Connection closed")
        try:
            raw = await self._ws.recv()
            if isinstance(raw, str):
                raw = raw.encode("utf-8")
            if len(raw) < 4:
                raise CryptoError("Frame too short")
            length = struct.unpack(">I", raw[:4])[0]
            ct = raw[4:]
            if len(ct) != length:
                raise CryptoError(
                    f"Frame length mismatch: header says {length}, got {len(ct)}"
                )
            return self._session.decrypt(ct)
        except websockets.ConnectionClosed:
            self._closed = True
            raise ConnectionError("WebSocket connection closed")

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


def _get_x25519_private(identity: PeerIdentity):
    """Derive X25519 private key from Ed25519 identity for Noise handshake."""
    from cryptography.hazmat.primitives.serialization import Encoding, NoEncryption, PrivateFormat

    ed_bytes = identity.private_key.private_bytes(
        Encoding.Raw, PrivateFormat.Raw, NoEncryption(),
    )
    return generate_x25519_from_ed25519_seed(ed_bytes)


async def _do_noise_initiator(
    ws, identity: PeerIdentity,
) -> tuple[NoiseSession, str]:
    """Perform Noise_XX handshake as initiator over a WebSocket."""
    x25519_priv = _get_x25519_private(identity)
    initiator = NoiseInitiator(x25519_priv)

    # Message 1: → e
    msg1 = initiator.write_message_1()
    # Prepend our PeerID so responder knows who we claim to be
    await ws.send(identity.peer_id.encode("utf-8") + b"\x00" + msg1)

    # Message 2: ← e, ee, s, es + responder PeerID
    raw2 = await ws.recv()
    if isinstance(raw2, str):
        raw2 = raw2.encode("utf-8")
    sep = raw2.index(b"\x00")
    remote_peer_id = raw2[:sep].decode("utf-8")
    msg2 = raw2[sep + 1:]

    msg3 = initiator.read_message_2(msg2)

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

    session = initiator.finalize()
    return session, remote_peer_id


async def _do_noise_responder(
    ws, identity: PeerIdentity,
) -> tuple[NoiseSession, str]:
    """Perform Noise_XX handshake as responder over a WebSocket."""
    x25519_priv = _get_x25519_private(identity)
    responder = NoiseResponder(x25519_priv)

    # Message 1: ← e + initiator PeerID
    raw1 = await ws.recv()
    if isinstance(raw1, str):
        raw1 = raw1.encode("utf-8")
    sep = raw1.index(b"\x00")
    remote_peer_id = raw1[:sep].decode("utf-8")
    msg1 = raw1[sep + 1:]

    msg2 = responder.read_message_1(msg1)

    # Message 2: → e, ee, s, es + our PeerID
    await ws.send(identity.peer_id.encode("utf-8") + b"\x00" + msg2)

    # Message 3: ← s, se
    raw3 = await ws.recv()
    if isinstance(raw3, str):
        raw3 = raw3.encode("utf-8")
    responder.read_message_3(raw3)

    session = responder.finalize()
    return session, remote_peer_id


class WSTransport(Transport):
    """WebSocket transport with Noise_XX encryption."""

    def __init__(self, identity: PeerIdentity) -> None:
        self._identity = identity
        self._server: websockets.asyncio.server.Server | None = None
        self._connections: list[WSConnection] = []

    async def listen(self, host: str, port: int) -> AsyncIterator[Connection]:
        """Listen for incoming WebSocket connections.

        Each incoming connection goes through the Noise_XX handshake
        as responder before being yielded.
        """
        conn_queue: asyncio.Queue[WSConnection] = asyncio.Queue()

        async def handler(ws):
            remote = f"{ws.remote_address[0]}:{ws.remote_address[1]}" if ws.remote_address else "unknown"
            try:
                session, remote_peer_id = await asyncio.wait_for(
                    _do_noise_responder(ws, self._identity),
                    timeout=30.0,
                )
                conn = WSConnection(ws, session, remote_peer_id, remote)
                self._connections.append(conn)
                await conn_queue.put(conn)
                # Keep the handler alive while connection is open
                while conn.is_open:
                    await asyncio.sleep(1)
            except asyncio.TimeoutError:
                logger.warning("Noise handshake timed out from %s", remote)
            except Exception:
                logger.debug("Noise handshake failed from %s", remote, exc_info=True)

        self._server = await websockets.asyncio.server.serve(
            handler, host, port,
            max_size=10 * 1024 * 1024,
        )
        logger.info("A2A WebSocket server listening on %s:%d", host, port)

        try:
            while True:
                conn = await conn_queue.get()
                yield conn
        except asyncio.CancelledError:
            pass

    async def connect(self, addr: str) -> Connection:
        """Connect to a remote WebSocket address and perform Noise handshake."""
        parsed = urlparse(addr)
        ws_url = addr
        if parsed.scheme not in ("ws", "wss"):
            ws_url = f"ws://{addr}"

        try:
            ws = await websockets.asyncio.client.connect(
                ws_url,
                max_size=10 * 1024 * 1024,
                open_timeout=10,
            )
        except Exception as e:
            raise PeerUnreachable("unknown", reason=str(e)) from e

        try:
            session, remote_peer_id = await asyncio.wait_for(
                _do_noise_initiator(ws, self._identity),
                timeout=30.0,
            )
        except Exception as e:
            await ws.close()
            raise CryptoError(f"Noise handshake failed: {e}") from e

        remote = f"{parsed.hostname}:{parsed.port}"
        conn = WSConnection(ws, session, remote_peer_id, remote)
        self._connections.append(conn)
        return conn

    async def shutdown(self) -> None:
        """Close server and all connections."""
        for conn in self._connections:
            await conn.close()
        self._connections.clear()
        if self._server:
            self._server.close()
            await self._server.wait_closed()
            self._server = None
