from __future__ import annotations

import asyncio
import socket
from collections.abc import Iterable
from ssl import SSLContext
from typing import TYPE_CHECKING, Any, NamedTuple

from aiohttp import ClientConnectorError, TCPConnector
from aiohttp.abc import AbstractResolver, ResolveResult
from aiohttp.client_proto import ResponseHandler

if TYPE_CHECKING:  # pragma: no cover
    from aiohttp import AddrInfoType, ClientRequest, ClientTimeout

import python_socks
from python_socks import ProxyType, parse_proxy_url
from python_socks.async_.asyncio.v2 import Proxy

from ._errors import ProxyConnectionError, ProxyError, ProxyTimeoutError


class NoResolver(AbstractResolver):
    async def resolve(
        self,
        host: str,
        port: int = 0,
        family: socket.AddressFamily = socket.AF_INET,
    ) -> list[ResolveResult]:
        return [
            {
                "hostname": host,
                "host": host,
                "port": port,
                "family": family,
                "proto": 0,
                "flags": 0,
            }
        ]

    async def close(self) -> None:
        pass  # pragma: no cover


class _ResponseHandler(ResponseHandler):
    """
    To fix issue https://github.com/romis2012/aiohttp-socks/issues/27
    In Python>=3.11.5 we need to keep a reference to the StreamWriter
    so that the underlying transport is not closed during garbage collection.
    See StreamWriter.__del__ method (was added in Python 3.11.5)
    """

    def __init__(
        self, loop: asyncio.AbstractEventLoop, writer: asyncio.StreamWriter
    ) -> None:
        super().__init__(loop)
        self._writer = writer


class _BaseProxyConnector(TCPConnector):
    async def _wrap_create_connection(
        self,
        *args: Any,  # noqa: ARG002
        addr_infos: list[AddrInfoType],
        req: ClientRequest,  # noqa: ARG002
        timeout: ClientTimeout,
        client_error: type[Exception] = ClientConnectorError,  # noqa: ARG002
        **kwargs: Any,
    ) -> tuple[asyncio.Transport, ResponseHandler]:
        try:
            host: str = addr_infos[0][4][0]
            port: int = addr_infos[0][4][1]
        except IndexError as e:  # pragma: no cover
            raise ValueError("Invalid arg: `addr_infos`") from e

        ssl: SSLContext | None = kwargs.get("ssl")
        try:
            return await self._connect_via_proxy(
                host=host,
                port=port,
                ssl=ssl,
                timeout=timeout.sock_connect,
            )
        except python_socks.ProxyConnectionError as e:
            raise ProxyConnectionError(str(e)) from e
        except python_socks.ProxyTimeoutError as e:
            raise ProxyTimeoutError(str(e)) from e
        except python_socks.ProxyError as e:
            raise ProxyError(str(e), error_code=e.error_code) from e

    async def _connect_via_proxy(
        self,
        host: str,
        port: int,
        ssl: SSLContext | None = None,
        timeout: float | None = None,
    ) -> tuple[asyncio.Transport, ResponseHandler]:
        raise NotImplementedError


class ProxyConnector(_BaseProxyConnector):
    def __init__(
        self,
        host: str,
        port: int,
        proxy_type: ProxyType = ProxyType.SOCKS5,
        username: str | None = None,
        password: str | None = None,
        rdns: bool | None = None,  # noqa: FBT001
        proxy_ssl: SSLContext | None = None,
        **kwargs: Any,
    ) -> None:
        kwargs["resolver"] = NoResolver()
        super().__init__(**kwargs)

        self._proxy_type = proxy_type
        self._proxy_host = host
        self._proxy_port = port
        self._proxy_username = username
        self._proxy_password = password
        self._rdns = rdns
        self._proxy_ssl = proxy_ssl

    async def _connect_via_proxy(
        self,
        host: str,
        port: int,
        ssl: SSLContext | None = None,
        timeout: float | None = None,
    ) -> tuple[asyncio.Transport, ResponseHandler]:
        proxy = Proxy(
            proxy_type=self._proxy_type,
            host=self._proxy_host,
            port=self._proxy_port,
            username=self._proxy_username,
            password=self._proxy_password,
            rdns=self._rdns,
            proxy_ssl=self._proxy_ssl,
        )

        stream = await proxy.connect(
            dest_host=host,
            dest_port=port,
            dest_ssl=ssl,
            timeout=timeout,
        )

        transport = stream.writer.transport
        protocol: ResponseHandler = _ResponseHandler(
            loop=self._loop,
            writer=stream.writer,
        )

        transport.set_protocol(protocol)
        protocol.connection_made(transport)

        return transport, protocol  # type:ignore[return-value]

    @classmethod
    def from_url(cls, url: str, **kwargs: Any) -> ProxyConnector:
        proxy_type, host, port, username, password = parse_proxy_url(url)
        return cls(
            proxy_type=proxy_type,
            host=host,
            port=port,
            username=username,
            password=password,
            **kwargs,
        )


class ProxyInfo(NamedTuple):
    proxy_type: ProxyType
    host: str
    port: int
    username: str | None = None
    password: str | None = None
    rdns: bool | None = None


class ChainProxyConnector(_BaseProxyConnector):
    def __init__(self, proxy_infos: Iterable[ProxyInfo], **kwargs: Any) -> None:
        kwargs["resolver"] = NoResolver()
        super().__init__(**kwargs)

        self._proxy_infos = proxy_infos

    async def _connect_via_proxy(
        self,
        host: str,
        port: int,
        ssl: SSLContext | None = None,
        timeout: float | None = None,
    ) -> tuple[asyncio.Transport, ResponseHandler]:
        forward = None
        proxy = None
        for info in self._proxy_infos:
            proxy = Proxy(
                proxy_type=info.proxy_type,
                host=info.host,
                port=info.port,
                username=info.username,
                password=info.password,
                rdns=info.rdns,
                forward=forward,
            )
            forward = proxy

        assert proxy is not None

        stream = await proxy.connect(
            dest_host=host,
            dest_port=port,
            dest_ssl=ssl,
            timeout=timeout,
        )

        transport = stream.writer.transport
        protocol: ResponseHandler = _ResponseHandler(
            loop=self._loop,
            writer=stream.writer,
        )

        transport.set_protocol(protocol)
        protocol.connection_made(transport)

        return transport, protocol  # type:ignore[return-value]

    @classmethod
    def from_urls(cls, urls: Iterable[str], **kwargs: Any) -> ChainProxyConnector:
        infos = []
        for url in urls:
            proxy_type, host, port, username, password = parse_proxy_url(url)
            proxy_info = ProxyInfo(
                proxy_type=proxy_type,
                host=host,
                port=port,
                username=username,
                password=password,
            )
            infos.append(proxy_info)

        return cls(infos, **kwargs)