# -*- coding: utf-8 -*-
"""Manages the pull/push with queues"""
from __future__ import unicode_literals

from typing import Dict, List, Union

import functools
import logging

from builtins import str as text
from inspect import signature
import queue
from queue import Empty, Full
from threading import Thread  # For non-multi-processing queues

# Multiprocessing imports
import multiprocessing
from ctypes import c_int

# Internal imports
from .classtemplate import ClassTemplate
from .exceptionhandler import ProcessNotStarted
from .metrics import Metrics, MetricType
from .common import DIRECTION
from .common import new_id, is_threaded, safe_get
from .common import (
    PRIORITY_QUEUES,
    SIMPLE_QUEUES,
    UNION_MULTIPROCESSING_QUEUES,
    UNION_SUPPORTED_EVENTS,
    UNION_SUPPORTED_LOCKS,
    UNION_SUPPORTED_QUEUES)


def validate_direction(func):
    """Decorator to check that 'direction' is an acceptable value.

    One of DIRECTION.FROM or DIRECTION.TO.

    :raises TypeError
    """
    @functools.wraps(func)
    def wrapper(*args, **kwargs):

        # Extract the function signature so we can check for direction
        sig = signature(func)
        bound = sig.bind_partial(*args, **kwargs)

        # Validate direction if present
        if "direction" in bound.arguments:
            arg_direction = bound.arguments['direction']

            # If the direction is valid, just keep going
            if arg_direction in DIRECTION:
                pass

            # If the direction is invalid, throw an error
            else:
                raise TypeError(
                    "destination must be a value from DIRECTION")

        # Call the normal function
        return func(*args, **kwargs)

    return wrapper


class LimitedLengthQueue(object):
    """Wrapper for an existing Queue to keep the length limited.

    NOT thread/process safe, only intended to be used by a single producer (i.e., Publishers).

    NOTE: This class is not currently used in the codebase. The metrics path previously used
    it but was replaced by bounded Queue(maxsize=) + put_nowait() in FEAT-002. Retained as a
    general-purpose utility in case future callers need a bounded blocking wrapper.
    """
    def __init__(self,
                 queue_instance: multiprocessing.queues.Queue,
                 max_size: int=0,
                 start_method: str=None):
        # Which multiprocess context to use
        self.multiprocessing_ctx = multiprocessing.get_context(start_method)

        self._queue = queue_instance
        self.max_size = max_size
        self._queue_length = self.multiprocessing_ctx.Value(c_int, 0)

        if not self._queue.empty():
            raise ValueError('Cannot start with a non-empty Queue!')

    def put(self, *args, **kwargs):
        """Add an element to the queue"""
        while self._queue_length.value >= self.max_size:
            self.get()

        self._queue.put(*args, **kwargs)
        with self._queue_length.get_lock():
            self._queue_length.value += 1

    def get(self, *args, **kwargs):
        """Retrieve an element from the queue

        :throws Empty
        """
        value = self._queue.get(*args, **kwargs)
        with self._queue_length.get_lock():
            self._queue_length.value -= 1

        return value

    def qsize(self):
        """Returns the number of elements in the queue

        Based on the wrapper's counter; does not rely on the underlying queue implementation.
        """
        return self._queue_length.value

    def __getattr__(self, item):
        if item in ('_queue', '_queue_length', 'max_size', 'get', 'put'):
            return self.__getattribute__(item)

        return getattr(self._queue, item)


