"""Test the `tinydns` and `dnsq` programs."""

import dataclasses
import enum
import os
import pathlib
import socket
import struct
import subprocess
import sys
import time

from typing import (
    Callable,  # noqa: H301
    Dict,
    List,
    Tuple,
    Union,
    TYPE_CHECKING,
)

from . import defs
from . import test_get


if TYPE_CHECKING:
    PopenStr = subprocess.Popen[str]  # pylint: disable=unsubscriptable-object
else:
    PopenStr = subprocess.Popen


class ProcessFDType(str, enum.Enum):
    """File descriptor types, enough for this very limited test."""

    IPV4 = "IPv4"


@dataclasses.dataclass(frozen=True)
class ProcessFDInfo:
    """Information about a single file descriptor in a process."""

    fileno: int
    access: str
    lock: str
    ftype: ProcessFDType


@dataclasses.dataclass(frozen=True)
class ProcessSocketFDInfo(ProcessFDInfo):
    """Information about a single socket in a process."""

    proto: str
    address: str
    port: int


@dataclasses.dataclass(frozen=True)
class ProcessInfo:
    """Information about a process and (some of) its file descriptors."""

    pid: int
    pgid: int
    tid: int
    comm: str

    uid: int
    username: str

    files: List[ProcessFDInfo]


@dataclasses.dataclass
class ParseProcessInfo:
    """Parse the output of `lsof -F` for our very limited testing needs."""

    cfg: defs.Config
    cproc: Dict[str, str] = dataclasses.field(default_factory=dict)
    cfile: Dict[str, str] = dataclasses.field(default_factory=dict)
    cfiles: List[ProcessFDInfo] = dataclasses.field(default_factory=list)
    result: Dict[int, ProcessInfo] = dataclasses.field(default_factory=dict)

    def finish_file_ipv4(self) -> ProcessFDInfo:
        """Store the information about a socket."""
        fields = self.cfile["addrport"].split(":", 1)
        return ProcessSocketFDInfo(
            fileno=int(self.cfile["fileno"]),
            access=self.cfile["access"],
            lock=self.cfile["lock"],
            ftype=ProcessFDType.IPV4,
            proto=self.cfile["proto"],
            address=fields[0],
            port=int(fields[1]),
        )

    def finish_file(self) -> None:
        """Store the information about the current file descriptor."""
        if self.cfile["type"] == "IPv4":
            fdata = self.finish_file_ipv4()
        else:
            sys.exit(f"Internal error: unhandled file type: {self.cfile!r}")

        self.cfile.clear()

        self.cfiles.append(fdata)

    def finish_process(self) -> None:
        """Store the information about the current process."""
        if self.cfile:
            self.finish_file()

        pdata = ProcessInfo(
            pid=int(self.cproc["pid"]),
            pgid=int(self.cproc["pgid"]),
            tid=int(self.cproc["tid"]),
            comm=self.cproc["comm"],
            uid=int(self.cproc["uid"]),
            username=self.cproc["username"],
            files=list(self.cfiles),
        )

        self.cfiles.clear()
        self.cproc.clear()

        if pdata.pid in self.result:
            sys.exit(
                f"lsof output a duplicate process set for {pdata.pid}: "
                f"first {self.result[pdata.pid]!r} now {pdata!r}"
            )
        self.result[pdata.pid] = pdata

    def parse(self, lines: List[str]) -> Dict[int, ProcessInfo]:
        """Parse some lsof output lines."""

        def h_process(_first: str, rest: str) -> None:
            """A new process set."""
            if self.cproc:
                self.finish_process()

            self.cproc["pid"] = rest

        def h_file(_first: str, rest: str) -> None:
            """A new file descriptor set."""
            if self.cfile:
                self.finish_file()

            self.cfile["fileno"] = rest

        handlers: Dict[
            str, Union[Callable[[str, str], None], Tuple[Dict[str, str], str]]
        ] = {
            "p": h_process,
            "g": (self.cproc, "pgid"),
            "R": (self.cproc, "tid"),
            "c": (self.cproc, "comm"),
            "u": (self.cproc, "uid"),
            "L": (self.cproc, "username"),
            "f": h_file,
            "a": (self.cfile, "access"),
            "l": (self.cfile, "lock"),
            "t": (self.cfile, "type"),
            "P": (self.cfile, "proto"),
            "n": (self.cfile, "addrport"),
        }

        for line in lines:
            if not line:
                sys.exit("lsof output an empty line")
            first, rest = line[0], line[1:]
            handler = handlers.get(first)
            if not handler:
                continue

            if isinstance(handler, tuple):
                handler[0][handler[1]] = rest
            else:
                handler(first, rest)

        if self.cproc:
            self.finish_process()

        return self.result


def check_single_udp_socket(
    processes: Dict[int, ProcessInfo], pid: int, ipaddr: str
) -> bool:
    """Make sure tinydns is only listening on a single UDP socket."""
    pdata = processes.get(pid)
    if pdata is None or len(pdata.files) != 1:
        return False

    data = pdata.files[0]
    if not isinstance(data, ProcessSocketFDInfo):
        return False

    return data.proto == "UDP" and data.address == ipaddr and data.port == 53


