"""
Socket.IO connection for IQ hosts.

Handles persistent connections for worker registration and task execution.
"""

from __future__ import annotations

import asyncio
import base64
import sys
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any

import socketio  # type: ignore[import-untyped]

if TYPE_CHECKING:
    from iq_sdk.config import SiteConfig


def _debug(msg: str, *, prefix: str = "[Worker]") -> None:
    """Print debug message to stderr."""
    print(f"{prefix} {msg}", file=sys.stderr)


@dataclass
class ConnectionState:
    """Current state of connection."""

    connected: bool = False
    subscribed_tools: list[str] = field(default_factory=list)
    worker_id: str | None = None
    error: str | None = None


class IQConnection:
    """
    Low-level Socket.IO connection to IQ host.

    Manages connection lifecycle and basic Socket.IO operations.
    """

    DOCTYPE = "Automation Tool"
    PING_TIMEOUT = 5.0  # Seconds to wait for pong response

    def __init__(
        self,
        host: str,
        socket_io_url: str,
        api_key: str,
        api_secret: str,
        namespace: str,
        worker_id: str | None = None,
    ):
        """
        Initialize connection.

        Args:
            host: Origin host (e.g. http://localhost:8000)
            socket_io_url: Socket.IO server URL
            api_key: API key for Basic Auth
            api_secret: API secret for Basic Auth
            namespace: Socket.IO namespace (from check_permissions socket_io_ns)
            worker_id: Unique worker identifier
        """
        self.host = host
        self.socket_io_url = socket_io_url
        self.api_key = api_key
        self.api_secret = api_secret
        self.namespace = namespace
        self.worker_id = worker_id

        self._sio: socketio.AsyncClient | None = None
        self._state = ConnectionState()
        self._debug = False
        self._disconnect_callback: asyncio.Event | None = None

    @property
    def state(self) -> ConnectionState:
        """Get current connection state."""
        return self._state

    @property
    def is_connected(self) -> bool:
        """Check if connected."""
        return self._state.connected

    async def connect(self, debug: bool = False) -> bool:
        """
        Establish connection to host.

        Args:
            debug: Enable debug output

        Returns:
            True if connection successful
        """
        self._debug = debug

        if self._sio is not None and self._state.connected:
            if debug:
                _debug("Already connected, reusing connection")
            return True

        if debug:
            _debug(f"Connecting to {self.socket_io_url}")
            _debug(f"  Namespace: {self.namespace}")
            _debug(f"  Worker ID: {self.worker_id}")

        # Don't use logger/engineio_logger - they expose secrets in logs
        self._sio = socketio.AsyncClient(
            reconnection=False,  # We handle reconnection ourselves
        )

        # Register connect/disconnect handlers
        @self._sio.on("connect", namespace=self.namespace)  # type: ignore[misc]
        async def on_connect() -> None:
            self._state.connected = True
            self._state.error = None
            if debug:
                _debug("✓ Connected")

        @self._sio.on("disconnect", namespace=self.namespace)  # type: ignore[misc]
        async def on_disconnect() -> None:
            self._state.connected = False
            if debug:
                _debug("Disconnected")
            # Signal disconnect to any waiting loops
            if self._disconnect_callback:
                self._disconnect_callback.set()

        try:
            # Build auth headers
            credentials = f"{self.api_key}:{self.api_secret}"
            encoded = base64.b64encode(credentials.encode()).decode()

            headers = {
                "Authorization": f"Basic {encoded}",
                "Cookie": "auth=it;",  # Dummy cookie required
                "Origin": self.host,
            }

            await self._sio.connect(
                self.socket_io_url,
                socketio_path="/socket.io",
                headers=headers,
                namespaces=[self.namespace],
                wait_timeout=10,
            )

            return True

        except Exception as e:
            self._state.error = str(e)
            if debug:
                _debug(f"✗ Connection failed: {e}")
            return False

    async def disconnect(self) -> None:
        """Disconnect from host."""
        if self._sio:
            try:
                await self._sio.disconnect()
            except Exception:
                pass
            self._sio = None
            self._state.connected = False
            self._state.subscribed_tools = []

    async def ping(self, timeout: float = PING_TIMEOUT) -> bool:
        """
        Send ping and wait for pong to verify connection and auth.

        Args:
            timeout: Seconds to wait for pong

        Returns:
            True if pong received within timeout
        """
        if not self._sio or not self._state.connected:
            return False

        pong_received = asyncio.Event()

        @self._sio.on("pong", namespace=self.namespace)  # type: ignore[misc]
        async def on_pong(data: Any = None) -> None:
            pong_received.set()

        try:
            await self._sio.emit("ping", namespace=self.namespace)

            await asyncio.wait_for(pong_received.wait(), timeout=timeout)
            return True

        except TimeoutError:
            self._state.error = f"Ping timeout after {timeout}s"
            return False
        except Exception as e:
            self._state.error = str(e)
            return False

    async def subscribe_tools(self, tool_names: list[str]) -> None:
        """
        Subscribe to multiple tools for wake-up notifications.

        Args:
            tool_names: List of full tool names
        """
        if not self._sio or not self._state.connected:
            return

        for tool_name in tool_names:
            await self._sio.emit(
                "doc_subscribe",
                (self.DOCTYPE, tool_name),
                namespace=self.namespace,
            )
            self._state.subscribed_tools.append(tool_name)

        if self._debug:
            _debug(f"Subscribed to {len(tool_names)} tools")

    async def unsubscribe_tools(self, tool_names: list[str]) -> None:
        """
        Unsubscribe from tools.

        Args:
            tool_names: List of full tool names
        """
        if not self._sio or not self._state.connected:
            return

        for tool_name in tool_names:
            await self._sio.emit(
                "doc_unsubscribe",
                (self.DOCTYPE, tool_name),
                namespace=self.namespace,
            )
            if tool_name in self._state.subscribed_tools:
                self._state.subscribed_tools.remove(tool_name)

    def on_docinfo_update(self, callback: Any) -> None:
        """
        Register handler for docinfo_update events.

        Args:
            callback: Async callback function(data)
        """
        if self._sio:

            @self._sio.on("docinfo_update", namespace=self.namespace)  # type: ignore[misc]
            async def on_wake(data: Any) -> None:
                await callback(data)

    def set_disconnect_callback(self, event: asyncio.Event) -> None:
        """Set event to signal on disconnect."""
        self._disconnect_callback = event

    @classmethod
    def from_site(
        cls,
        site: SiteConfig,
        socket_io_url: str,
        namespace: str,
        worker_id: str | None = None,
    ) -> IQConnection:
        """
        Create connection from site config.

        Args:
            site: Site configuration
            socket_io_url: Socket.IO URL from check_permissions
            namespace: Socket.IO namespace from check_permissions (socket_io_ns)
            worker_id: Worker ID

        Returns:
            Configured IQConnection
        """
        return cls(
            host=site.get_host(),
            socket_io_url=socket_io_url,
            api_key=site.get_api_key() or "",
            api_secret=site.load_api_secret() or "",
            namespace=namespace,
            worker_id=worker_id,
        )


