import logging
import multiprocessing as mp
import selectors
import socket
import subprocess
import threading
import time
from datetime import datetime
from typing import Iterator

import numpy as np

from ..config import config
from .csi_frame import CSI
from .protocols import MergedCSI

Host = tuple[str, int]


class FeitHost:
    def __init__(self, host: Host, command: str) -> None:
        self.command = command
        self.host = host
        self.logger = logging.getLogger(
            f"{__name__}.{self.__class__.__name__}-{self.host[0]}"
        )
        self.active = True
        self.checker = threading.Thread(target=self.check_continuous)
        self.checker.start()
        self.selectors: list[selectors.BaseSelector] = []

    def check_connection(self) -> bool:
        try:
            feitcsi_status = subprocess.run(
                f"ssh root@{self.host[0]} pgrep feitcsi",
                check=False,
                stdout=subprocess.DEVNULL,
                shell=True,
                timeout=3,
            )
            return feitcsi_status.returncode == 0
        except subprocess.TimeoutExpired:
            return False

    def connect(self) -> None:
        self.server = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.server.connect(self.host)
        self.server.send(b"stop\n")
        self.server.send(self.command.encode())

        for selector in self.selectors:
            selector.register(self.server, selectors.EVENT_READ, data=self)

        self.logger.info(f"Connected to {self.host}")

    def check_continuous(self) -> None:
        """
        Repeatedly check if the FeitCSI service is running

        Because FeitCSI is using TCP, we will get no information if the service stops,
        or the computer is not reachable. In order to make debugging easier, this checks
        and logs continuously if the service is running.
        """
        last_status = False
        while self.active:
            self.logger.debug(f"Checking connection to {self.host[0]}")
            if not self.check_connection():
                self.logger.error(f"FeitCSI is not running on {self.host[0]}")
                last_status = False
            else:
                if not last_status:
                    self.connect()
                last_status = True
            time.sleep(1)
        self.logger.info("Stopping host checker")


class FeitTransmitter(FeitHost):
    def __init__(self) -> None:
        command = (
            f"feitcsi --frequency {config.central_freq} "
            f"--channel-width {config.channel_width} "
            f"--format {config.frame_format} "
            f"--mode inject -s 1 --verbose "
            f"--inject-delay {1_000_000 // config.collection_sample_rate}"
        )
        super().__init__(config.transmit_host, command)


class FeitReceiver(FeitHost):
    """
    A class used to connect to the host running FeitCSI and receive CSI data over a UDP
    socket.
    """

    def __init__(self, host: Host) -> None:
        command = (
            f"feitcsi --frequency {config.central_freq} "
            f"--channel-width {config.channel_width} "
            f"--format {config.frame_format} "
            f"--mode measure"
        )
        super().__init__(host, command)

    def recv(self) -> CSI:
        """
        Block until a CSI frame is received and return it.
        """

        # Receive the data from the socket
        data = self.server.recv(65535)
        # Decode the data using the CSI frame format
        csi = CSI(data)
        return csi

    def register(self, selector: selectors.BaseSelector) -> None:
        """
        Register the receiver as a selector. This will allow us to multiplex the I/O
        operations on a single thread.

        This allows us to block until any of the receivers receive data
        """
        self.selectors.append(selector)

        if hasattr(self, "server"):
            selector.register(self.server, selectors.EVENT_READ, data=self)


class CSIAntennaArray:
    """
    Represents a set of receiver hosts that are used to receive CSI data, assumed to be
    part of a linear antenna array.

    This class is used to buffer a set of CSI data from multiple receivers, and once a
    set of readings is ready for each receiver, it will be merged into a single sample
    using the antenna order in the configuration.

    When the data is ready, the following callbacks are called:
        - `pre_merge_callback`: Called with the data before merging. This is useful for
          per-antenna analysis, e.g. finding the correct antenna order
        - `sample_callback`: Called with the merged data
    """

    def __init__(self, receivers: list[FeitReceiver]) -> None:
        self.pending_data: dict[Host, tuple[datetime, CSI]] = {}
        self.pending_data_lock = mp.Lock()
        self.logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}")
        self.last_processed = datetime.now()
        self.receivers = receivers
        self.active = True

    def add_data(self, host: Host, data: CSI) -> None:
        if (
            host in self.pending_data
            and self.last_processed < self.pending_data[host][0]
        ):
            self.logger.warning(
                f"Skipping data from {host} at {self.pending_data[host][0]}"
            )

        with self.pending_data_lock:
            self.pending_data[host] = (datetime.now(), data)

    def is_ready(self) -> bool:
        for host in config.receive_hosts:
            if (
                host not in self.pending_data
                or self.pending_data[host][0] <= self.last_processed
            ):
                return False
        return True

    def process_data(self) -> None | MergedCSI:
        if not self.is_ready():
            self.logger.debug("Not all data is ready")
            return None
        self.last_processed = datetime.now()

        frames = {host: data[1] for host, data in self.pending_data.items()}
        antenna_data = [
            np.expand_dims(self.pending_data[ip][1].matrix[:, antenna], axis=2)
            for ip, antenna in config.antennas.order
        ]

        # We have data from all servers
        all_data = np.concat(antenna_data, axis=1)
        return MergedCSI(frames=frames, matrix=all_data)

    def process_forever(self) -> Iterator[MergedCSI]:
        # Register the receivers with the selector
        self.logger.info("Starting to process data from receivers")
        selector = selectors.DefaultSelector()

        for receiver in self.receivers:
            receiver.register(selector)

        while self.active:
            events = selector.select(timeout=0.1)
            if not events:
                self.logger.warning("No events received in the last 0.1 seconds")
                continue

            for key, _ in events:
                receiver = key.data
                self.add_data(receiver.host, receiver.recv())

            sample = self.process_data()
            if sample is not None:
                yield sample


class RealtimeCSIProducer:
    def __init__(self) -> None:
        self.receivers = [FeitReceiver(ip) for ip in config.receive_hosts]
        self.transmitter = FeitTransmitter()

        # Start injecting CSI frames
        self.buffer = CSIAntennaArray(self.receivers)
        self.is_live = True

    def __call__(self) -> Iterator[MergedCSI]:
        return self.buffer.process_forever()

    def stop(self) -> None:
        for r in self.receivers:
            r.active = False
        self.transmitter.active = False
        self.buffer.active = False
