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, NoopCSIProducer, ingest
from where_fi.collection.csi_frame import CSI
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
from where_fi.visualise.server import figures


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, visualise_raw: bool = False) -> None:
        self.logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}")

        self.producer = NoopCSIProducer()

        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()
        self.custom_figures: dict[
            visualise.figures.FigureId, visualise.figures.SpecificFigure
        ] = {}

    def set_producer(self, producer: CSIProducer) -> None:
        """
        Set the producer for the application. This is used to change the data source
        at runtime.
        """
        self.producer = producer

    def register_figure(
        self,
        figure_id: visualise.figures.FigureId,
        figure: visualise.figures.SpecificFigure,
    ) -> None:
        self.custom_figures[figure_id] = figure
        figures.all_figures[figure_id] = figure

    def visualise_data(
        self, data: npt.NDArray[Any], dtype: visualise.figures.FigureId
    ) -> None:
        """
        Update a visualisation with the given data. The available visualisations are as
        per visualise.server.all_figures.
        """
        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:
            func(data)

        self.preprocessed_csi_callback = decorator
        return decorator

    def on_pre_merge(
        self, func: Callable[[dict[ingest.Host, CSI]], None]
    ) -> Callable[[dict[ingest.Host, CSI]], None]:
        """
        Decorator to register a callback for the samples before merging.
        """
        self.pre_merge_callback = func
        return func

    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, sample.frames, visualiser=self.visualise_data
        )
        # 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()

        def get_proc_sample() -> None:
            sample = self.preprocessor.last_sample
            if self.visualise_raw and sample is not None:
                self.visualise_data(
                    sample,
                    visualise.figures.Figure.PROCESSED_CSI,
                )
            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")
