from __future__ import annotations

import logging
import time
import uuid
from collections import OrderedDict
from concurrent.futures import ProcessPoolExecutor
from concurrent.futures._base import Executor, Future
from concurrent.futures.thread import ThreadPoolExecutor
from inspect import isclass
from logging.handlers import QueueHandler
from multiprocessing import Manager, get_context
from multiprocessing.managers import SyncManager
from queue import Queue
from threading import Event, Thread
from typing import Any, Callable, Generator, Mapping, Optional
from typing import OrderedDict as OrderedDictType
from typing import Sequence, Tuple, Type, Union

from smartpipeline.containers import (
    BatchConcurrentStageContainer,
    BatchStageContainer,
    ConcurrentStageContainer,
    ConnectedStageMixin,
    ContainerType,
    SourceContainer,
    StageContainer,
)
from smartpipeline.defaults import CONCURRENCY_WAIT, MAX_QUEUES_SIZE
from smartpipeline.error.handling import ErrorManager, RetryManager
from smartpipeline.item import Item, Stop
from smartpipeline.stage import BatchStage, ItemsQueue, Source, Stage, StageType
from smartpipeline.utils import LogsReceiver, ProcessCounter, ThreadCounter, last_key

__author__ = "Giacomo Berardi <giacbrd.com>"


def _stage_initialization_with_logger(
    logs_queue: Queue, stage_class: Type[StageType], args: Sequence, kwargs: Mapping
) -> StageType:
    root_logger = logging.getLogger()
    # only by comparing string of queues we obtain their "original" address
    if not any(
        isinstance(handler, QueueHandler) and str(handler.queue) == str(logs_queue)
        for handler in root_logger.handlers
    ):
        handler = QueueHandler(logs_queue)
        root_logger.addHandler(handler)
    return stage_class(*args, **kwargs)


class EmptyStage(Stage):
    def process(self, item: Item) -> Item:
        return item


