"""Persistent heartbeat connection to license server.

Maintains a long-lived TCP/TLS connection to license.muravey.app
that survives CSCA network blocking (established connections persist).
Sends periodic ping frames to keep the connection alive.

Usage:
    from src.utils.heartbeat import Heartbeat

    hb = Heartbeat(host="license.muravey.app")
    hb.start()
    ...
    hb.stop()
"""

from __future__ import annotations

import base64
import hashlib
import logging
import os
import socket
import ssl
import struct
import threading
import time

log = logging.getLogger(__name__)

_WS_MAGIC = b"258EAFA5-E914-47DA-95CA-5AB9F0631073"


class Heartbeat:
    """Persistent heartbeat connection using WebSocket protocol.

    Establishes a TLS WebSocket connection and sends ping frames
    every `ping_interval` seconds. Auto-reconnects on disconnect.
    """

    def __init__(
        self,
        host: str = "license.muravey.app",
        port: int = 443,
        path: str = "/ws/heartbeat",
        ping_interval: float = 20.0,
        reconnect_delay: float = 5.0,
        max_reconnect_delay: float = 60.0,
        session_key: str = "",
    ) -> None:
        self.host = host
        self.port = port
        self.path = path
        self.ping_interval = ping_interval
        self.reconnect_delay = reconnect_delay
        self.max_reconnect_delay = max_reconnect_delay
        self.session_key = session_key

        self._sock: socket.socket | None = None
        self._ssl_sock: ssl.SSLSocket | None = None
        self._thread: threading.Thread | None = None
        self._stop_event = threading.Event()
        self._connected = threading.Event()
        self._consecutive_failures: int = 0

    @property
    def is_connected(self) -> bool:
        return self._connected.is_set()

    def start(self) -> None:
        """Start heartbeat in background thread."""
        if self._thread and self._thread.is_alive():
            return

        self._stop_event.clear()
        self._thread = threading.Thread(
            target=self._run_loop,
            daemon=True,
            name="heartbeat",
        )
        self._thread.start()
        log.info("Heartbeat started → %s:%d%s", self.host, self.port, self.path)

    def stop(self) -> None:
        """Stop heartbeat and close connection."""
        self._stop_event.set()
        self._close_connection()
        if self._thread:
            self._thread.join(timeout=5)
            self._thread = None
        log.info("Heartbeat stopped")

    def _run_loop(self) -> None:
        """Main loop: connect → ping → reconnect on failure."""
        while not self._stop_event.is_set():
            try:
                if self._connect():
                    self._consecutive_failures = 0
                    self._ping_loop()
                else:
                    self._consecutive_failures += 1
            except Exception as exc:
                log.debug("Heartbeat error: %s", exc)
                self._consecutive_failures += 1
            finally:
                self._close_connection()

            if self._stop_event.is_set():
                break

            # Exponential backoff for reconnect
            delay = min(
                self.reconnect_delay * (2 ** min(self._consecutive_failures, 5)),
                self.max_reconnect_delay,
            )
            log.debug(
                "Heartbeat reconnecting in %.1fs (attempt %d)",
                delay,
                self._consecutive_failures + 1,
            )
            self._stop_event.wait(delay)

    def _connect(self) -> bool:
        """Establish TLS + WebSocket connection."""
        try:
            # TCP connect
            self._sock = socket.create_connection(
                (self.host, self.port),
                timeout=10.0,
            )

            # Enable TCP keepalive
            self._sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)

            # TLS wrap
            ctx = ssl.create_default_context()
            self._ssl_sock = ctx.wrap_socket(self._sock, server_hostname=self.host)

            # WebSocket handshake
            ws_key = base64.b64encode(os.urandom(16)).decode()

            headers = [
                f"GET {self.path} HTTP/1.1",
                f"Host: {self.host}",
                "Upgrade: websocket",
                "Connection: Upgrade",
                f"Sec-WebSocket-Key: {ws_key}",
                "Sec-WebSocket-Version: 13",
            ]
            if self.session_key:
                headers.append(f"Authorization: Bearer {self.session_key}")
            headers.append("")
            headers.append("")

            self._ssl_sock.sendall("\r\n".join(headers).encode())

            # Read response (we don't strictly validate — just check for 101)
            response = b""
            while b"\r\n\r\n" not in response:
                chunk = self._ssl_sock.recv(4096)
                if not chunk:
                    return False
                response += chunk
                if len(response) > 8192:
                    return False

            status_line = response.split(b"\r\n")[0].decode()
            if "101" not in status_line:
                # Server doesn't support WebSocket — fall back to plain keepalive
                log.debug(
                    "WebSocket upgrade rejected (%s), using TCP keepalive only",
                    status_line.strip(),
                )
                self._connected.set()
                return True

            self._connected.set()
            log.debug("WebSocket connected to %s:%d", self.host, self.port)
            return True

        except Exception as exc:
            log.debug("Heartbeat connect failed: %s", exc)
            return False

    def _ping_loop(self) -> None:
        """Send periodic ping frames."""
        while not self._stop_event.is_set():
            self._stop_event.wait(self.ping_interval)
            if self._stop_event.is_set():
                break

            try:
                self._send_ping()
            except Exception:
                log.debug("Heartbeat ping failed — connection lost")
                break

    def _send_ping(self) -> None:
        """Send a WebSocket ping frame."""
        if not self._ssl_sock:
            raise ConnectionError("Not connected")

        # WebSocket ping frame: FIN=1, opcode=0x9 (ping), MASK=1
        payload = struct.pack("!d", time.time())  # 8-byte timestamp
        mask_key = os.urandom(4)

        # Frame header
        frame = bytearray()
        frame.append(0x89)  # FIN + ping opcode

        payload_len = len(payload)
        if payload_len < 126:
            frame.append(0x80 | payload_len)  # MASK bit + length
        else:
            frame.append(0x80 | 126)
            frame.extend(struct.pack("!H", payload_len))

        frame.extend(mask_key)

        # Mask payload
        masked = bytearray(payload_len)
        for i in range(payload_len):
            masked[i] = payload[i] ^ mask_key[i % 4]
        frame.extend(masked)

        self._ssl_sock.sendall(bytes(frame))

    def _close_connection(self) -> None:
        """Close the connection cleanly."""
        self._connected.clear()
        if self._ssl_sock:
            try:
                self._ssl_sock.close()
            except Exception:
                pass
            self._ssl_sock = None
        if self._sock:
            try:
                self._sock.close()
            except Exception:
                pass
            self._sock = None

    def __del__(self) -> None:
        self.stop()