def run_tinydns(
    cfg: defs.Config,
    tempd: pathlib.Path,
    callback: Callable[[PopenStr, str], None],
) -> None:
    """Test the `tinydns` and `dnsq` programs."""
    ipaddr = defs.RECORDS[0].address
    svcdir = tempd / defs.SVCDIR
    if not svcdir.is_dir():
        sys.exit(f"Expected {svcdir} to be a directory")
    os.chdir(svcdir)

    with subprocess.Popen(
        ["./run"],
        shell=False,
        env=cfg.subenv,
        bufsize=0,
        encoding=defs.MINENC,
        stdout=subprocess.PIPE,
    ) as proc:
        try:
            print(f"- spawned process {proc.pid}")
            assert proc.stdout is not None
            print("- waiting for the 'starting tinydns' line")
            line = proc.stdout.readline()
            print(f"- got line {line!r}")
            if line != "starting tinydns\n":
                sys.exit(f"Unexpected first line from tinydns: {line!r}")

            lines = cfg.check_output(
                [
                    "lsof",
                    "-a",
                    "-n",
                    "-P",
                    "-p",
                    str(proc.pid),
                    "-F",
                    "-i4udp:53",
                ]
            ).splitlines()
            cfg.diag(f"lsof output: {lines!r}")
            processes = ParseProcessInfo(cfg).parse(lines)
            if not check_single_udp_socket(processes, proc.pid, ipaddr):
                sys.exit(
                    f"tinydns should listen on a single UDP socket: "
                    f"{processes!r}"
                )
            print(f"- tinydns is listening at {ipaddr}:53")
            callback(proc, ipaddr)
        finally:
            print("- we spawned a process, checking if it has exited")
            res = proc.poll()
            if res is None:
                print(f"Terminating process {proc.pid}")
                proc.terminate()
                time.sleep(0.5)
                res = proc.poll()
                if res is None:
                    print(f"Killing process {proc.pid}")
                    proc.kill()

    res = proc.wait()
    print(f"Process {proc.pid}: exit code {res}")


def test_tinydns_run(cfg: defs.Config, tempd: pathlib.Path) -> None:
    """Test the `tinydns` and `dnsq` programs."""
    print("\n==== test_tinydns_run")
    if os.geteuid() != 0:
        print("- not running as root, skipped")
        return

    dnsq = cfg.bindir / "dnsq"
    if not dnsq.is_file() or not os.access(dnsq, os.X_OK):
        sys.exit(f"Not an executable file: {dnsq}")

    def test_dnsq(proc: PopenStr, ipaddr: str) -> None:
        """Run `dnsq` against the specified `tinydns` instance."""
        assert proc.stdout is not None

        for rec in defs.RECORDS:
            rdef = defs.TYPES[rec.rtype]
            print(f"- querying {ipaddr} for {rdef.query} {rec.name}")
            test_get.test_tinydns_query(
                cfg,
                [
                    "timelimit",
                    "-p",
                    "-t1",
                    "-T1",
                    "--",
                    dnsq,
                    rdef.query,
                    rec.name,
                    ipaddr,
                ],
                rec,
            )

            print("  - there should be a single line of output from tinydns")
            line = proc.stdout.readline()
            cfg.diag(f"line: {line!r}")

    run_tinydns(cfg, tempd, test_dnsq)


def test_tinydns_run_udp(cfg: defs.Config, tempd: pathlib.Path) -> None:
    """Test the `tinydns` program by sending UDP packets."""
    print("\n==== test_tinydns_run_udp")
    if os.geteuid() != 0:
        print("- not running as root, skipped")
        return

    def test_dnsq(proc: PopenStr, ipaddr: str) -> None:
        """Run `dnsq` against the specified `tinydns` instance."""
        assert proc.stdout is not None

        pkt = struct.pack(
            ">HHHHHH8p4p1pHH",
            17,
            0,
            1,
            0,
            0,
            0,
            "example".encode("us-ascii"),
            "com".encode("us-ascii"),
            b"",
            1,
            1,
        )
        print(f"- query packet: {pkt!r}")

        sock = socket.socket(
            socket.AF_INET, socket.SOCK_DGRAM, socket.IPPROTO_UDP
        )
        sock.connect((ipaddr, 53))

        for idx in range(1000000):
            pkt = struct.pack(">H", (idx + 17) % 65536) + pkt[2:]
            sent = sock.send(pkt)
            if sent != len(pkt):
                sys.exit(f"Only sent {sent} of {len(pkt)} bytes")

            line = proc.stdout.readline()
            if "+" not in line:
                sys.exit(
                    f"Query {idx + 1}: expected a '+' in "
                    f"the tinydns log line {line!r}"
                )

            resp = sock.recv(4096)
            if resp[:2] != pkt[:2]:
                sys.exit(
                    f"Query {idx + 1}: unexpected ID: "
                    f"query {pkt!r} response {resp!r}"
                )
            if resp[4:8] != b"\x00\x01\x00\x01":
                sys.exit(
                    f"Query {idx + 1}: unexpected query/response count: "
                    f"query {pkt!r} response {resp!r}"
                )

            if (idx % 10000) == 9999:
                print(f"- sent {idx + 1} queries to {ipaddr}")

    run_tinydns(cfg, tempd, test_dnsq)
