"""Connector stuffs."""

import random
from asyncio import sleep as asyncio_sleep
from asyncio import wait_for
from typing import TYPE_CHECKING, Any, Dict, Optional, Union
from urllib.parse import ParseResult

from onecache import ExpirableCache

from aiosonic.exceptions import ConnectTimeout, HttpParsingError, TimeoutException
from aiosonic.pools import BasePool, Http2MultiplexPool, PoolConfig, SmartPool
from aiosonic.resolver import DefaultResolver
from aiosonic.timeout import Timeouts

if TYPE_CHECKING:
    from aiosonic.connection import Connection


class TCPConnector:
    """TCPConnector.

    Holds the main logic for making connections to destination hosts.

    Parameters:
        pool_configs (Optional[Dict[str, PoolConfig]]):
            Map of host domains to pool configurations. Keys are host domains
            (e.g., "https://example.com" or "example.com") and values are PoolConfig
            instances. A special key ":default" is used for hosts without a specific configuration.
        timeouts (Optional[Timeouts]):
            Global timeouts for connections. Defaults to a Timeouts instance with default args.
        connection_cls:
            Connection class to be used. Defaults to Connection.
        pool_cls:
            Pool class to be used. Defaults to SmartPool, or Http2MultiplexPool when http2=True.
        resolver:
            DNS resolver to be used. Defaults to DefaultResolver.
        ttl_dns_cache (int):
            TTL in milliseconds for DNS cache. Defaults to 10000 (10 seconds).
        use_dns_cache (bool):
            Flag to indicate usage of DNS cache. Defaults to True.
        http2 (bool):
            When True and pool_cls is not explicitly provided, defaults pool_cls to
            Http2MultiplexPool. Defaults to False.
    """

    def __init__(
        self,
        pool_configs: Optional[Dict[str, Union[PoolConfig, Dict[str, Any]]]] = None,
        timeouts: Optional[Timeouts] = None,
        connection_cls=None,
        pool_cls=None,
        resolver=None,
        ttl_dns_cache=10000,
        use_dns_cache=True,
        http2: bool = False,
    ):
        from aiosonic.connection import Connection  # avoid circular dependency

        self.connection_cls = connection_cls or Connection
        self.pool_cls = pool_cls or (Http2MultiplexPool if http2 else SmartPool)
        self.timeouts = timeouts or Timeouts()

        if pool_configs is None:
            pool_configs = {}

        if ":default" not in pool_configs:
            pool_configs[":default"] = PoolConfig()

        self.pool_configs = _check_pool_configs(pool_configs)

        # Pre-create pools based on provided pool_configs keys.
        # Keys are expected to be in the form "<scheme>://<host>" or ":default".
        self.pools: Dict[str, BasePool] = {}
        for key, config in self.pool_configs.items():
            self.pools[key] = self.pool_cls(config, self.connection_cls, self.timeouts)

        self.resolver = resolver or DefaultResolver()
        self.use_dns_cache = use_dns_cache
        if self.use_dns_cache:
            self.cache = ExpirableCache(512, ttl_dns_cache)

    async def acquire(self, urlparsed: ParseResult, verify, ssl, timeouts, http2) -> "Connection":
        """Acquire a connection from the appropriate pool."""
        if not urlparsed.hostname:
            raise HttpParsingError("missing hostname")

        host_key = f"{urlparsed.scheme}://{urlparsed.hostname}"

        # Use host-specific pool if available; otherwise, fall back to the default pool.
        if host_key in self.pools:
            pool = self.pools[host_key]
        else:
            pool = self.pools[":default"]

        conn = await pool.acquire(urlparsed)
        return await self.after_acquire(urlparsed, conn, verify, ssl, timeouts, http2)

    async def after_acquire(self, urlparsed, conn, verify, ssl, timeouts, http2):
        """Process connection after acquisition."""
        try:
            dns_info = await self.__resolve_dns(urlparsed.hostname, urlparsed.port)
            await wait_for(
                conn.connect(urlparsed, dns_info, verify, ssl, http2),
                timeout=timeouts.sock_connect,
            )
        except TimeoutException:
            self.release(conn)
            raise ConnectTimeout()
        except BaseException as ex:
            self.release(conn)
            raise ex
        return conn

    def release(self, conn):
        """Release connection back to its pool."""
        conn.release()

    async def wait_free_pool(self):
        """Wait until all pools are free."""
        while True:
            if all(pool.is_all_free() for pool in self.pools.values()):
                return True
            await asyncio_sleep(0.02)  # pragma: no cover

    async def cleanup(self):
        """Cleanup connector connections."""
        for pool in self.pools.values():
            await pool.cleanup()

    async def __resolve_dns(self, host: str, port: int):
        key = f"{host}-{port}"
        dns_data = self.cache.get(key)
        if not dns_data:
            dns_data = await self.resolver.resolve(host, port)
            self.cache.set(key, dns_data)
        assert isinstance(dns_data, list)
        return random.choice(dns_data)


def _check_pool_configs(configs) -> Dict[str, PoolConfig]:
    result = {}  # temporary variable to store results
    for key, val in configs.items():
        if isinstance(val, PoolConfig):
            result[key] = val
        else:
            result[key] = PoolConfig(**val)
    return result
