import asyncio
from socket import (
    socket,
    IPPROTO_ICMP,
    IPPROTO_IP,
    SOCK_RAW,
    AF_INET,
    IP_TTL,
    SOL_SOCKET,
)
import struct
from functools import partial
from random import randint
from time import perf_counter
from dataclasses import dataclass

SO_BINDTODEVICE = 25

# https://datatracker.ietf.org/doc/html/rfc792
ICMP_ECHO_REPLY = 0
ICMP_DEST_UNREACHABLE = 3
ICMP_ECHO_REQUEST = 8
ICMP_TTL_EXCEEDED = 11

ICMP_RESPONSES = [ICMP_ECHO_REPLY, ICMP_DEST_UNREACHABLE, ICMP_TTL_EXCEEDED]


@dataclass
class IcmpPingResponse:
    rtt: float | None = None
    message: str | None = None

    @property
    def rtt_ms(self) -> str:
        return f"{round(self.rtt, 3):.3f} ms" if self.rtt else ""

    def __bool__(self) -> bool:
        return bool(self.rtt and not self.message)


class IcmpPacketSender:
    def __init__(self) -> None:
        self.loop = asyncio.get_running_loop()

    @staticmethod
    def calculate_checksum(data: bytes) -> int:
        checksum = 0

        if len(data) % 2:
            data += b"\x00"

        for i in range(0, len(data), 2):
            checksum += (data[i] << 8) + data[i + 1]

        checksum = (checksum >> 16) + (checksum & 0xFFFF)
        checksum += checksum >> 16

        return (~checksum) & 0xFFFF

    @staticmethod
    def generate_header(id: int, *, checksum: int | None = None) -> bytes:
        icmp_type = ICMP_ECHO_REQUEST
        icmp_code = 0
        icmp_checksum = checksum or 0
        icmp_sequence = 1

        icmp_header = struct.pack(
            "!BBHHH",
            icmp_type,
            icmp_code,
            icmp_checksum,
            id,
            icmp_sequence,
        )

        return icmp_header

    async def receive_ping(
        self, raw_socket: socket, id: int, timeout: float
    ) -> IcmpPingResponse:
        try:
            while True:
                received = await asyncio.wait_for(
                    self.loop.sock_recv(raw_socket, 1500), timeout
                )

                time_received = perf_counter()

                offset = 20
                icmp_header_size = 8

                icmp_header = received[offset : offset + icmp_header_size]

                # On ttl and dest unreachable, only the type, code, and checksum
                # will be present. The packet_id and sequence are zero. Then
                # follows the ip header, then the original header + data

                try:
                    type, code, checksum, packet_id, sequence = struct.unpack(
                        "!BBHHH", icmp_header
                    )
                except (struct.error, TypeError):
                    continue

                msg = None
                rtt = None

                if type == ICMP_ECHO_REPLY:
                    if packet_id != id:
                        continue

                    data = received[
                        offset + 8 : offset + 8 + struct.calcsize("d")
                    ]
                    time_sent = struct.unpack("d", data)[0]

                    rtt = (time_received - time_sent) * 1000
                elif type == ICMP_TTL_EXCEEDED:
                    msg = "TTL Exceeded"
                elif type == ICMP_DEST_UNREACHABLE:
                    msg = "Dest Unreachable"

                return IcmpPingResponse(rtt, msg)

        except asyncio.TimeoutError:
            self.loop.remove_writer(raw_socket)
            self.loop.remove_reader(raw_socket)
            raw_socket.close()

            raise TimeoutError("Ping timeout")

    def send_packet(
        self,
        packet: bytes,
        socket: socket,
        future: asyncio.Future,
        dest: tuple[str, int],
    ) -> None:
        try:
            socket.sendto(packet, dest)
        except (BlockingIOError, InterruptedError):
            return  # The callback will be retried
        except Exception as e:
            self.loop.remove_writer(socket)
            future.set_exception(e)
        else:
            self.loop.remove_writer(socket)
            future.set_result(None)

    async def send_ping(
        self, raw_socket: socket, dest_addr: tuple[str, int], id: int
    ) -> bool:
        dummy_header = self.generate_header(id)

        fmt = "d"
        rtt_size = struct.calcsize(fmt)
        data = (192 - rtt_size) * "0"

        payload = struct.pack(fmt, perf_counter()) + data.encode("utf-8")
        checksum = self.calculate_checksum(dummy_header + payload)

        header = self.generate_header(id, checksum=checksum)

        future: asyncio.Future[float] = asyncio.Future()

        callback = partial(
            self.send_packet,
            packet=header + payload,
            socket=raw_socket,
            dest=dest_addr,
            future=future,
        )
        self.loop.add_writer(raw_socket, callback)

        try:
            await future
        except OSError:
            return False

        return True

    async def ping(
        self,
        dest_addr: str,
        *,
        timeout: float = 5.0,
        ttl: int = 64,
        interface: str | None = None,
    ) -> IcmpPingResponse:
        addr = (dest_addr, 0)

        # info = await self.loop.getaddrinfo(addr)

        # let this raise
        sock = socket(AF_INET, SOCK_RAW, IPPROTO_ICMP)

        # if interface:
        #     sock.bind((interface_ip, 0))

        if interface:
            # linux only
            encoded = str(interface + "\0").encode("utf-8")
            sock.setsockopt(SOL_SOCKET, SO_BINDTODEVICE, encoded)

        sock.setblocking(False)
        sock.setsockopt(IPPROTO_IP, IP_TTL, struct.pack("I", ttl))

        icmp_id = randint(1, 65535)

        ok = await self.send_ping(sock, addr, icmp_id)

        if not ok:
            return IcmpPingResponse(message="OSError")

        try:
            response = await self.receive_ping(sock, icmp_id, timeout)
        except asyncio.TimeoutError:
            response = IcmpPingResponse(message="Timeout exceeded")
        finally:
            sock.close()

        return response


if __name__ == "__main__":

    async def main():
        icmp = IcmpPacketSender()

        while True:
            res = await icmp.ping("8.8.8.8", ttl=14)

            if res:
                print(res.rtt_ms)
            else:
                print(res.message)

            await asyncio.sleep(1)

    asyncio.run(main())