class QueueLink(ClassTemplate):
    """Manages publishing from source and client queues"""

    def __init__(self,
                 source: UNION_SUPPORTED_QUEUES=None,
                 destination: Union[UNION_SUPPORTED_QUEUES, List[UNION_SUPPORTED_QUEUES]]=None,
                 *,  # Remainder are keyword arguments
                 name: str=None,
                 log_name: str=None,
                 start_method: str=None,
                 thread_only: bool=False,
                 link_timeout: float=0.01):
        """Manages the pull/push with queues

        Args:
            source: Queue to act as the source
            destination: Queue or list of queues to act as a destination
            name: Appended to class name in log records (``cl.name``)
            log_name: Appended to class name.name (``cl.name.log_name``)
                or just class name if ``name`` is not provided
            start_method: For multiprocess use: fork, spawn or forkserver
            thread_only: Only use threads, not separate processes, for links
            link_timeout: Tune the queue.get(timeout=link_timeout) value; default 0.01 sec
        """
        # Unique ID
        # Not used for cryptographic purposes, so excluding from Bandit
        self.id = new_id()

        self.name = name
        self.log_name = log_name

        self._initialize_logging_with_log_name(__name__)

        # Timeout for publisher queue.get calls
        self.link_timeout = link_timeout

        # Thread only usage (no multiprocessing)
        self.thread_only = thread_only

        # Start method for process-based use
        self.start_method = start_method

        # Use the correct types when creating multithreaded/process Events
        self.multiprocess_ctx = multiprocessing.get_context(self.start_method)

        # Indicate whether we have ever been started (Event())
        self.started = self.Event()

        # List of locks allows decoupled IO while also preventing the set of
        # queues from changing during IO operations
        self.queues_lock = self.Lock()
        self.last_queue_id = 0
        self.client_queues_source = {}
        self.client_queues_destination = {}
        self.client_pair_publishers = {}
        self.publisher_stops = {}

        # Metrics — two queues to avoid feeder-thread races:
        # Thread publishers use a threading queue (synchronous put, no feeder thread).
        # Process publishers use a multiprocessing queue (cross-process safe), created
        # lazily on first process-based publisher start to avoid unnecessary OS resources.
        self._metrics_queue_thread = queue.Queue(maxsize=100)
        self._metrics_queue_process = None

        # Short-circut usage
        if source:
            self.register_queue(q=source, direction=DIRECTION.FROM)

        if destination:
            if not isinstance(destination, list):
                destination = [destination]

            for q_inst in destination:
                self.register_queue(q=q_inst, direction=DIRECTION.TO)


    # Class contains Locks and Queues which cannot be pickled
    def __getstate__(self):
        """Prevent _QueueHandleAdapterBase from being pickled across Processes

        Raises:
            Exception
        """
        raise TypeError("Don't pickle me!")

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def __del__(self):
        self.close()

    def close(self) -> None:
        """Remove managed objects started within QueueLink"""

        # Make sure processes are stopped
        if hasattr(self, 'stop') and hasattr(self, 'started'):
            if self.started:
                self.stop()

        # Unset these variables to release any objects they held
        unset_list = ['_metrics_queue_thread',  # Metrics (thread publishers)
                      '_metrics_queue_process',  # Metrics (process publishers)
                      'started',  # Started event
                      'queues_lock',  # Queues lock
                      'publisher_stops',  # Publisher stop events
                      'client_queues_source',  # Queue references
                      'client_queues_destination',  # Queue references
                      'client_pair_publishers']  # Queue references

        for var in unset_list:
            if hasattr(self, var):
                setattr(self, var, None)

    def _initialize_logging_with_log_name(self, class_name: str):
        """Need to reverse the print order of log_name and name"""
        if hasattr(self, '_log'):
            if self._log is not None:
                return

        # Make a helpful name
        if self.name is None:
            name = class_name
        else:
            name = f'{class_name}.{self.name}'

        if self.log_name is not None:
            name = f'{name}.{self.log_name}'

        self._log = logging.getLogger(name)
        self.add_logging_handler(logging.NullHandler())

    def Event(self, *args, **kwargs) -> UNION_SUPPORTED_EVENTS:  # pylint: disable=invalid-name
        """Create a context-appropriate Event"""
        return self.multiprocess_ctx.Event(*args, **kwargs)

    def Lock(self, *args, **kwargs) -> UNION_SUPPORTED_LOCKS:  # pylint: disable=invalid-name
        """Create a context-appropriate Lock"""
        return self.multiprocess_ctx.Lock(*args, **kwargs)

    def Process(self, *args, **kwargs) -> multiprocessing.Process:  # pylint: disable=invalid-name
        """Create a context-appropriate Process"""
        return self.multiprocess_ctx.Process(*args, **kwargs)

    def Queue(self, *args, **kwargs) -> UNION_MULTIPROCESSING_QUEUES:  # pylint: disable=invalid-name
        """Create a context-appropriate Queue"""
        return self.multiprocess_ctx.Queue(*args, **kwargs)

    def _get_process_metrics_queue(self) -> UNION_MULTIPROCESSING_QUEUES:
        """Return the metrics queue for process-based publishers, creating it on first call.

        Deferred creation avoids allocating OS-level pipe resources on QueueLink instances
        that only ever use thread-based publishers.
        """
        if self._metrics_queue_process is None:
            self._metrics_queue_process = self.multiprocess_ctx.Queue(maxsize=100)

        return self._metrics_queue_process

    def stop(self) -> None:
        """Use to stop somewhat gracefully"""
        self._stop_publishers()

    @staticmethod
    def _publisher(stop_event: UNION_SUPPORTED_EVENTS,
                   source_id: str,
                   source_queue: UNION_SUPPORTED_QUEUES,
                   dest_queues_dict: Dict[str, UNION_SUPPORTED_QUEUES],
                   timeout: float,
                   metrics_queue: UNION_SUPPORTED_QUEUES,
                   metric_interval: int = 100,
                   name: str = None) -> None:
        """Move messages from the source queue to the destination queue

        Args:
            stop_event: Used to determine whether to stop the process
            source_id: Source queue ID
            source_queue: Source queue object
            dest_queues_dict: Dictionary holding a set of destination queues
            timeout: How long to wait for a message
            metrics_queue: Where to place metric snapshots (bounded; when full the
                oldest entry is evicted to make room for the newest)
            metric_interval: How many messages between periodic metric emissions
            name: Name to use in logging
        """
        def _emit_metrics():
            """Emit a metrics snapshot, evicting the oldest entry if the queue is full.

            Newer snapshots are more valuable than older ones — they reflect the most
            recent timing and count data — so we prefer to keep the newest rather than
            silently discard it.
            """
            data = metrics.get_all_data()
            try:
                metrics_queue.put_nowait(data)
            except Full:
                try:
                    metrics_queue.get_nowait()  # evict oldest to make room
                except Empty:
                    pass  # consumer drained concurrently; next put will succeed
                try:
                    metrics_queue.put_nowait(data)
                except Full:
                    pass  # still full after eviction (race); accept the loss

        # Make a helpful logger name
        if name is None:
            logger_name = f'{__name__}.publisher.{source_id}'
        else:
            logger_name = f'{__name__}.publisher.{name}.{source_id}'

        log = logging.getLogger(logger_name)
        log.addHandler(logging.NullHandler())

        # Set up metrics
        counter = 0
        metrics = Metrics()

        latency_metric_id = metrics.add_element(MetricType.TIMING, name='pipe latency')
        metrics.start(latency_metric_id)

        counting_metric_id = metrics.add_element(MetricType.COUNTING, name='movement count')

        while True:
            # Check for stop
            if stop_event.is_set():
                log.info("Stopping due to stop event")
                _emit_metrics()
                return

            try:
                log.debug("Trying to get line from source in pair %s",
                          source_id)
                line = safe_get(source_queue, timeout=timeout, stop_event=stop_event)

                # Distribute the line to all downstream queues
                for dest_id, dest_queue in dest_queues_dict.items():
                    log.info("Writing line from source %s to dest %s",
                             source_id, dest_id)
                    dest_queue.put(line)

                # Mark that we've finished processing the item from the queue
                if hasattr(source_queue, 'task_done'):
                    source_queue.task_done()

                # Mark latency
                counter += 1
                metrics.lap(latency_metric_id)
                metrics.increment(counting_metric_id)

                # Send metric info
                if counter % metric_interval == 0:
                    _emit_metrics()

            except Empty:
                log.debug("No lines to get from source in pair %s", source_id)
                continue

            except EOFError:
                log.debug("Source in pair %s no longer available", source_id)

            except BrokenPipeError:
                log.debug("One pipe in pair %s is no longer available",
                          name if name else source_id)
                return

    def get_queue(self, queue_id: Union[str, int]) -> UNION_SUPPORTED_QUEUES:
        """Retrieve a client's queue reference

        Args:
            queue_id: ID of the queue reference

        Returns:
            Reference to a queue instance
        """
        if queue_id in self.client_queues_source:
            queue_list = self.client_queues_source
        else:
            queue_list = self.client_queues_destination

        return queue_list[text(queue_id)]

    @validate_direction
    def register_queue(self,
                       q: UNION_SUPPORTED_QUEUES,
                       direction: DIRECTION) -> str:
        """Register a queue to this link

        For a new ``FROM`` (source) queue, a publishing process will be created to send
        all additions down to destination queues.

        For a new ``TO`` (destination) queue, all new additions to source queues
        will be added to this queue.

        Returns the numeric ID for the queue reference, which must be used in all
        future interactions.

        Args:
            q: Queue object to register
            direction: DIRECTION.FROM or DIRECTION.TO

        Returns:
            The client's ID for access to this queue
        """
        # Calculate the opposite direction
        op_direction = DIRECTION.TO if direction == DIRECTION.FROM else DIRECTION.FROM

        # Warnings
        # Priority Queues
        if isinstance(q, tuple(PRIORITY_QUEUES)):
            self._log.warning('Entries in a PriorityQueue must have the correct format (tuple or '
                           'simple class with a priority value) or the queue will block.')

        # SimpleQueues
        if direction == DIRECTION.FROM and isinstance(q, tuple(SIMPLE_QUEUES)):
            self._log.warning('Using multiple readers on a SimpleQueue that is a source for a '
                              'QueueLink instance (including other QueueLink instances) can cause '
                              'a deadlock. Queue type: %s.', type(q))

        with self.queues_lock:
            # Get the queue list and opposite queue list
            queue_dict = getattr(self, f'client_queues_{direction}')
            op_queue_dict = getattr(self, f'client_queues_{op_direction}')

            # Make sure we don't accidentally create a loop, or add multiple
            # times
            if q in queue_dict.values():
                raise ValueError("Cannot add this queue again")

            if q in op_queue_dict.values():
                raise ValueError('This queue is in the opposite list. Cannot'
                                 f' add to the {direction} list because it would cause'
                                 ' a circular reference.')

            # Increment the current queue_id
            queue_id = self.last_queue_id + 1
            self.last_queue_id = queue_id

            # Store the queue proxy in the appropriate queue list
            queue_dict[text(queue_id)] = q

            # (Re)create the publishing processes
            # New source:
            #   Just add a new publishing process and send it the list of
            #   destinations.
            if direction == DIRECTION.FROM:
                self._log.debug("Registering source client")
                source_id = queue_id

                # Create and store a new stop event
                stop_event = self.Event()
                self.publisher_stops[text(source_id)] = stop_event

                # Start the publisher
                # This doesn't do anything unless there are some available
                # destinations
                self._start_publisher(source_id)

            # New destination:
            #   Stop all existing publishing processes and restart with updated
            #   destination queue list.
            else:
                self._log.debug("Registering destination client")

                # Stop current processes
                self._stop_publishers()

                # Restart publishers with the updated destinations
                for source_id in self.client_queues_source:
                    self._log.debug("Starting publisher %s", source_id)
                    self._start_publisher(source_id)

        return text(queue_id)

    def read(self, q: UNION_SUPPORTED_QUEUES) -> str:
        """Register a source queue

        Args:
            q: Queue or proxy object to a queue

        Returns:
            The client's ID for access to this queue
        """
        return self.register_queue(q=q, direction=DIRECTION.FROM)

    def write(self, q: UNION_SUPPORTED_QUEUES) -> str:
        """Register a destination queue

        Args:
            q: Queue or proxy object to a queue

        Returns:
            The client's ID for access to this queue
        """
        return self.register_queue(q=q, direction=DIRECTION.TO)

    def _start_publisher(self, source_id: str):
        """Eliminate duplicated Process() call in register_queue

        Start a publisher process to move items from the source queue to
        the destination queues

        Should only be called by register_queue and unregisterQueue (they hold
        a lock for the queue dictionaries while this runs.)
        """
        stop_event = self.publisher_stops[text(source_id)]
        source_queue = self.client_queues_source[text(source_id)]
        dest_queues_dict = self.client_queues_destination
        destination_names = "-".join(dest_queues_dict.keys())

        # Make sure we don't overwrite a running publisher (Process or Thread)
        try:
            if self.client_pair_publishers[text(source_id)].is_alive():
                raise ValueError("Cannot overwrite a running publisher!")

        # If this is new, there won't be an existing value here
        except KeyError:
            pass

        # Make sure there is at least one desination queue
        if len(dest_queues_dict) == 0:
            return

        # Start a publisher process to move items from the source
        # queue to the destination queues
        self._log.info("Starting queue link for source %s to %s",
                       source_id, destination_names)

        # Name the process
        if self.log_name is not None:
            process_name = f'ClientPublisher-{self.log_name}-{source_id}'
        else:
            process_name = f'ClientPublisher-{source_id}'

        # Determine whether to use a thread or process based publisher
        threaded = is_threaded([source_queue, *dest_queues_dict.values()])

        # Decide whether to use a local thread or process-based publisher
        # pylint: disable=invalid-name
        Parallel = Thread if threaded or self.thread_only else self.Process
        # pylint: enable=invalid-name

        # Thread publishers use the threading queue (no feeder-thread race);
        # process publishers use the multiprocessing queue (cross-process safe).
        metrics_q = (self._metrics_queue_thread
                     if (threaded or self.thread_only)
                     else self._get_process_metrics_queue())

        # Start the publisher
        proc = Parallel(target=self._publisher,
                        name=process_name,
                        kwargs={"name": self.name,
                                "stop_event": stop_event,
                                "source_id": source_id,
                                "source_queue": source_queue,
                                "dest_queues_dict": dest_queues_dict,
                                "timeout": self.link_timeout,
                                "metrics_queue": metrics_q})
        proc.daemon = True
        proc.start()
        self.started.set()
        self.client_pair_publishers[text(source_id)] = proc

    def _stop_publisher(self, source_id: str) -> None:
        """Stop a current publisher process

        Args:
            source_id: ID of the source queue
        """
        self._log.debug("Stopping publisher %s", source_id)
        stop_event = self.publisher_stops[text(source_id)]
        stop_event.set()

        # Wait for the process to stop
        try:
            proc_old = self.client_pair_publishers[text(source_id)]

            while True:
                proc_old.join(timeout=1)

                if proc_old.is_alive():
                    self._log.info("Waiting for client %s to stop", source_id)
                else:
                    break

        # If no current publishers are running, a KeyError will be raised
        except KeyError:
            pass

        # Flip the "stop" event back to normal
        # Clear the stop event so it can be reused if this publisher is restarted
        # (e.g., when a new destination queue is registered)
        stop_event.clear()

        # Metrics queues are intentionally NOT replaced here. They persist across
        # publisher restarts so that data emitted by the stopping publisher remains
        # available to get_metrics(). The bounded maxsize=100 on each queue ensures
        # that restarted publishers' put_nowait calls discard safely if the consumer
        # falls behind.

    def _stop_publishers(self) -> None:
        """Stop all current publisher processes"""
        self._log.debug("Stopping any current publishers")
        for source_id in self.client_queues_source:
            self._stop_publisher(source_id)

    @validate_direction
    def unregister_queue(self, queue_id: Union[str, int], direction: DIRECTION) -> str:
        """Detach a queue from this link

        Returns the clientId that was removed

        Args:
            queue_id: ID of the client
            direction: source or destination

        Returns:
            ID of the client queue
        """
        queue_id_txt = text(queue_id)

        with self.queues_lock:
            queue_list = getattr(self, f'client_queues_{direction}')
            if queue_id_txt in queue_list:
                queue_list.pop(queue_id_txt)

            if direction == DIRECTION.FROM:
                source_id = queue_id

                # Stop the source publisher
                # Wait for the process to stop
                self._stop_publisher(source_id)

                # Remove the stop
                self.publisher_stops.pop(queue_id_txt)

            else:
                # Stop current processes
                self._stop_publishers()

                # Restart publishers with the updated destinations
                for source_id in self.client_queues_source:
                    self._start_publisher(source_id)

        return queue_id_txt

    def is_empty(self, queue_id: Union[str, None]=None) -> bool:
        """Checks whether the source or destination queues are empty

        Args:
            queue_id: (optional) ID of the client

        Returns:
            If queue_id is None, returns True ONLY if ALL queues are empty.
            If queue_id is provided, True ONLY if both main queue and specified client queue are
            empty.
        """
        with self.queues_lock:
            if queue_id is not None:
                self._log.debug("Checking if %s is empty", queue_id)
                empty = True

                # If this is a downstream queue, we need to make sure all
                # upstream queues are empty, too
                if queue_id in self.client_queues_destination:
                    self._log.debug("First checking upstream queue(s) of %s", queue_id)

                    for source_id in self.client_queues_source:
                        source_empty = self.get_queue(source_id).empty()
                        self._log.info("Source %s is empty: %s",
                                       source_id, source_empty)
                        empty = empty and source_empty

                queue_empty = self.get_queue(queue_id).empty()
                empty = empty and queue_empty
                self._log.info("%s is empty: %s", queue_id, queue_empty)

            else:
                empty = True

                for client_queue in list(self.client_queues_source.values()):
                    empty = empty and client_queue.empty()

                for client_queue in list(self.client_queues_destination.values()):
                    empty = empty and client_queue.empty()

            self._log.debug("Reporting queue link empty: %s", empty)
            return empty

    def is_alive(self) -> bool:
        """Whether all of the publishers are alive

        Returns:
            False if any publisher has stopped

        Raises:
            ProcessNotStarted if we've never started a publisher
        """
        alive = True

        if not self.started.is_set():
            raise ProcessNotStarted(f'{self.name} has not started')

        for proc in self.client_pair_publishers.values():
            alive = alive and proc.is_alive()

        return alive

    def is_drained(self, queue_id: str=None):
        """Check alive and empty

        Attempts clean semantic response to "is there, or will there be, data
        to read?"

        :returns bool
        """
        drained = True

        # Don't want to check "alive", as these processes run until a parent
        # processes stops them (or is stopped).
        # If we haven't started yet, we are not drained
        if not self.started.is_set():
            return False

        drained = drained and self.is_empty(queue_id=queue_id)

        return drained

    @validate_direction
    def destructive_audit(self, direction: DIRECTION):
        """Print a line from each client Queue from the provided direction

        Args:
            direction (DIRECTION): source or destination

        This is a destructive operation, as it *removes* a line from each Queue

        :raises Empty
        """
        with self.queues_lock:
            queue_list = getattr(self, f'client_queues_{direction}')

            for q_id in list(queue_list):
                try:
                    target_queue = self.get_queue(q_id)
                    self._log.info("queue_id %s: %s",
                                   text(q_id), safe_get(target_queue, block=False))
                except Empty:
                    self._log.info("queue_id %s is empty", text(q_id))

    def get_metrics(self) -> dict:
        """Retrieve accumulated metrics from all publishers.

        Drains both the thread and process metrics queues and merges all
        snapshots by element_id. Because each snapshot is keyed by element_id,
        later snapshots for the same element overwrite earlier ones, so the
        result reflects the most recent data for each metric element. Returns
        an empty dict if no metrics have been emitted yet.

        :returns dict mapping element_id -> metric data dict
        """
        value = {}

        queues = [self._metrics_queue_thread]
        if self._metrics_queue_process is not None:
            queues.append(self._metrics_queue_process)

        for mq in queues:
            try:
                while True:
                    value.update(mq.get_nowait())
            except Empty:
                pass

        return value