async def test_worker_connection(
    site: SiteConfig,
    project_name: str,
    tool_name: str,
    worker_id: str,
    debug: bool = True,
) -> bool:
    """
    Test worker connection with ping/pong.

    Connects to Socket.IO, verifies auth with ping/pong, then disconnects.

    Args:
        site: Site configuration
        project_name: Project namespace for permission check
        tool_name: Full tool name (e.g. "demo.my_tool.1.0")
        worker_id: Worker ID to identify with
        debug: Enable debug output

    Returns:
        True if ping/pong successful
    """
    from iq_sdk.client import IQApiClient

    if debug:
        _debug("=" * 50)
        _debug(f"Testing worker connection for site: {site.id}")
        _debug(f"Project: {project_name}")
        _debug(f"Worker ID: {worker_id}")
        _debug(f"Tool: {tool_name}")
        _debug("=" * 50)

    # Step 1: Check permissions and get Socket.IO URL
    if debug:
        _debug("\nStep 1: Checking permissions...")

    async with IQApiClient.from_site(site) as client:
        status = await client.check_status(namespace=project_name)

    if not status.reachable:
        if debug:
            _debug(f"✗ Site not reachable: {status.error}")
        return False

    if not status.authenticated:
        if debug:
            _debug(f"✗ {status.error}")
        return False

    if not status.can_register_worker:
        if debug:
            _debug(f"✗ No permission to register worker for Project: {project_name}")
        return False

    if not status.socket_io_url or not status.socket_io_ns:
        if debug:
            _debug("✗ No Socket.IO URL or namespace provided")
        return False

    if debug:
        _debug("✓ Permissions OK")
        _debug(f"  Socket.IO URL: {status.socket_io_url}")
        _debug(f"  Namespace: {status.socket_io_ns}")

    # Step 2: Connect via Socket.IO
    if debug:
        _debug("\nStep 2: Connecting via Socket.IO...")

    conn = IQConnection.from_site(
        site, status.socket_io_url, status.socket_io_ns, worker_id=worker_id
    )

    connected = await conn.connect(debug=debug)
    if not connected:
        if debug:
            _debug(f"✗ Socket.IO connection failed: {conn.state.error}")
        return False

    # Step 3: Ping/Pong test
    if debug:
        _debug("\nStep 3: Testing ping/pong...")

    try:
        pong_ok = await conn.ping(timeout=5.0)

        if pong_ok:
            if debug:
                _debug("✓ Ping/pong successful!")
                _debug("\n" + "=" * 50)
                _debug("✓ Connection test PASSED")
                _debug("=" * 50)
            return True
        else:
            if debug:
                _debug(f"✗ Ping failed: {conn.state.error}")
            return False

    finally:
        await conn.disconnect()
        if debug:
            _debug("\nConnection closed.")
