"""Connection management for GQLDB Python driver."""

from __future__ import annotations

import threading
import time
from dataclasses import dataclass
from typing import Dict, Optional

import grpc

from .config import GqldbConfig
from .errors import AllHostsFailedError, ConnectionClosedError, NoConnectionError


@dataclass
class Connection:
    """Represents a single gRPC connection."""
    host: str
    channel: grpc.Channel
    healthy: bool = True
    last_ping: float = 0.0


class ConnectionPool:
    """Manages a pool of connections to GQLDB servers."""

    def __init__(self, config: GqldbConfig):
        """Initialize the connection pool.

        Args:
            config: Configuration for connections.

        Raises:
            AllHostsFailedError: If all hosts fail to connect.
        """
        self._config = config
        self._connections: Dict[str, Connection] = {}
        self._lock = threading.RLock()
        self._closed = False
        self._stop_event = threading.Event()
        self._health_thread: Optional[threading.Thread] = None

        # Initialize connections to all hosts
        for host in config.hosts:
            try:
                conn = self._create_connection(host)
                self._connections[host] = conn
            except Exception:
                # Log error but continue with other hosts
                continue

        if not self._connections:
            raise AllHostsFailedError()

        # Start health check thread
        if config.health_check_interval > 0:
            self._health_thread = threading.Thread(
                target=self._health_check_loop,
                daemon=True
            )
            self._health_thread.start()

    def _create_connection(self, host: str) -> Connection:
        """Create a new gRPC connection to a host.

        Args:
            host: Host address in format "host:port".

        Returns:
            New Connection object.
        """
        options = [
            ("grpc.max_receive_message_length", self._config.max_recv_size),
            ("grpc.max_send_message_length", self._config.max_recv_size),
        ]

        if self._config.ssl_context:
            # Create secure channel with SSL
            credentials = grpc.ssl_channel_credentials()
            channel = grpc.secure_channel(host, credentials, options=options)
        else:
            # Create insecure channel
            channel = grpc.insecure_channel(host, options=options)

        return Connection(
            host=host,
            channel=channel,
            healthy=True,
            last_ping=time.time()
        )

    def get_connection(self) -> grpc.Channel:
        """Get a healthy connection from the pool.

        Returns:
            A gRPC channel.

        Raises:
            ConnectionClosedError: If the pool is closed.
            NoConnectionError: If no healthy connection is available.
        """
        with self._lock:
            if self._closed:
                raise ConnectionClosedError()

            # Find a healthy connection
            for conn in self._connections.values():
                if conn.healthy:
                    return conn.channel

            raise NoConnectionError()

    def get_connection_for_host(self, host: str) -> grpc.Channel:
        """Get a connection for a specific host.

        Args:
            host: Host address.

        Returns:
            A gRPC channel.

        Raises:
            NoConnectionError: If no connection is available for the host.
        """
        with self._lock:
            conn = self._connections.get(host)
            if not conn or not conn.healthy:
                raise NoConnectionError()
            return conn.channel

    def _health_check_loop(self) -> None:
        """Periodically check the health of all connections."""
        while not self._stop_event.wait(self._config.health_check_interval):
            self._check_health()

    def _check_health(self) -> None:
        """Check the health of all connections."""
        with self._lock:
            hosts = list(self._connections.keys())

        for host in hosts:
            with self._lock:
                conn = self._connections.get(host)
                if not conn:
                    continue

            try:
                # Check channel state
                state = conn.channel._channel.check_connectivity_state(False)
                healthy = state in (
                    grpc.ChannelConnectivity.READY,
                    grpc.ChannelConnectivity.IDLE
                )

                with self._lock:
                    conn.healthy = healthy
                    if healthy:
                        conn.last_ping = time.time()

                # Try to reconnect if unhealthy
                if not healthy:
                    self._reconnect(host)

            except Exception:
                with self._lock:
                    if conn:
                        conn.healthy = False

    def _reconnect(self, host: str) -> None:
        """Attempt to reconnect to a host.

        Args:
            host: Host address.
        """
        with self._lock:
            if self._closed:
                return

            # Close existing connection
            old_conn = self._connections.pop(host, None)
            if old_conn:
                try:
                    old_conn.channel.close()
                except Exception:
                    pass

            # Create new connection
            try:
                new_conn = self._create_connection(host)
                self._connections[host] = new_conn
            except Exception:
                pass

    def close(self) -> None:
        """Close all connections in the pool."""
        with self._lock:
            if self._closed:
                return

            self._closed = True
            self._stop_event.set()

            # Close all connections
            for conn in self._connections.values():
                try:
                    conn.channel.close()
                except Exception:
                    pass

            self._connections.clear()

        # Wait for health check thread
        if self._health_thread:
            self._health_thread.join(timeout=5.0)

    def healthy_host_count(self) -> int:
        """Return the number of healthy hosts.

        Returns:
            Number of healthy hosts.
        """
        with self._lock:
            return sum(1 for conn in self._connections.values() if conn.healthy)

    def hosts(self) -> list[str]:
        """Return the list of configured hosts.

        Returns:
            List of host addresses.
        """
        with self._lock:
            return list(self._connections.keys())

    @property
    def is_closed(self) -> bool:
        """Check if the pool is closed."""
        with self._lock:
            return self._closed
