"""SniTun worker for traffics."""

from __future__ import annotations

import asyncio
import contextlib
from ipaddress import IPv4Address, IPv6Address
import logging
from multiprocessing import Manager
from multiprocessing.context import ForkServerProcess
from socket import socket
from threading import Thread
import time
from typing import TYPE_CHECKING, Any

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(ForkServerProcess):
    """Worker for multiplexer.

    A worker runs an asyncio loop in a background thread and owns a
    multiprocessing.Manager, so the parent process is already multi-threaded by
    the time it spawns workers. Forking a multi-threaded process is unsafe and,
    as of Python 3.14, no longer the default start method on Linux, so the
    worker is spawned via forkserver (by subclassing ForkServerProcess) and
    forks from a clean single-threaded server process. The server only runs on
    Linux (it drives connections with select.epoll), where forkserver is always
    available. forkserver pickles the worker to hand it to the child, so
    ServerWorker must stay picklable; see __getstate__.
    """

    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
        self._manager: SyncManager = Manager()
        self._new = self._manager.Queue()
        self._sync = self._manager.dict()
        self._peer_count = self._manager.Value("peer_count", 0)

    def __getstate__(self) -> dict[str, Any]:
        """Return picklable state for spawn/forkserver.

        The SyncManager lives in the parent process and cannot be pickled (it
        holds a weakref via its Finalize). The child only needs the proxies
        (_new/_sync/_peer_count), which pickle and reconnect to the manager
        process on their own, so drop the manager itself from the state.
        """
        state = self.__dict__.copy()
        state["_manager"] = None
        return state

    @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.set(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.set(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,
        peer_address: IPv4Address | IPv6Address | None = None,
    ) -> None:
        """Move new connection to worker."""
        self._new.put_nowait((con, data, sni, peer_address))

    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("Stopping 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,
        peer_address: IPv4Address | IPv6Address | None = 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,
                    peer_address=peer_address,
                ),
            )
        else:
            assert self._list_peer is not None, "PeerListener not initialized"
            self._loop.create_task(
                self._list_peer.handle_connection(reader, writer, data=data),
            )
