"""High-level quicproquo client.

Provides both async (QUIC transport) and sync (FFI transport) APIs for
interacting with a quicproquo server.
"""

from __future__ import annotations

from typing import Optional

from quicproquo.types import (
    ConnectOptions,
    Envelope,
    ChannelResult,
    HealthInfo,
    ConnectionError,
)
from quicproquo.transport import QuicTransport
from quicproquo.ffi import FfiTransport
from quicproquo import proto, wire


class QpqClient:
    """High-level quicproquo client.

    Use ``QpqClient.connect()`` for the async QUIC transport, or
    ``QpqClient.connect_ffi()`` for the synchronous Rust FFI backend.

    Example (async)::

        client = await QpqClient.connect(ConnectOptions(addr="127.0.0.1:5001"))
        health = await client.health()
        token = await client.login_start("alice", opaque_request)
        await client.close()

    Example (FFI)::

        client = QpqClient.connect_ffi(ConnectOptions(addr="127.0.0.1:5001"))
        client.ffi_login("alice", "password123")
        client.ffi_send("bob", b"hello")
        client.close()
    """

    def __init__(
        self,
        quic: Optional[QuicTransport] = None,
        ffi: Optional[FfiTransport] = None,
    ) -> None:
        self._quic = quic
        self._ffi = ffi
        self._session_token: bytes = b""
        self._device_id: bytes = b""

    # ------------------------------------------------------------------
    # Constructors
    # ------------------------------------------------------------------

    @staticmethod
    async def connect(opts: ConnectOptions) -> "QpqClient":
        """Connect to a server using the async QUIC transport."""
        transport = await QuicTransport.connect(
            opts.addr,
            ca_cert_path=opts.ca_cert_path,
            server_name=opts.server_name,
            insecure_skip_verify=opts.insecure_skip_verify,
            connect_timeout_ms=opts.connect_timeout_ms,
            request_timeout_ms=opts.request_timeout_ms,
        )
        return QpqClient(quic=transport)

    @staticmethod
    def connect_ffi(opts: ConnectOptions) -> "QpqClient":
        """Connect using the synchronous Rust FFI backend."""
        transport = FfiTransport.connect(
            opts.addr,
            ca_cert_path=opts.ca_cert_path,
            server_name=opts.server_name,
        )
        return QpqClient(ffi=transport)

    # ------------------------------------------------------------------
    # Session management
    # ------------------------------------------------------------------

    def set_session_token(self, token: bytes) -> None:
        """Set an externally-obtained session token for authenticated RPCs."""
        self._session_token = token

    def set_device_id(self, device_id: bytes) -> None:
        """Set the device ID for multi-device scoping."""
        self._device_id = device_id

    # ------------------------------------------------------------------
    # Async RPC methods (QUIC transport)
    # ------------------------------------------------------------------

    async def health(self) -> HealthInfo:
        """Check server health (async)."""
        data = await self._rpc(wire.HEALTH, proto.encode_health())
        info = proto.decode_health_response(data)
        return HealthInfo(
            status=str(info["status"]),
            node_id=str(info["node_id"]),
            version=str(info["version"]),
            uptime_secs=int(info["uptime_secs"]),
            storage_backend=str(info["storage_backend"]),
        )

    async def register_start(self, username: str, request: bytes) -> bytes:
        """Start OPAQUE registration. Returns server response bytes."""
        payload = proto.encode_opaque_register_start(username, request)
        data = await self._rpc(wire.OPAQUE_REGISTER_START, payload)
        return proto.decode_opaque_register_start_response(data)

    async def register_finish(
        self, username: str, upload: bytes, identity_key: bytes
    ) -> bool:
        """Complete OPAQUE registration."""
        payload = proto.encode_opaque_register_finish(username, upload, identity_key)
        data = await self._rpc(wire.OPAQUE_REGISTER_FINISH, payload)
        return proto.decode_opaque_register_finish_response(data)

    async def login_start(self, username: str, request: bytes) -> bytes:
        """Start OPAQUE login. Returns server response bytes."""
        payload = proto.encode_opaque_login_start(username, request)
        data = await self._rpc(wire.OPAQUE_LOGIN_START, payload)
        return proto.decode_opaque_login_start_response(data)

    async def login_finish(
        self, username: str, finalization: bytes, identity_key: bytes
    ) -> bytes:
        """Complete OPAQUE login. Returns and stores session token."""
        payload = proto.encode_opaque_login_finish(username, finalization, identity_key)
        data = await self._rpc(wire.OPAQUE_LOGIN_FINISH, payload)
        token = proto.decode_opaque_login_finish_response(data)
        self._session_token = token
        return token

    async def resolve_user(self, username: str) -> tuple[bytes, bytes]:
        """Resolve username to (identity_key, inclusion_proof)."""
        payload = proto.encode_resolve_user(username)
        data = await self._rpc(wire.RESOLVE_USER, payload)
        return proto.decode_resolve_user_response(data)

    async def resolve_identity(self, identity_key: bytes) -> str:
        """Resolve identity key to username."""
        payload = proto.encode_resolve_identity(identity_key)
        data = await self._rpc(wire.RESOLVE_IDENTITY, payload)
        return proto.decode_resolve_identity_response(data)

    async def create_channel(self, peer_key: bytes) -> ChannelResult:
        """Create a 1:1 DM channel with a peer."""
        payload = proto.encode_create_channel(peer_key)
        data = await self._rpc(wire.CREATE_CHANNEL, payload)
        channel_id, was_new = proto.decode_create_channel_response(data)
        return ChannelResult(channel_id=channel_id, was_new=was_new)

    async def send(
        self,
        recipient_key: bytes,
        payload: bytes,
        *,
        channel_id: bytes = b"",
        ttl_secs: int = 0,
        message_id: bytes = b"",
    ) -> tuple[int, bytes]:
        """Enqueue a message. Returns (seq, delivery_proof)."""
        req = proto.encode_enqueue(recipient_key, payload, channel_id, ttl_secs, message_id)
        data = await self._rpc(wire.ENQUEUE, req)
        seq, proof, _ = proto.decode_enqueue_response(data)
        return seq, proof

    async def receive(
        self,
        recipient_key: bytes,
        *,
        channel_id: bytes = b"",
        limit: int = 0,
        device_id: bytes = b"",
    ) -> list[Envelope]:
        """Fetch queued messages."""
        req = proto.encode_fetch(recipient_key, channel_id, limit, device_id)
        data = await self._rpc(wire.FETCH, req)
        return [Envelope(seq=s, data=d) for s, d in proto.decode_fetch_response(data)]

    async def receive_wait(
        self,
        recipient_key: bytes,
        *,
        timeout_ms: int = 5000,
        channel_id: bytes = b"",
        limit: int = 0,
        device_id: bytes = b"",
    ) -> list[Envelope]:
        """Long-poll for messages with a timeout."""
        req = proto.encode_fetch_wait(recipient_key, channel_id, timeout_ms, limit, device_id)
        data = await self._rpc(wire.FETCH_WAIT, req)
        return [Envelope(seq=s, data=d) for s, d in proto.decode_fetch_wait_response(data)]

    async def ack(
        self,
        recipient_key: bytes,
        seq_up_to: int,
        *,
        channel_id: bytes = b"",
        device_id: bytes = b"",
    ) -> None:
        """Acknowledge messages up to a sequence number."""
        req = proto.encode_ack(recipient_key, seq_up_to, channel_id, device_id)
        await self._rpc(wire.ACK, req)

    async def upload_key_package(self, identity_key: bytes, package: bytes) -> bytes:
        """Upload an MLS key package. Returns fingerprint."""
        req = proto.encode_upload_key_package(identity_key, package)
        data = await self._rpc(wire.UPLOAD_KEY_PACKAGE, req)
        return proto.decode_upload_key_package_response(data)

    async def fetch_key_package(self, identity_key: bytes) -> bytes:
        """Fetch an MLS key package."""
        req = proto.encode_fetch_key_package(identity_key)
        data = await self._rpc(wire.FETCH_KEY_PACKAGE, req)
        return proto.decode_fetch_key_package_response(data)

    async def upload_hybrid_key(self, identity_key: bytes, hybrid_public_key: bytes) -> None:
        """Upload a hybrid (X25519 + ML-KEM-768) public key."""
        req = proto.encode_upload_hybrid_key(identity_key, hybrid_public_key)
        await self._rpc(wire.UPLOAD_HYBRID_KEY, req)

    async def fetch_hybrid_key(self, identity_key: bytes) -> bytes:
        """Fetch a hybrid public key."""
        req = proto.encode_fetch_hybrid_key(identity_key)
        data = await self._rpc(wire.FETCH_HYBRID_KEY, req)
        return proto.decode_fetch_hybrid_key_response(data)

    async def delete_account(self) -> bool:
        """Permanently delete the authenticated account."""
        data = await self._rpc(wire.DELETE_ACCOUNT, proto.encode_delete_account())
        return proto.decode_delete_account_response(data)

    # ------------------------------------------------------------------
    # FFI (synchronous) methods
    # ------------------------------------------------------------------

    def ffi_login(self, username: str, password: str) -> None:
        """Authenticate via OPAQUE using the FFI backend (synchronous)."""
        if not self._ffi:
            raise ConnectionError("no FFI transport; use QpqClient.connect_ffi()")
        self._ffi.login(username, password)

    def ffi_send(self, recipient: str, message: bytes) -> None:
        """Send a message via the FFI backend (synchronous)."""
        if not self._ffi:
            raise ConnectionError("no FFI transport; use QpqClient.connect_ffi()")
        self._ffi.send(recipient, message)

    def ffi_receive(self, timeout_ms: int = 5000) -> list[str]:
        """Receive messages via the FFI backend (synchronous)."""
        if not self._ffi:
            raise ConnectionError("no FFI transport; use QpqClient.connect_ffi()")
        return self._ffi.receive(timeout_ms)

    # ------------------------------------------------------------------
    # Lifecycle
    # ------------------------------------------------------------------

    async def close(self) -> None:
        """Close all transports."""
        if self._quic:
            self._quic.close()
            self._quic = None
        if self._ffi:
            self._ffi.close()
            self._ffi = None

    def close_sync(self) -> None:
        """Close all transports (synchronous variant)."""
        if self._quic:
            self._quic.close()
            self._quic = None
        if self._ffi:
            self._ffi.close()
            self._ffi = None

    # ------------------------------------------------------------------
    # Internal
    # ------------------------------------------------------------------

    async def _rpc(self, method_id: int, payload: bytes) -> bytes:
        if not self._quic:
            raise ConnectionError("no QUIC transport; use QpqClient.connect()")
        return await self._quic.rpc(method_id, payload)
