from __future__ import annotations

import enum
import ipaddress
import socket
from dataclasses import dataclass

from .._helpers import is_ipv4_address
from .errors import ReplyError

RSV = NULL = 0x00
SOCKS_VER = 0x04


class Command(enum.IntEnum):
    CONNECT = 0x01
    BIND = 0x02


class ReplyCode(enum.IntEnum):
    REQUEST_GRANTED = 0x5A
    REQUEST_REJECTED_OR_FAILED = 0x5B
    CONNECTION_FAILED = 0x5C
    AUTHENTICATION_FAILED = 0x5D


ReplyMessages = {
    ReplyCode.REQUEST_GRANTED: "Request granted",
    ReplyCode.REQUEST_REJECTED_OR_FAILED: "Request rejected or failed",
    ReplyCode.CONNECTION_FAILED: "Request rejected because SOCKS server cannot connect to identd on the client",  # noqa: E501
    ReplyCode.AUTHENTICATION_FAILED: "Request rejected because the client program and identd report different user-ids",  # noqa: E501,
}


@dataclass
class ConnectRequest:
    host: str  # hostname or IPv4 address
    port: int
    user_id: str | None

    def dumps(self) -> bytes:
        port_bytes = self.port.to_bytes(2, "big")
        include_hostname = False

        if is_ipv4_address(self.host):
            host_bytes = ipaddress.IPv4Address(self.host).packed
        else:
            include_hostname = True
            host_bytes = bytes([NULL, NULL, NULL, 1])

        data = bytearray([SOCKS_VER, Command.CONNECT])
        data += port_bytes
        data += host_bytes

        if self.user_id:
            data += self.user_id.encode("ascii")

        data.append(NULL)

        if include_hostname:
            data += self.host.encode("idna")
            data.append(NULL)

        return bytes(data)


@dataclass
class ConnectReply:
    SIZE = 8

    rsv: int
    reply: ReplyCode
    host: str  # should be ignored when using Command.CONNECT
    port: int  # should be ignored when using Command.CONNECT

    @classmethod
    def loads(cls, data: bytes) -> ConnectReply:
        if len(data) != cls.SIZE:
            raise ReplyError("Malformed connect reply")

        rsv = data[0]
        if rsv != RSV:  # pragma: no cover
            raise ReplyError(f"Unexpected reply version: {data[0]:#02X}")

        try:
            reply = ReplyCode(data[1])
        except ValueError as e:
            raise ReplyError(f"Invalid reply code: {data[1]:#02X}") from e

        if reply != ReplyCode.REQUEST_GRANTED:  # pragma: no cover
            msg = ReplyMessages.get(reply, "Unknown error")
            raise ReplyError(msg, error_code=reply)

        try:
            port = int.from_bytes(data[2:4], byteorder="big")
        except ValueError as e:
            raise ReplyError("Invalid port data") from e

        try:
            host = socket.inet_ntop(socket.AF_INET, data[4:8])
        except ValueError as e:
            raise ReplyError("Invalid port data") from e

        return cls(rsv=rsv, reply=reply, host=host, port=port)


class Connection:
    def send(self, request: ConnectRequest) -> bytes:
        return request.dumps()

    def receive(self, data: bytes) -> ConnectReply:
        return ConnectReply.loads(data)