import asyncio
import asyncio.subprocess
import pathlib
import sys

from attrs import Factory, define, field

from dnstap_pb2 import Message
from rabdomancia.dnstap import FstrmDnsTapClientProtocol
from rabdomancia.utils import Tasker


@define(auto_attribs=True)
class RabdomanciaController:
    # graphic / gource settings
    bloom_multiplier: float = 2.0
    bloom_intensity: float = 0.2
    elasticity: float = 0.5
    dir_name_position: float = 0.9
    logo: pathlib.Path = pathlib.Path("logo.png")
    # sound settings
    player: str = "mpv"
    # input settings
    disable_dnstap: bool = False
    raw_input: bool = False
    domain_lists_dir: pathlib.Path = pathlib.Path("domain-lists")
    # DNSTAP settings
    host: str = ""
    port: int = 5342

    # Internal data
    domain_lists: dict[str, str] = Factory(dict)
    tld_list: set[str] = Factory(set)
    visit_counters: dict[str, int] = Factory(dict)
    # Internal workings
    initialised: asyncio.Event = Factory(asyncio.Event)
    tasker: Tasker = Factory(Tasker)
    gource: asyncio.subprocess.Process = field(init=False)
    sound: asyncio.subprocess.Process = field(init=False)

    def __attrs_post_init__(self):
        self.tasker.one_off(self.first_start())
        if not self.raw_input:
            self.load_lists()

    @property
    def lists_loaded(self) -> bool:
        return len(self.domain_lists) * len(self.tld_list) > 0

    @property
    def gource_args(self) -> list[str]:
        return (
            [
                "--bloom-multiplier",
                str(self.bloom_multiplier),
                "--bloom-intensity",
                str(self.bloom_intensity),
                "--elasticity",
                str(self.elasticity),
                "--dir-name-position",
                str(self.dir_name_position),
            ]
            + (["--logo", str(self.logo)] if self.logo.is_file() else [])
            + [
                # "--hide",
                # "bloom",
                "--log-format",
                "custom",
                "-",
            ]
        )

    @property
    def sound_args(self) -> list[str]:
        return ["--player", self.player]

    def load_lists(self):
        domain_lists_dir = self.domain_lists_dir / pathlib.Path("data")
        if not domain_lists_dir.is_dir():
            return
        for dlf in domain_lists_dir.iterdir():
            if not dlf.is_file():
                continue
            with dlf.open() as f:
                for line in f:
                    if not line.strip() or line.startswith("#"):
                        continue
                    self.domain_lists[line.strip()] = dlf.name

        tld_list = self.domain_lists_dir / pathlib.Path("public_suffix_list.dat")
        if not tld_list.is_file():
            return
        with tld_list.open() as f:
            for line in f:
                if not line.strip() or line.startswith("//"):
                    continue
                self.tld_list.add(line.strip())

    def line_to_gource(self, line: str) -> str:
        t, ip, _mod, domain = line.strip().split("|", maxsplit=4)
        ip = ip.strip()
        domain = domain.strip()
        # Normalise domain to avoid duplicates (capitalisation, etc.)
        n_domain = domain.lower()
        mode = "M" if n_domain in self.visit_counters else "A"
        # We now want to generate the path with the hierarchy we like
        domain_parts = n_domain.split(".")
        # Get public TLD
        tld_candidate = n_domain
        tld, group, two_level_domain = "", "", ""
        for i in range(len(domain_parts)):
            star_candidate = f"*.{tld_candidate}"
            if not tld:
                two_level_domain = tld_candidate
            tld_candidate = ".".join(domain_parts[i:])
            if not group:
                group = self.domain_lists.get(tld_candidate, "")
            if not tld:
                for tc in [star_candidate, tld_candidate]:
                    if tc in self.tld_list:
                        tld = tc
        # Ensure we always have a TLD
        if not tld:
            tld = domain_parts[-1]
            two_level_domain = ".".join(domain_parts[-2:])

        entry_top = "/".join(reversed(group.split("."))) if group else tld
        entry_path = f"{entry_top}/{two_level_domain}/{domain}"

        return f"{t}|{ip}|{mode}|{entry_path}\n"

    def first_start(self):
        return self._reset(first_run=True)

    def restart(self):
        return self._reset(first_run=False)

    async def _reset(self, first_run):
        self.initialised.clear()
        self.visit_counters.clear()
        if not first_run:
            for p in [self.sound, self.gource]:
                # Try to terminate process
                if p.returncode is None:
                    p.terminate()
                # Process has not exited yet
                if p.returncode is None:
                    p.kill()
        self.gource = await asyncio.create_subprocess_exec(
            "gource",
            *self.gource_args,
            # subprocess args
            stdin=asyncio.subprocess.PIPE,
            stdout=None,
            stderr=None,
            close_fds=True,
        )
        self.sound = await asyncio.create_subprocess_exec(
            "rabdomancia-sound",
            *self.sound_args,
            # subprocess args
            stdin=asyncio.subprocess.PIPE,
            stdout=None,
            stderr=None,
            close_fds=True,
        )
        self.initialised.set()

    async def ensure_process_sync(self):
        """
        If graphics or sound are restarted,
        we want to restart the other component for consistency.
        """
        for p in [self.sound, self.gource]:
            if p.returncode is not None:
                await self.restart()
                break

    async def process_line(self, line: str) -> None:
        init_wait = self.initialised.wait()
        proc_sync = self.ensure_process_sync()
        lb: bytes
        if not self.raw_input:
            # convert line to gource input
            line = self.line_to_gource(line)
        lb = line.encode("utf-8")
        # Do not continue until fully initialised and syncd
        await asyncio.gather(init_wait, proc_sync)
        if not lb:
            return
        # Finally write lines
        drains = []
        for p in [self.gource, self.sound]:
            if p.stdin is not None:
                p.stdin.write(lb)
                drains.append(p.stdin.drain())
        # Wait for streams to have flushed
        await asyncio.gather(*drains)

    def process_message(self, msg: Message, addr: str, domain: str):
        self.tasker.one_off(self.process_line(f"{msg.query_time_sec}|{addr}|M|{domain}"))

    async def run_forever(self) -> None:
        loop = asyncio.get_event_loop()
        if self.disable_dnstap:
            # stdin mode
            while True:
                line = await loop.run_in_executor(None, sys.stdin.readline)
                await self.process_line(line)
        else:
            # DNSTAP mode
            server = await loop.create_server(
                lambda: FstrmDnsTapClientProtocol(process_message_cb=lambda x, y, z: self.process_message(x, y, z)),
                self.host,
                self.port,
            )
            await server.serve_forever()
