import logging
import multiprocessing as mp
import threading
import time
from typing import Any, Callable, NamedTuple

import numpy.typing as npt

from where_fi.collection import CSIMatrix, ingest
from where_fi.collection.protocols import CSIProducer, MergedCSI
from where_fi.config import config
from where_fi.processing.preprocess import Preprocessor
from where_fi.visualise import server as visualise


class Receiver(NamedTuple):
    ip: ingest.Host
    receiver: ingest.FeitReceiver
    thread: threading.Thread


class Scheduler:
    """
    A simple scheduler that runs a function at a given rate.
    """

    def __init__(self, rate: float, func: Callable[[], None]) -> None:
        self.rate = rate
        self.func = func
        self.active = True

    def run(self) -> None:
        """
        Run the scheduler in a loop, calling the function at the given rate.
        """
        while self.active:
            self.func()
            time.sleep(1 / self.rate)

    def start(self) -> None:
        """
        Start the scheduler in a separate thread.
        """
        self.thread = threading.Thread(target=self.run)
        self.thread.start()

    def stop(self) -> None:
        """
        Stop the scheduler.
        """
        self.active = False


class CSIApplication:
    """
    A high-level application that manages CSI data ingestion and processing.

    This allows a simple interface to be used for the main CLI logic, without having to
    deal with the threads and queues directly.

    It can be used using decortors to register callbacks for different stages of the
    processing pipeline:

        - `on_pre_merge`: Called with the raw CSI data (including headers) from each
                          receiver, before it is merged into a single matrix.
        - `on_sample`: Called once per sample with the merged CSI data.
        - `on_process`: Called at the processing sample rate, without any data (this is
                        mostly used as a scheduler).
    """

    def __init__(self, producer: CSIProducer, visualise_raw: bool = False) -> None:
        self.logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}")

        self.producer = producer

        self.webapp_queue: "mp.Queue[visualise.VisualiserData]" = mp.Queue(
            config.processing_sample_rate
        )
        self.webapp = visualise.Webapp()
        self.webapp_thread = threading.Thread(
            target=self.webapp.start, args=(self.webapp_queue,)
        )
        self.visualise_raw = visualise_raw

        self.pre_merge_callback = None
        self.raw_csi_callback = None
        self.preprocessed_csi_callback = None
        self.processing_callback = None

        self.preprocessor = Preprocessor()

    def visualise_data(
        self, data: npt.NDArray[Any], dtype: visualise.figures.Figure
    ) -> None:
        """
        Update a visualisation with the given data. The available visualisations are as
        per visualise.server.all_figures.
        """
        if not self.webapp_queue.full():
            self.webapp_queue.put(visualise.VisualiserData(data, dtype))

    def on_raw(self, func: Callable[[CSIMatrix], None]) -> Callable[[CSIMatrix], None]:
        """
        Decorator to register a callback for the sample preprocessing.
        """

        def decorator(data: CSIMatrix) -> None:
            func(data)

        self.raw_csi_callback = decorator
        return decorator

    def on_sample(
        self, func: Callable[[CSIMatrix], None]
    ) -> Callable[[CSIMatrix], None]:
        """
        Decorator to register a callback for the sample preprocessing.
        """

        def decorator(data: CSIMatrix) -> None:
            self.visualise_data(data, visualise.figures.Figure.RAW_CSI)
            func(data)

        self.preprocessed_csi_callback = decorator
        return decorator

    def on_pre_merge(
        self,
    ) -> Callable[[Callable[[dict[ingest.Host, ingest.CSI]], None]], None]:
        def decorator(func: Callable[[dict[ingest.Host, ingest.CSI]], None]) -> None:
            self.pre_merge_callback = func

        return decorator

    def on_process(
        self, func: Callable[[CSIMatrix], None]
    ) -> Callable[[CSIMatrix], None]:
        """
        Decorator to register a callback for the sample processing.
        """
        self.processing_callback = func
        return func

    def process_sample(self, sample: MergedCSI) -> None:
        """Data pipeline for processing a sample.
        This is the entry point for a sample being received from the producer. Depending
        on the callbacks which are registered, this will call the appropriate functions
        to pre-process and visualise the sample.
        """
        if self.visualise_raw:
            self.visualise_data(sample.matrix, visualise.figures.Figure.RAW_CSI)
        if self.raw_csi_callback is not None:
            self.raw_csi_callback(sample.matrix)
        if self.pre_merge_callback is not None:
            self.pre_merge_callback(sample.frames)
        processed = self.preprocessor.preprocess(sample.matrix)
        if self.visualise_raw:
            self.visualise_data(processed, visualise.figures.Figure.PROCESSED_CSI)
        if self.preprocessed_csi_callback is not None:
            self.preprocessed_csi_callback(processed)

    def listen(self) -> None:
        """
        Listen for incoming data from the selected ingestor, and call the appropriate
        callback.
        """
        data_gen = self.producer()
        while True:
            try:
                sample = next(data_gen)
                self.process_sample(sample)
            except StopIteration:
                break
            except Exception as e:
                self.logger.error(f"Error processing sample: {e}", exc_info=True)
                break

    def start(self) -> None:
        # for receiver in self.receivers:
        #     receiver.thread.start()
        self.buffer_thread = threading.Thread(
            target=self.listen,
        )
        self.buffer_thread.start()
        self.webapp_thread.start()

        if self.processing_callback is not None:

            def get_proc_sample() -> None:
                sample = self.preprocessor.last_sample
                if sample is not None and self.processing_callback is not None:
                    self.processing_callback(sample)
                else:
                    self.logger.warning("No samples to process")

            self.scheduler = Scheduler(config.processing_sample_rate, get_proc_sample)
            self.scheduler.start()
        try:
            self.webapp_thread.join()
        except KeyboardInterrupt:
            self.stop()

    def stop(self) -> None:
        self.logger.info("Stopping application")

        self.producer.stop()

        if hasattr(self, "scheduler"):
            self.scheduler.stop()
            self.scheduler.thread.join()
        self.buffer_thread.join()

        self.webapp.active = False
        self.webapp_thread.join()
        self.logger.info("Application stopped")