import asyncio
import struct
from abc import ABC, abstractmethod
from socket import inet_aton

from .errors import BadResponseError, BadStatusError
from .utils import get_headers, get_status_code

__all__ = [
    "Socks5Ngtr",
    "Socks4Ngtr",
    "Connect80Ngtr",
    "Connect25Ngtr",
    "HttpsNgtr",
    "HttpNgtr",
    "NGTRS",
]


SMTP_READY = 220


def _CONNECT_request(host, port, **kwargs):
    kwargs.setdefault("User-Agent", get_headers()["User-Agent"])
    kw = {
        "host": host,
        "port": port,
        "headers": "\r\n".join((f"{k}: {v}" for k, v in kwargs.items())),
    }
    req = (
        (
            "CONNECT {host}:{port} HTTP/1.1\r\nHost: {host}\r\n"
            "{headers}\r\nConnection: keep-alive\r\n\r\n"
        )
        .format(**kw)
        .encode()
    )
    return req


class BaseNegotiator(ABC):
    """Base Negotiator."""

    name = None
    check_anon_lvl = False
    use_full_path = False

    def __init__(self, proxy):
        self._proxy = proxy

    @abstractmethod
    async def negotiate(self, **kwargs):
        """Negotiate with proxy."""


class Socks5Ngtr(BaseNegotiator):
    """SOCKS5 Negotiator."""

    name = "SOCKS5"

    async def negotiate(self, **kwargs):
        await self._proxy.send(struct.pack("3B", 5, 1, 0))
        resp = await self._proxy.recv(2)

        if not isinstance(resp, (bytes, str)):
            raise TypeError(f"{type(resp).__name__} is not supported")
        if resp[0] == 0x05 and resp[1] == 0xFF:
            self._proxy.log("Failed (auth is required)", err=BadResponseError)
            raise BadResponseError
        elif resp[0] != 0x05 or resp[1] != 0x00:
            self._proxy.log("Failed (invalid data)", err=BadResponseError)
            raise BadResponseError

        bip = inet_aton(kwargs.get("ip"))
        port = kwargs.get("port", 80)

        await self._proxy.send(struct.pack(">8BH", 5, 1, 0, 1, *bip, port))
        resp = await self._proxy.recv(10)

        if resp[0] != 0x05 or resp[1] != 0x00:
            self._proxy.log("Failed (invalid data)", err=BadResponseError)
            raise BadResponseError
        else:
            self._proxy.log("Request is granted")


class Socks4Ngtr(BaseNegotiator):
    """SOCKS4 Negotiator."""

    name = "SOCKS4"

    async def negotiate(self, **kwargs):
        bip = inet_aton(kwargs.get("ip"))
        port = kwargs.get("port", 80)

        await self._proxy.send(struct.pack(">2BH5B", 4, 1, port, *bip, 0))
        resp = await self._proxy.recv(8)
        if isinstance(resp, asyncio.Future):
            resp = await resp
        assert not isinstance(resp, asyncio.Future)

        if resp[0] != 0x00 or resp[1] != 0x5A:
            self._proxy.log("Failed (invalid data)", err=BadResponseError)
            raise BadResponseError
        # resp = b'\x00Z\x00\x00\x00\x00\x00\x00' // ord('Z') == 90 == 0x5A
        else:
            self._proxy.log("Request is granted")


class Connect80Ngtr(BaseNegotiator):
    """CONNECT Negotiator."""

    name = "CONNECT:80"

    async def negotiate(self, **kwargs):
        await self._proxy.send(_CONNECT_request(kwargs.get("host"), 80))
        resp = await self._proxy.recv(head_only=True)
        code = get_status_code(resp)
        if code != 200:
            self._proxy.log(f"Connect: failed. HTTP status: {code}", err=BadStatusError)
            raise BadStatusError


class Connect25Ngtr(BaseNegotiator):
    """SMTP Negotiator (connect to 25 port)."""

    name = "CONNECT:25"

    async def negotiate(self, **kwargs):
        await self._proxy.send(_CONNECT_request(kwargs.get("host"), 25))
        resp = await self._proxy.recv(head_only=True)
        code = get_status_code(resp)
        if code != 200:
            self._proxy.log(f"Connect: failed. HTTP status: {code}", err=BadStatusError)
            raise BadStatusError

        resp = await self._proxy.recv(length=3)
        code = get_status_code(resp, start=0, stop=3)
        if code != SMTP_READY:
            self._proxy.log(f"Failed (invalid data): {code}", err=BadStatusError)
            raise BadStatusError


class HttpsNgtr(BaseNegotiator):
    """HTTPS Negotiator (CONNECT + SSL)."""

    name = "HTTPS"

    async def negotiate(self, **kwargs):
        await self._proxy.send(_CONNECT_request(kwargs.get("host"), 443))
        resp = await self._proxy.recv(head_only=True)
        code = get_status_code(resp)
        if code != 200:
            self._proxy.log(f"Connect: failed. HTTP status: {code}", err=BadStatusError)
            raise BadStatusError
        await self._proxy.connect(ssl=True)


class HttpNgtr(BaseNegotiator):
    """HTTP Negotiator."""

    name = "HTTP"
    check_anon_lvl = True
    use_full_path = True

    async def negotiate(self, **kwargs):
        pass


NGTRS = {
    "HTTP": HttpNgtr,
    "HTTPS": HttpsNgtr,
    "SOCKS4": Socks4Ngtr,
    "SOCKS5": Socks5Ngtr,
    "CONNECT:80": Connect80Ngtr,
    "CONNECT:25": Connect25Ngtr,
}
