import argparse
import asyncio
import ipaddress
import sys
from collections.abc import Callable
from typing import cast

import fstrm  # type: ignore
from attrs import Factory, define, field

from dnstap_pb2 import INET6, Dnstap, Message


@define
class FstrmDnsTapClientProtocol(asyncio.Protocol):
    content_type: bytes = b"protobuf:dnstap.Dnstap"
    stream: fstrm.FstrmCodec = Factory(fstrm.FstrmCodec)
    transport: asyncio.Transport = field(init=False)
    process_message_cb: Callable[[Message, str, str], None] | None = None

    def connection_made(self, transport: asyncio.BaseTransport):
        self.transport = cast("asyncio.Transport", transport)

    def data_received(self, data: bytes):
        self.stream.append(data=data)
        # Process the buffer if we have a complete frame
        while self.stream.process():
            ctrl, _ct, payload = self.stream.decode()
            if ctrl == fstrm.FSTRM_CONTROL_READY:
                self.transport.write(self.stream.encode_ctrlaccept(self.content_type))
            elif ctrl == fstrm.FSTRM_DATA_FRAME:
                dnstap = Dnstap()
                dnstap.ParseFromString(payload)
                msg = dnstap.message
                if msg.type == Message.CLIENT_QUERY:
                    self.preprocess_message(msg)
            elif ctrl == fstrm.FSTRM_CONTROL_STOP:
                self.stream.reset()
                self.transport.write(self.stream.encode(ctrl=fstrm.FSTRM_CONTROL_FINISH))

    def decode_domain(self, query_message: bytes) -> str:
        buf = query_message[12:]  # 12 is DNS length
        parts: list[bytes] = []
        while len(buf):
            length = buf[0]
            # break on end
            if length == 0x00:
                break
            # save part
            parts.append(buf[1 : length + 1])
            # advance
            buf = buf[length + 1 :]
        return b".".join(parts).decode("utf-8")

    def preprocess_message(self, msg: Message):
        addr_int = int.from_bytes(msg.query_address)
        if msg.socket_family == INET6:
            addr = str(ipaddress.IPv6Address(addr_int))
        else:
            addr = str(ipaddress.IPv4Address(addr_int))
        domain = self.decode_domain(msg.query_message)
        if self.process_message_cb:
            self.process_message_cb(msg, addr, domain)


def post_process_message(msg: Message, addr: str, domain: str):
    tld = domain.split(".")[-1]
    # Emulate basic behaviour from dowse-to-source
    sys.stdout.write(f"{msg.query_time_sec}|{addr}|M|{tld}/{domain}\n")


async def main(host: str | None, port: int):
    """
    Simple implementation of a DNSTAP receiver that produces gource input
    """
    loop = asyncio.get_event_loop()
    server = await loop.create_server(
        lambda: FstrmDnsTapClientProtocol(process_message_cb=post_process_message), host, port
    )
    await server.serve_forever()


def run():
    parser = argparse.ArgumentParser(description=main.__doc__)
    parser.add_argument("--host", default="")
    parser.add_argument("--port", default=5342, type=int)

    ns = parser.parse_args()

    asyncio.run(main(**vars(ns)))


if __name__ == "__main__":
    run()