class Pipeline:
    def __init__(
        self,
        max_init_workers: Optional[int] = None,
        max_queues_size: int = MAX_QUEUES_SIZE,
    ):
        """
        :param max_init_workers: Number of workers to use for concurrent initialization of stages, default the number of CPUs
        :param max_queues_size: Maximum size of any queue instanced for the pipeline (stage input and output queues)
        """
        self._containers: OrderedDictType[str, ContainerType] = OrderedDict()
        self._error_manager = ErrorManager()
        self._max_init_workers = max_init_workers
        self._init_executor: Optional[Executor] = None
        self._wait_previous_executor: Optional[Executor] = None
        self._pipeline_executor: Optional[Thread] = None
        self._max_queues_size = max_queues_size
        self._out_queue = None
        self._enqueue_source = False
        # we approach sharing memory between processes exclusively through a `multiprocessing.Manager`
        self._sync_manager: Optional[SyncManager] = None
        self._count = 0
        self._executors_ready = False
        self._logs_receiver: Optional[LogsReceiver] = None
        self._name = f"{self.__class__.__name__}-{str(uuid.uuid4())[:8]}"
        self._logger = logging.getLogger(self._name)
        self._source_name = f"SourceOf{self._name}"
        # Initialize with a trivial "support" source, on which we can eventually send items when using `process` methods
        self._source_container = SourceContainer(self._source_name)
        # A fake container for the pipeline to support data structures initialization
        self._fake_container = StageContainer(
            f"FakeContainerOf{self._name}", EmptyStage(), ErrorManager(), RetryManager()
        )

    def _start_logs_receiver(self) -> LogsReceiver:
        if self._logs_receiver is None:
            if self._sync_manager is None:
                self._sync_manager = Manager()
            logs_queue = self._sync_manager.Queue()
            self._logs_receiver = LogsReceiver(logs_queue)
            self._logs_receiver.start()
        return self._logs_receiver

    def _stop_logs_receiver(self):
        if self._logs_receiver is not None:
            self._logs_receiver.stop()
            self._logs_receiver = None

    def _get_logs_receiver_queue(self) -> Queue:
        if self._logs_receiver is not None:
            return self._logs_receiver.queue
        else:
            raise AttributeError("Logs queue for multiprocessing never initialized")

    @property
    def name(self) -> str:
        """
        Pipeline unique name, equivalent to its logger name
        """
        return self._name

    def _new_mp_queue(self) -> ItemsQueue:
        """
        Construct queue for multiprocessing communication
        """
        # why not use multiprocessing.Queue or JoinableQueue instead of the manager?
        # answer https://stackoverflow.com/a/45236748 and the lack of "safety" of task_done() and join()
        if self._sync_manager is None:
            self._sync_manager = Manager()
        return self._sync_manager.Queue(maxsize=self._max_queues_size)

    def _new_queue(self) -> ItemsQueue:
        """
        Construct queue for communication
        """
        return Queue(maxsize=self._max_queues_size)

    def _new_mp_event(self) -> Event:
        """
        Construct synchronization event for multiprocessing
        """
        if self._sync_manager is None:
            self._sync_manager = Manager()
        return self._sync_manager.Event()

    @staticmethod
    def _new_event() -> Event:
        """
        Construct synchronization event
        """
        return Event()

    def _new_mp_counter(self) -> ProcessCounter:
        """
        Construct a safe counter for multiprocessing
        """
        if self._sync_manager is None:
            self._sync_manager = Manager()
        return ProcessCounter(self._sync_manager)

    @staticmethod
    def _new_counter() -> ThreadCounter:
        """
        Construct a safe counter for threads
        """
        return ThreadCounter()

    def _wait_executors(self, wait_seconds: float = CONCURRENCY_WAIT):
        """
        Wait for all containers to start

        :param wait_seconds: Recurrently wait these seconds for all stage initializers to finish
        """
        if self._executors_ready:
            return
        if self._init_executor is not None:
            self._init_executor.shutdown(wait=True)
            self._init_executor = None
        while not all(self._containers.values()):
            time.sleep(wait_seconds)
        if self._wait_previous_executor is not None:
            self._wait_previous_executor.shutdown(wait=True)
        for name, container in self._containers.items():
            if isinstance(
                container, (ConcurrentStageContainer, BatchConcurrentStageContainer)
            ):
                container.run()
        # finalize initialization of the error manager shared by this and other stage threads
        self._error_manager.on_start()
        self._executors_ready = True
        self._logger.debug("Pipeline ready to run")

    def shutdown(self):
        """
        Execute shutdown of various pool executors and multiprocessing stuff.
        The developer should not need to call it explicitly.
        """
        if self._out_queue is not None:
            self._out_queue.join()
        if self._init_executor is not None:
            self._init_executor.shutdown()
        for name, container in self._containers.items():
            if isinstance(
                container, (ConcurrentStageContainer, BatchConcurrentStageContainer)
            ):
                container.shutdown()
        if self._sync_manager is not None:
            self._sync_manager.shutdown()

    def build(self) -> Pipeline:
        """
        Pipeline builder method
        """
        if not any(self._containers):
            raise ValueError("Must append at least a stage")
        self._logger.debug("Building the pipeline on stages: %s", self._log_stages())
        self._wait_executors()
        return self

    def _exit_run(self, source_thread, terminator_thread):
        self.stop()
        if source_thread is not None:
            source_thread.join()
        # TODO in case of errors we loose pending items!
        self._terminate_all(force=True)
        if terminator_thread is not None:
            terminator_thread.join()
        self.shutdown()
        self._count += 1

    def run(self) -> Generator[Item, None, None]:
        """
        Run the pipeline given a source and a concatenation of stages.
        Get the sequence of items through iteration

        :return: Iterator over processed items
        :raises ValueError: When a source has not been set for the pipeline
        """
        if not self._source_container.is_set():
            raise ValueError("Set the data source for this pipeline")
        self._logger.debug("Running the pipeline on stages: %s", self._log_stages())
        counter = 0
        last_stage_name = self._last_stage_name()
        terminator_thread = None
        source_errors = None
        source_thread = None
        # in case the first stage is concurrent
        if self._enqueue_source:
            source_errors = Queue()
            source_thread = Thread(
                target=self._source_container.pop_into_queue,
                kwargs={"errors_queue": source_errors},
            )
            source_thread.start()
        while True:
            for name, container in self._containers.items():
                try:
                    # first, check errors from the source
                    if source_errors is not None and not source_errors.empty():
                        source_error = source_errors.get(block=False)
                        if source_error is not None:
                            raise source_error
                    # concurrent stages run by themselves in threads/processes
                    if not isinstance(
                        container,
                        (ConcurrentStageContainer, BatchConcurrentStageContainer),
                    ):
                        container.process()
                    # but me must periodically check for errors
                    else:
                        container.check_errors()
                except Exception as e:
                    self._exit_run(source_thread, terminator_thread)
                    raise e
                # retrieve finally processed items from the last stage
                if name == last_stage_name:
                    for _ in range(
                        container.size
                        if isinstance(container, BatchStageContainer)
                        else 1
                    ):
                        item = container.get_processed()
                        if item is not None:
                            if not isinstance(item, Stop):
                                try:
                                    yield item
                                except GeneratorExit:
                                    # when the caller exits the run loop terminates everything
                                    self._exit_run(source_thread, terminator_thread)
                                    return
                                finally:
                                    counter += 1
                                self._count += 1
                            # if a stop is finally signaled, start termination of all containers
                            elif (
                                not self._all_terminated() and terminator_thread is None
                            ):
                                terminator_thread = Thread(target=self._terminate_all)
                                terminator_thread.start()
                        # an item is None if the final output queue is empty
                        else:
                            break
            # exit the loop only when all items have been returned
            if self._all_empty() and counter >= self._source_container.count():
                if source_thread is not None:
                    source_thread.join()
                if terminator_thread is not None:
                    terminator_thread.join()
                    self.shutdown()
                return

    @property
    def count(self) -> int:
        """
        Get the number of items processed by all executed runs, also for items which have failed

        :return: Count of processed items
        """
        return self._count

    def _terminate_all(
        self, force: bool = False, wait_seconds: float = CONCURRENCY_WAIT
    ):
        """
        Terminate all running containers

        :param force: If True do not wait for a container to process all items produced by the source
        :param wait_seconds: Time to wait before pinging again a container for its termination
        """
        self._logger.debug("Terminating the pipeline")
        # scroll the pipeline by its order (this is critical for correctness)
        # and terminate stages after the relative queues are empty
        for container in self._containers.values():
            if not force:
                # ensure the stage has processed all source items
                while container.count() < self._source_container.count():
                    time.sleep(wait_seconds)
            container.terminate()
            if isinstance(container, ConcurrentStageContainer):
                if force:
                    # empty the queues, losing pending items
                    container.empty_queues()
                while not container.is_terminated():
                    time.sleep(wait_seconds)
                container.queues_join()
                while not container.queues_empty():
                    time.sleep(wait_seconds)
            if (
                not isinstance(container, ConcurrentStageContainer)
                or not container.parallel
            ):
                container.stage.on_end()
        self._error_manager.on_end()
        self._stop_logs_receiver()
        self._logger.debug("Termination done")

    def _all_terminated(self) -> bool:
        """
        Check if all containers have been alerted for termination and are exited
        """
        return all(container.is_terminated() for container in self._containers.values())

    def _all_empty(self) -> bool:
        """
        Check if all containers are terminated and there are not items left in the queues
        """
        return self._all_terminated() and all(
            container.queues_empty()
            for container in self._containers.values()
            if isinstance(
                container, (ConcurrentStageContainer, BatchConcurrentStageContainer)
            )
        )

    def process(self, item: Item) -> Item:
        """
        Process a single item synchronously (no concurrency) through the pipeline
        """
        self._logger.debug("Processing %s on stages: %s", item, self._log_stages())
        last_stage_name = last_key(self._containers)
        self._source_container.prepend_item(item)
        ret = None
        for name, container in self._containers.items():
            container.process()
            if name == last_stage_name:
                ret = container.get_processed(block=True)
        return ret or item

    def process_async(
        self, item: Item, callback: Optional[Callable[[Item], Any]] = None
    ):
        """
        Process a single item asynchronously through the pipeline, stages may run concurrently.
        The call returns immediately, processed items are retrieved with :meth:`.Pipeline.get_item`

        :param callback: A function to call after a successful process of the item
        """
        self._logger.debug(
            "Processing asynchronously %s on stages: %s", item, self._log_stages()
        )
        if callback is not None:
            item.set_callback(callback)
        self._source_container.prepend_item(item)
        self._start_pipeline_executor()

    def stop(self):
        """
        Tell the source to stop to generate items and consequently terminate the pipeline
        after all "remaining" items are processed.
        """
        self._source_container.stop()

    def get_item(self, block: bool = True) -> Item:
        """
        Get a single item from the asynchronous execution of the pipeline on single items from :meth:`.Pipeline.process_async`

        :param block: If True wait indefinitely for the next processed item
        :raises ValueError: When there is not output queue set, the pipeline is not running asynchronously
        :raises queue.Empty: When we do not block and the queue is empty
        """
        if self._out_queue is not None:
            item = self._out_queue.get(block)
            self._out_queue.task_done()
            return item
        else:
            raise ValueError(
                "No pipeline is running asynchronously, not item can be retrieved from the output queue"
            )

    def set_source(self, source: Source) -> Pipeline:
        """
        Set the source of the pipeline: a subclass of :class:`.stage.Source`
        """
        source.set_name(self._source_name)
        self._source_container.set(source)
        return self

    def set_error_manager(self, error_manager: ErrorManager) -> Pipeline:
        """
        Set the error manager for handling errors from each stage item processing
        """
        self._error_manager = error_manager
        for container in self._containers.values():
            container.set_error_manager(self._error_manager)
        return self

    def _last_stage_name(self) -> Optional[str]:
        if self._containers:
            return last_key(self._containers)
        return None

    def _wait_for_previous(
        self,
        container: ConnectedStageMixin,
        last_stage_name: Optional[str],
        wait_seconds: float = CONCURRENCY_WAIT,
    ):
        """
        Given a container we want to append to the pipeline, wait for the last one (added to the pipeline) to be created

        :param container: A container to add to the pipeline
        :param last_stage_name: Name of the last stage currently in the pipeline
        :param wait_seconds: Time to recurrently wait the construction of the container relative to the last stage in the pipeline
        """

        def _waiter():
            if last_stage_name is not None:
                while self._containers[last_stage_name] is self._fake_container:
                    time.sleep(wait_seconds)
                container.set_previous(self._containers[last_stage_name])
            else:
                container.set_previous(self._source_container)

        executor = self._get_wait_previous_executor()
        executor.submit(_waiter)

    def _build_container(
        self,
        name: str,
        stage: StageType,
        concurrency: int,
        parallel: bool,
        retryable_errors: Tuple[Type[Exception], ...],
        max_retries: int,
        backoff: Union[float, int],
    ) -> ContainerType:
        """
        Get a new container instance according to the pipeline configuration

        :param name: Stage name
        :param stage: A stage instance
        :param concurrency: Number of concurrent stage executions, if 0 then just create the non-concurrent containers
        :param parallel: If True use multiprocessing, otherwise threads
        :param retryable_errors: list of exceptions for which the stage applies an exponential backoff strategy. When the maximum number of retries is hit, a :class:`.error.exceptions.RetryError` is raised
        :param max_retries: maximum number of retries for the stage before raising a RetryError(SoftError) (default 0)
        :param backoff: backoff factor for the `exponential backoff strategy` used by the stage when it raises one of the exceptions declared in `retryable_errors` param (default 0.0)
        """
        retry_manager = RetryManager(
            backoff=backoff, max_retries=max_retries, retryable_errors=retryable_errors
        )
        if concurrency <= 0:
            constructor = (
                BatchStageContainer if isinstance(stage, BatchStage) else StageContainer
            )
            # if not concurrent we must explicitly finalize initialization of this single stage object
            stage.on_start()
            return constructor(name, stage, self._error_manager, retry_manager)
        else:
            constructor = (
                BatchConcurrentStageContainer
                if isinstance(stage, BatchStage)
                else ConcurrentStageContainer
            )
            if parallel:
                self._start_logs_receiver()
                logs_queue = self._get_logs_receiver_queue()
                return constructor(
                    name,
                    stage,
                    self._error_manager,
                    retry_manager,
                    self._new_mp_queue,
                    self._new_mp_counter,
                    self._new_mp_event,
                    concurrency,
                    parallel,
                    logs_queue,
                )
            else:
                # if the stage is executed on multiple threads we must finalize initialization once,
                # while on multiprocessing each process executor calls it for its own copy of the stage
                stage.on_start()
                return constructor(
                    name,
                    stage,
                    self._error_manager,
                    retry_manager,
                    self._new_queue,
                    self._new_counter,
                    self._new_event,
                    concurrency,
                    parallel,
                )

    def get_stage(self, name: str) -> Optional[StageType]:
        """
        Get a stage instance by its name
        """
        ret = self._containers.get(name)
        return ret.stage if ret else None

    def append(
        self,
        name: str,
        stage: StageType,
        concurrency: int = 0,
        parallel: bool = False,
        retryable_errors: Tuple[Type[Exception], ...] = tuple(),
        max_retries: int = 0,
        backoff: Union[float, int] = 0.0,
    ) -> Pipeline:
        """
        Append a stage to the pipeline just after the last one appended, or after the source if it is the first stage

        :param name: Name for identify the stage in the pipeline, it is also set in the stage and it must be unique in the pipeline
        :param stage: Instance of a stage
        :param concurrency: Number of concurrent stage executions, if 0 then threads/processes won't be involved for this stage
        :param parallel: If True use multiprocessing, otherwise threads
        :param retryable_errors: list of exceptions for which the stage applies an exponential backoff strategy. When the maximum number of retries is hit, a :class:`.error.exceptions.RetryError` is raised
        :param max_retries: maximum number of retries for the stage before raising a RetryError(SoftError) (default 0)
        :param backoff: backoff factor for the `exponential backoff strategy` used by the stage when it raises one of the exceptions declared in `retryable_errors` param (default 0.0)
        """
        self._executors_ready = False
        # FIXME here we force a BatchStage to run on a thread, but we would leave it on the main thread
        if concurrency < 1 and isinstance(stage, BatchStage):
            parallel = False
            concurrency = 1
        self._check_stage_name(name)
        self._check_retries_params(retryable_errors, max_retries, backoff)
        container = self._build_container(
            name, stage, concurrency, parallel, retryable_errors, max_retries, backoff
        )
        if concurrency > 0:
            # if it is concurrent and it is the first stage, make the source work on a output queue
            if not self._containers:
                self._enqueue_source = True
        self._wait_for_previous(
            container, self._last_stage_name()
        )  # wait that previous stage is initialized
        self._containers[name] = container
        return self

    # deprecated
    append_stage = append  # pragma: no cover

    def append_concurrently(
        self,
        name: str,
        stage_class: Type[StageType],
        args: Optional[Sequence] = None,
        kwargs: Optional[Mapping] = None,
        concurrency: int = 0,
        parallel: bool = False,
        retryable_errors: Tuple[Type[Exception], ...] = tuple(),
        max_retries: int = 0,
        backoff: Union[float, int] = 0.0,
    ) -> Pipeline:
        """
        Append a stage class to the pipeline just after the last one appended, or after the source if it is the first stage.
        The stage construction will be executed concurrently respect to the general pipeline construction

        :param name: Name for identify the stage in the pipeline, it is also set in the stage and it must be unique in the pipeline
        :param stage_class: Class of a stage
        :param args: List of arguments for the stage constructor
        :param kwargs: Dictionary of keyed arguments for the stage constructor
        :param concurrency: Number of concurrent stage executions, if 0 then threads/processes won't be involved for this stage
        :param parallel: If True use multiprocessing, otherwise threads (also for concurrent construction)
        :param retryable_errors: list of exceptions for which the stage applies an exponential backoff strategy. When the maximum number of retries is hit, a :class:`.error.exceptions.RetryError` is raised
        :param max_retries: maximum number of retries for the stage before raising a RetryError(SoftError) (default 0)
        :param backoff: backoff factor for the `exponential backoff strategy` used by the stage when it raises one of the exceptions declared in `retryable_errors` param (default 0.0)
        """
        if not isclass(stage_class):
            raise ValueError(
                "Must pass stage classes, and not stage instances, as argument"
            )
        self._executors_ready = False
        parallel_construction = parallel
        # FIXME here we force a BatchStage to run on a thread, but we would like to leave it on the main thread
        if concurrency < 1 and issubclass(stage_class, BatchStage):
            parallel = False
            concurrency = 1
        if kwargs is None:
            kwargs = {}
        if args is None:
            args = []
        self._check_stage_name(name)
        self._check_retries_params(retryable_errors, max_retries, backoff)
        # if it is concurrent, and it is the first stage, make the source work on a output queue
        if concurrency > 0 and not self._containers:
            self._enqueue_source = True
        last_stage_name = self._last_stage_name()
        # set it immediately with an empty stage container so the order of the calls of this method is followed in `_containers`
        self._containers[name] = self._fake_container
        self._start_logs_receiver()
        logs_queue = self._get_logs_receiver_queue()
        executor = self._get_init_executor(parallel_construction)
        if parallel_construction:
            future = executor.submit(
                _stage_initialization_with_logger, logs_queue, stage_class, args, kwargs
            )
        else:
            future = executor.submit(stage_class, *args, **kwargs)

        def append_stage(stage_future: Future):
            stage = stage_future.result()
            container = self._build_container(
                name,
                stage,
                concurrency,
                parallel,
                retryable_errors,
                max_retries,
                backoff,
            )
            self._wait_for_previous(container, last_stage_name)
            self._containers[name] = container

        future.add_done_callback(append_stage)
        return self

    # deprecated
    append_stage_concurrently = append_concurrently

    def _get_init_executor(self, parallel: bool = False) -> Executor:
        """
        Get a pool executor for concurrent stage initialization

        :param parallel: True if the executor uses multiprocessing, otherwise treads
        """
        if self._init_executor is None:
            if not parallel:
                self._init_executor = ThreadPoolExecutor(
                    max_workers=self._max_init_workers
                )
            else:
                self._init_executor = ProcessPoolExecutor(
                    max_workers=self._max_init_workers, mp_context=get_context("spawn")
                )
        return self._init_executor

    def _get_wait_previous_executor(self) -> Executor:
        """
        Get a pool executor for the function that will recurrently wait for a container to be ready
        """
        if self._wait_previous_executor is None:
            self._wait_previous_executor = ThreadPoolExecutor()
        return self._wait_previous_executor

    def _start_pipeline_executor(self) -> Thread:
        """
        Get a thread where to run a pipeline that accepts asynchronous processing of single items
        """
        if self._pipeline_executor is None:
            self._init_out_queue()

            def pipeline_runner():
                for item in self.run():
                    item.callback()
                    self._out_queue.put(item)

            self._pipeline_executor = Thread(target=pipeline_runner)
            self._pipeline_executor.start()
        return self._pipeline_executor

    def _check_stage_name(self, name: str):
        """
        Check if a stage name is not already defined in the pipeline
        :raises ValueError: Stage name is already defined in the pipeline
        """
        if name in self._containers:
            raise ValueError(f"The stage name {name} is already used in this pipeline")

    @staticmethod
    def _check_retries_params(
        retryable_errors: Tuple[Type[Exception], ...],
        max_retries: int,
        backoff: Union[float, int],
    ):
        """
        Check the values used to apply the retry strategy for a stage
        """
        if not (isinstance(backoff, (float, int)) and backoff >= 0):
            raise ValueError(
                "The `backoff` parameter must be either a float or int and its value must be >= 0"
            )
        if not (isinstance(max_retries, int) and max_retries >= 0):
            raise ValueError(
                "The `max_retries` parameter must be an int and its value must be >= 0"
            )
        if not (
            isinstance(retryable_errors, tuple)
            and all(issubclass(err, Exception) for err in retryable_errors)
        ):
            raise ValueError(
                "The `retryable_errors` parameter must be a tuple and all its elements must be Exceptions"
            )

    def _init_out_queue(self):
        """
        Get the internal output pipeline for asynchronous processing of single items
        """
        self._out_queue = self._new_queue()

    def _log_stages(self):
        return ", ".join(self._containers.keys())
