"""
SSP Relay client for 2-of-2 multisig coordination.

Handles communication between fluxwallet (key #1) and SSP Key mobile app
(key #2) via the SSP Relay server. Uses Socket.IO for real-time push
events and httpx for action POSTs.
"""

from __future__ import annotations

import asyncio
import base64
import hashlib
import json
import logging
import secrets
import ssl
import time

import aiohttp
import certifi
import httpx
import socketio

from fluxwallet.keys import HDKey, sign_message

_logger = logging.getLogger(__name__)


class RelayError(Exception):
    def __init__(self, msg: str = ""):
        self.msg = msg
        _logger.error(msg)

    def __str__(self) -> str:
        return self.msg


class SSPRelayClient:
    """Async client for the SSP Relay server.

    Connects via Socket.IO to receive real-time events (txid, txrejected)
    and uses HTTP POST for sending actions to SSP Key.
    """

    def __init__(
        self,
        relay_url: str,
        identity_key: HDKey,
        witness_script: str = "",
        identity_pubkey: str = "",
    ):
        """
        :param relay_url: Relay server URL (e.g. 'relay.sspwallet.io')
        :param identity_key: Private key for signing (identity chain, typeIndex=10)
        :param witness_script: Hex witness script for multisig identity auth
        :param identity_pubkey: Hex compressed pubkey for multisig identity auth
        """
        self.relay_url = relay_url.rstrip("/")
        if not self.relay_url.startswith("http"):
            self.relay_url = f"https://{self.relay_url}"
        self.identity_key = identity_key
        self.identity_address = identity_key.address()
        self.identity_pubkey = identity_pubkey or identity_key.public_hex
        self.witness_script = witness_script

        self._ssl_context = ssl.create_default_context(cafile=certifi.where())
        self._http = httpx.AsyncClient(
            timeout=30,
            headers={"Accept": "application/json", "Content-Type": "application/json"},
            verify=self._ssl_context,
        )
        self._aio_session: aiohttp.ClientSession | None = None
        self._sio: socketio.AsyncClient | None = None
        self._connected = False

        # Async events for receiving push responses
        self._txid_event = asyncio.Event()
        self._txid_value: str | None = None
        self._txrejected_event = asyncio.Event()
        self._txrejected_value: str | None = None
        self._sync_event = asyncio.Event()
        self._sync_value: dict | None = None

    def _ensure_socketio(self) -> socketio.AsyncClient:
        """Lazily create aiohttp session and socketio client (requires event loop)."""
        if self._sio is None:
            connector = aiohttp.TCPConnector(ssl=self._ssl_context)
            self._aio_session = aiohttp.ClientSession(connector=connector)
            self._sio = socketio.AsyncClient(
                reconnection=True, reconnection_attempts=100,
                http_session=self._aio_session,
            )
            self._sio.on("txid", self._on_txid)
            self._sio.on("txrejected", self._on_txrejected)
            self._sio.on("sync", self._on_sync)
            self._sio.on("connect", self._on_connect)
            self._sio.on("disconnect", self._on_disconnect)
        return self._sio

    # --- Auth ---

    def _build_auth(self, action: str, wk_identity: str, body: dict | None = None) -> dict:
        """Build authentication fields for a relay request.

        Matches the SSP relay auth protocol: signs a JSON payload containing
        timestamp, action, identity (wkIdentity), nonce, and optional body hash.
        Uses Bitcoin message signing prefix.

        :param action: Action type ('action', 'join', 'sync', 'token')
        :param wk_identity: The multisig wkIdentity address
        :param body: Optional request body to hash
        :returns: Dict with signature, message, publicKey, witnessScript fields
        """
        payload = {
            "timestamp": int(time.time() * 1000),
            "action": action,
            "identity": wk_identity,
            "nonce": secrets.token_hex(32),
        }

        if body is not None:
            body_hash = hashlib.sha256(json.dumps(body, separators=(",", ":")).encode()).hexdigest()
            payload["data"] = body_hash

        message = json.dumps(payload, separators=(",", ":"))

        # Sign with Bitcoin message prefix (relay always uses BTC for identity)
        btc_prefix = b"\x18Bitcoin Signed Message:\n"
        sig_bytes, _ = sign_message(message, self.identity_key, message_prefix=btc_prefix)

        auth = {
            "signature": base64.b64encode(sig_bytes).decode(),
            "message": message,
            "publicKey": self.identity_pubkey,
        }
        if self.witness_script:
            auth["witnessScript"] = self.witness_script
        return auth

    # --- Socket.IO ---

    async def connect(self, wk_identity: str) -> None:
        """Connect to the relay via Socket.IO and join the wallet room.

        :param wk_identity: The 2-of-2 multisig identity address
        """
        if self._connected:
            return

        sio = self._ensure_socketio()
        await sio.connect(
            self.relay_url,
            socketio_path="/v1/socket/wallet",
            wait_timeout=10,
        )

        # Emit authenticated join
        join_data = {"wkIdentity": wk_identity}
        auth = self._build_auth("join", wk_identity)
        await sio.emit("join", {**join_data, **auth})
        _logger.info("Joined relay room for %s", wk_identity)

    async def disconnect(self) -> None:
        """Disconnect from the relay."""
        if self._connected and self._sio:
            await self._sio.disconnect()

    async def _on_connect(self) -> None:
        self._connected = True
        _logger.info("Connected to SSP Relay")

    async def _on_disconnect(self) -> None:
        self._connected = False
        _logger.info("Disconnected from SSP Relay")

    async def _on_txid(self, data: dict) -> None:
        _logger.info("Received txid: %s", data.get("payload", "")[:16])
        self._txid_value = data.get("payload", "")
        self._txid_event.set()

    async def _on_txrejected(self, data: dict) -> None:
        _logger.warning("Transaction rejected: %s", data.get("payload", ""))
        self._txrejected_value = data.get("payload", "")
        self._txrejected_event.set()

    async def _on_sync(self, data: dict) -> None:
        _logger.info("Received sync data")
        self._sync_value = data
        self._sync_event.set()

    # --- HTTP Actions ---

    async def post_action(
        self,
        action: str,
        payload: str,
        chain: str,
        path: str,
        wk_identity: str,
        **extra,
    ) -> dict:
        """POST an action to the relay for SSP Key to process.

        :param action: Action type ('tx', 'pair', etc.)
        :param payload: Payload string (e.g. partially signed tx hex)
        :param chain: Chain identifier ('flux')
        :param path: Derivation path ('0-0' = typeIndex-addressIndex)
        :param wk_identity: Multisig identity address
        :returns: Relay response
        """
        data = {
            "action": action,
            "payload": payload,
            "chain": chain,
            "path": path,
            "wkIdentity": wk_identity,
            **extra,
        }

        auth = self._build_auth("action", wk_identity, body=data)
        data.update(auth)

        response = await self._http.post(f"{self.relay_url}/v1/action", json=data)
        response.raise_for_status()
        return response.json()

    # --- Event Waiters ---

    async def wait_for_txid(self, timeout: float = 300) -> str:
        """Wait for SSP Key to broadcast and return a txid.

        :param timeout: Maximum seconds to wait
        :returns: Transaction ID
        :raises RelayError: On timeout or rejection
        """
        self._txid_event.clear()
        self._txrejected_event.clear()

        done, _ = await asyncio.wait(
            [
                asyncio.create_task(self._txid_event.wait()),
                asyncio.create_task(self._txrejected_event.wait()),
            ],
            timeout=timeout,
            return_when=asyncio.FIRST_COMPLETED,
        )

        if not done:
            raise RelayError("Timeout waiting for SSP Key response")

        if self._txrejected_event.is_set():
            raise RelayError(f"Transaction rejected by SSP Key: {self._txrejected_value}")

        return self._txid_value

    async def wait_for_sync(self, timeout: float = 300) -> dict:
        """Wait for SSP Key to respond with its xpub during pairing.

        :param timeout: Maximum seconds to wait
        :returns: Sync data dict with keyXpub, wkIdentity, etc.
        :raises RelayError: On timeout
        """
        self._sync_event.clear()

        try:
            await asyncio.wait_for(self._sync_event.wait(), timeout=timeout)
        except asyncio.TimeoutError:
            raise RelayError("Timeout waiting for SSP Key pairing")

        return self._sync_value

    # --- Fallback Polling ---

    async def poll_sync(
        self, identity: str, timeout: float = 300, chain: str | None = None,
    ) -> dict | None:
        """Poll GET /v1/sync/{identity} for peer's xpub.

        :param identity: Wallet identity address
        :param timeout: Maximum seconds to poll
        :param chain: If set, only return data matching this chain (e.g. 'btc', 'flux')
        :returns: Sync data or None
        """
        deadline = asyncio.get_event_loop().time() + timeout
        while asyncio.get_event_loop().time() < deadline:
            try:
                resp = await self._http.get(f"{self.relay_url}/v1/sync/{identity}")
                if resp.status_code == 200:
                    data = resp.json()
                    if data and data.get("keyXpub"):
                        if chain and data.get("chain") != chain:
                            await asyncio.sleep(2)
                            continue
                        return data
            except Exception as e:
                _logger.debug("Sync poll error: %s", e)
            await asyncio.sleep(2)
        return None

    async def poll_action(self, wk_identity: str, timeout: float = 300) -> dict | None:
        """Fallback: poll GET /v1/action/{wkIdentity} for response.

        :param wk_identity: Multisig identity address
        :param timeout: Maximum seconds to poll
        :returns: Action response or None
        """
        deadline = asyncio.get_event_loop().time() + timeout
        while asyncio.get_event_loop().time() < deadline:
            try:
                resp = await self._http.get(f"{self.relay_url}/v1/action/{wk_identity}")
                if resp.status_code == 200:
                    data = resp.json()
                    if data and data.get("action") == "txid":
                        return data
            except Exception as e:
                _logger.debug("Action poll error: %s", e)
            await asyncio.sleep(2)
        return None

    # --- Cleanup ---

    async def close(self) -> None:
        """Close all connections."""
        await self.disconnect()
        await self._http.aclose()
        if self._aio_session:
            await self._aio_session.close()
