"""SniTun worker for traffics."""

from __future__ import annotations

import asyncio
import contextlib
import logging
from multiprocessing import Manager, Process
from socket import socket
from threading import Thread
import time
from typing import TYPE_CHECKING

from ..metrics import MetricsCollector, MetricsFactory, create_noop_metrics_collector
from .listener_peer import PeerListener
from .listener_sni import SNIProxy
from .peer import Peer
from .peer_manager import PeerManager, PeerManagerEvent

_LOGGER = logging.getLogger(__name__)

if TYPE_CHECKING:
    from multiprocessing.managers import SyncManager


class ServerWorker(Process):
    """Worker for multiplexer."""

    def __init__(
        self,
        fernet_keys: list[str],
        throttling: int | None = None,
        metrics_factory: MetricsFactory | None = None,
        metrics_interval: int = 60,
    ) -> None:
        """Initialize worker & communication."""
        super().__init__()

        self._fernet_keys: list[str] = fernet_keys
        self._throttling: int | None = throttling
        self._metrics_factory = metrics_factory or create_noop_metrics_collector
        self._metrics_interval = metrics_interval

        # Used on the child
        self._peers: PeerManager | None = None
        self._list_sni: SNIProxy | None = None
        self._list_peer: PeerListener | None = None
        self._loop: asyncio.AbstractEventLoop | None = None
        self._metrics: MetricsCollector | None = None
        self._metrics_task: asyncio.Task | None = None

        # Communication between Parent/Child
        manager: SyncManager = Manager()
        self._new = manager.Queue()
        self._sync = manager.dict()
        self._peer_count = manager.Value("i", 0)

    @property
    def peer_size(self) -> int:
        """Return amount of managed peers."""
        return self._peer_count.value

    def is_responsible_peer(self, sni: str) -> bool:
        """Return True if worker is responsible for this peer domain."""
        return sni in self._sync

    async def _async_init(self) -> None:
        """Initialize child process data."""
        self._peers = PeerManager(
            self._fernet_keys,
            throttling=self._throttling,
            event_callback=self._event_stream,
        )

        # Initialize metrics collector in child process
        self._metrics = self._metrics_factory()

        self._list_sni = SNIProxy(self._peers)
        self._list_peer = PeerListener(self._peers, metrics=self._metrics)

        # Start metrics reporting task
        self._metrics_task = asyncio.create_task(self._report_metrics_loop())

    async def _report_metrics_loop(self) -> None:
        """Schedule periodic metrics reporting."""
        with contextlib.suppress(asyncio.CancelledError):
            next_report = time.monotonic() + self._metrics_interval
            while True:
                now = time.monotonic()
                sleep_time = next_report - now
                if sleep_time > 0:
                    await asyncio.sleep(sleep_time)
                await self._collect_and_report_metrics()
                next_report += self._metrics_interval

    async def _collect_and_report_metrics(self) -> None:
        """Collect current state and report metrics."""
        if not self._metrics:
            return

        if not self._peers:
            self._metrics.gauge("snitun.worker.peer_connections", 0)
            return

        protocol_version_counts: dict[int, int] = {
            0: 0,
            1: 0,
        }

        for peer in self._peers.iter_peers():
            if peer.protocol_version not in protocol_version_counts:
                protocol_version_counts[peer.protocol_version] = 0
                # Log out unknown protocol versions
                _LOGGER.warning(
                    "Unknown protocol version %d for peer %s",
                    peer.protocol_version,
                    peer.hostname,
                )
            protocol_version_counts[peer.protocol_version] += 1

        self._metrics.gauge(
            "snitun.worker.peer_connections",
            sum(protocol_version_counts.values()),
        )
        for version, count in protocol_version_counts.items():
            self._metrics.gauge(
                "snitun.worker.peer_connections",
                count,
                {"protocol_version": str(version)},
            )

    def _event_stream(self, peer: Peer, event: PeerManagerEvent) -> None:
        """Event stream peer connection data."""
        if event == PeerManagerEvent.CONNECTED:
            if peer.hostname not in self._sync:
                self._peer_count.value += 1
            for hostname in peer.all_hostnames:
                self._sync[hostname] = None
        else:
            if peer.hostname in self._sync:
                self._peer_count.value -= 1
            for hostname in peer.all_hostnames:
                self._sync.pop(hostname, None)

    def shutdown(self) -> None:
        """Shutdown child process."""
        self._new.put(None)
        self.join(10)

    def handover_connection(
        self,
        con: socket,
        data: bytes,
        sni: str | None = None,
    ) -> None:
        """Move new connection to worker."""
        self._new.put_nowait((con, data, sni))

    def run(self) -> None:
        """Run the worker process."""
        _LOGGER.info("Start worker: %s", self.name)

        # Init new event loop
        self._loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self._loop)

        # Start eventloop
        running_loop = Thread(target=self._loop.run_forever)
        running_loop.start()

        # Init backend
        asyncio.run_coroutine_threadsafe(self._async_init(), loop=self._loop).result()

        while True:
            new = self._new.get()
            if new is None:
                break

            new[0].setblocking(False)
            asyncio.run_coroutine_threadsafe(
                self._async_new_connection(*new),
                loop=self._loop,
            )

        # Shutdown worker
        _LOGGER.info("Stoping worker: %s", self.name)

        # Cancel metrics task if running
        if self._metrics_task and not self._metrics_task.done():
            self._metrics_task.cancel()
            # Wait for metrics task to finish
            # Create a coroutine that waits for the task

            async def wait_for_task() -> None:
                with contextlib.suppress(asyncio.CancelledError):
                    if self._metrics_task:
                        await self._metrics_task

            with contextlib.suppress(asyncio.CancelledError, asyncio.TimeoutError):
                asyncio.run_coroutine_threadsafe(
                    wait_for_task(),
                    loop=self._loop,
                ).result()

        assert self._peers is not None, "PeerManager not initialized"
        asyncio.run_coroutine_threadsafe(
            self._peers.close_connections(),
            loop=self._loop,
        ).result()
        self._loop.call_soon_threadsafe(self._loop.stop)
        running_loop.join(10)

    async def _async_new_connection(
        self,
        con: socket,
        data: bytes,
        sni: str | None,
    ) -> None:
        """Handle incoming connection."""
        try:
            reader, writer = await asyncio.open_connection(sock=con)
        except OSError:
            con.close()
            return

        # Select the correct handler for process connection
        assert self._loop is not None, "Event loop not initialized"
        if sni:
            assert self._list_sni is not None, "SNIProxy not initialized"
            self._loop.create_task(
                self._list_sni.handle_connection(reader, writer, data=data, sni=sni),
            )
        else:
            assert self._list_peer is not None, "PeerListener not initialized"
            self._loop.create_task(
                self._list_peer.handle_connection(reader, writer, data=data),
            )