"""Browser state tracking model."""

import asyncio
import logging
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, Dict, Optional

logger = logging.getLogger(__name__)


@dataclass
class BrowserConnection:
    """Represents a browser connection."""

    port: int  # client ephemeral port (internal tracking)
    server_port: int  # server listening port (user-facing, for MCP tools)
    connected_at: datetime
    last_message_at: Optional[datetime] = None
    message_count: int = 0
    url: Optional[str] = None
    user_agent: Optional[str] = None
    websocket: Any = None  # WebSocket connection object
    is_active: bool = True

    def update_activity(self) -> None:
        """Update last activity timestamp."""
        self.last_message_at = datetime.now()
        self.message_count += 1

    def disconnect(self) -> None:
        """Mark connection as disconnected."""
        self.is_active = False

    @property
    def connection_duration(self) -> float:
        """Get connection duration in seconds."""
        return (datetime.now() - self.connected_at).total_seconds()

    @property
    def idle_time(self) -> float:
        """Get idle time since last message in seconds."""
        if self.last_message_at:
            return (datetime.now() - self.last_message_at).total_seconds()
        return self.connection_duration


@dataclass
class BrowserState:
    """Manages state for all browser connections."""

    connections: Dict[int, BrowserConnection] = field(default_factory=dict)
    server_port_map: Dict[int, int] = field(
        default_factory=dict
    )  # server_port → client_port mapping
    _lock: asyncio.Lock = field(default_factory=asyncio.Lock)

    async def add_connection(
        self,
        port: int,
        server_port: int,
        websocket: Any,
        user_agent: Optional[str] = None,
    ) -> BrowserConnection:
        """Add a new browser connection.

        Args:
            port: Client ephemeral port number (internal tracking)
            server_port: Server listening port number (user-facing, for MCP tools)
            websocket: WebSocket connection object
            user_agent: Optional user agent string

        Returns:
            BrowserConnection instance
        """
        async with self._lock:
            connection = BrowserConnection(
                port=port,
                server_port=server_port,
                connected_at=datetime.now(),
                websocket=websocket,
                user_agent=user_agent,
            )
            self.connections[port] = connection
            # Map server port to client port for user-facing API
            self.server_port_map[server_port] = port
            return connection

    async def remove_connection(self, port: int) -> None:
        """Remove a browser connection.

        Args:
            port: Port number to remove (client port)
        """
        async with self._lock:
            if port in self.connections:
                # Clean up server_port_map first
                conn = self.connections[port]
                if conn.server_port in self.server_port_map:
                    del self.server_port_map[conn.server_port]

                # Then remove connection
                self.connections[port].disconnect()
                del self.connections[port]

    async def get_connection(self, port: int) -> Optional[BrowserConnection]:
        """Get a browser connection by port.

        Supports both server port (8851-8899) and client ephemeral port lookups.
        When a server port is provided, it's automatically mapped to the client port.

        Args:
            port: Port number (can be server port or client port)

        Returns:
            BrowserConnection if exists, None otherwise
        """
        async with self._lock:
            # Try server port mapping first (if it's a server port like 8852)
            lookup_port = self.server_port_map.get(port, port)
            return self.connections.get(lookup_port)

    async def update_connection_activity(self, port: int) -> None:
        """Update connection activity timestamp.

        Args:
            port: Port number
        """
        async with self._lock:
            if port in self.connections:
                self.connections[port].update_activity()

    async def update_connection_url(self, port: int, url: str) -> None:
        """Update the current URL for a connection.

        Args:
            port: Port number
            url: Current URL
        """
        async with self._lock:
            if port in self.connections:
                self.connections[port].url = url

    async def get_any_active_connection(self) -> Optional[BrowserConnection]:
        """Get any active connection (fallback when server port is used instead of client port).

        This is useful when MCP tools pass the server port (8851-8895) but connections
        are keyed by the client's ephemeral port (e.g., 57803).

        Returns:
            First active connection found, or None if no active connections exist
        """
        async with self._lock:
            logger.info(
                f"get_any_active_connection: {len(self.connections)} total connections"
            )
            for port, conn in self.connections.items():
                logger.info(
                    f"  Port {port}: is_active={conn.is_active}, websocket={conn.websocket is not None}"
                )
                if conn.is_active:
                    return conn
            return None

    async def get_active_connections(self) -> Dict[int, BrowserConnection]:
        """Get all active connections.

        Returns:
            Dictionary of active connections
        """
        async with self._lock:
            return {
                port: conn for port, conn in self.connections.items() if conn.is_active
            }

    async def get_connection_stats(self) -> Dict[str, Any]:
        """Get statistics about all connections.

        Returns:
            Dictionary with connection statistics
        """
        async with self._lock:
            active_connections = [c for c in self.connections.values() if c.is_active]
            total_messages = sum(c.message_count for c in self.connections.values())

            return {
                "total_connections": len(self.connections),
                "active_connections": len(active_connections),
                "total_messages": total_messages,
                "ports": list(self.connections.keys()),
                "connections": [
                    {
                        "port": c.port,
                        "connected_at": c.connected_at.isoformat(),
                        "message_count": c.message_count,
                        "url": c.url,
                        "is_active": c.is_active,
                        "duration": c.connection_duration,
                        "idle_time": c.idle_time,
                    }
                    for c in self.connections.values()
                ],
            }
