import json
import math
import signal
import time
import traceback
import uuid

import redis

from .events import TaskCompleteEvent, TaskErrorEvent
from .population.controls.static import StaticControl
from .population.spawners.static import StaticSpawner
from .tasks import BaseTask

__all__ = ['BaseWorker']


class BaseWorker:
    """
    A base class for defining worker behaviour.
    """

    def __init__(
        self,
        conn,
        task_cls,
        broadcast_channel='swm_events',
        max_spawn_time=0,
        max_status_interval=600,
        sleep_interval=1,
        idle_lifespan=0,
        population_control=None,
        population_spawner=None,
        node_id=None
    ):

        # The redis connection (instance) to use for communicating with other
        # workers, servers and the monitor/manager.
        self._conn = conn

        # The class of tasks the worker handles
        self._task_cls = task_cls

        # The channel that task events (e.g complete, error) will be broadcast
        # over.
        self._broadcast_channel = broadcast_channel

        # The maximum amount of time allocated to a worker to spawn new
        # workers. If set to 0 the worker will wait forever.
        self._max_spawn_time = max_spawn_time

        # The maximum number of seconds between status updates before the
        # worker is considered to have died.
        self._max_status_interval = max_status_interval

        # The interval between an idle worker checking between tasks
        self._sleep_interval = sleep_interval

        # The workers unique Id (acquired when started)
        self._id = None

        # The maximum period of time a worker can be idle for before it
        # attempts to shut itself down. If 0 the worker will never attempt to
        # shut itself down.
        self._idle_lifespan = idle_lifespan

        # The time the worker was last active (was started or completed a
        # task).
        self._idle_since = None

        # The population control used to regulate the size of the worker
        # population.
        self._population_control = population_control or StaticControl()

        # The population spawner used to spawn new workers
        self._population_spawner = population_spawner or StaticSpawner()

        # A unique ID to identify the node (hardware device) this worker is
        # running on.
        self._node_id = str(node_id or uuid.getnode())

    def __str__(self):
        return self._id

    @property
    def id(self):
        return self._id

    @property
    def node_id(self):
        return self._node_id

    def do_task(self, task):
        """Do the given task"""

    def get_tasks(self):
        """Return a list of all relevant tasks for the worker"""

        task_cls_list = self._task_cls
        if isinstance(task_cls_list, BaseTask):
            task_cls_list = [task_cls_list]

        tasks_map = {}

        for task_cls in task_cls_list:

            task_prefix = task_cls.get_id_prefix()
            task_ids = set([
                id for id in self._conn.scan_iter(f'{task_prefix}:*')
            ])

            if not task_ids:
                continue

            tasks = [
                task_cls.loads(t)
                for t in self._conn.mget(task_ids) if t
            ]
            tasks_map.update(zip(task_ids, tasks))

        return tasks_map

    def get_workers(self):
        """Return a list of registered workers"""
        return set([
            id for id in self._conn.scan_iter(f'{self.get_id_prefix()}:*')
        ])

    def get_node_workers(self, workers):
        """
        Filter a list of workers so that it only contains works on for this
        workers node.
        """
        nodeset = self._conn.smembers(
            f'{self.get_node_prefix()}{self.node_id}'
        )
        return [worker for worker in workers if worker in nodeset]

    def start(self):
        """Start the worker"""

        # Set a unique Id for the worker
        self._id = f'{self.get_id_prefix()}:{uuid.uuid4()}'

        # Register the worker
        self._conn.setex(self._id, self._max_status_interval, 'idle')
        self._conn.sadd(f'{self.get_node_prefix()}{self.node_id}', self._id)

        # Start the main loop
        signal.signal(signal.SIGTERM, self.shut_down)

        try:
            self._idle_since = time.time()
            self._loop()

        except KeyboardInterrupt:
            pass

        finally:
            self.shut_down()

    def shut_down(self, *args):
        """Shutdown the worker"""

        # Unregister the worker
        self._conn.delete(self._id)
        self._conn.srem(f'{self.get_node_prefix()}{self.node_id}', self._id)

    # Event broadcasting and error handling

    def on_complete(self, task_id, data):
        """Broadcast a task completed event"""
        event = TaskCompleteEvent(task_id, data)
        self._conn.publish(self._broadcast_channel, event.dumps())

    def on_error(self, task_id, error):
        """Broadcast a task error event"""
        event = TaskErrorEvent(task_id, str(error))
        self._conn.publish(self._broadcast_channel, event.dumps())

    def on_spawn_error(self, error):
        """Override this method to manage spawn errors"""
        print(traceback.format_exc())
        print(error)

    # Private

    def _acquire_lock(self, key, value, watch_key=None):
        """Attempt to require a lock"""

        watch_key = key if watch_key is None else watch_key

        with self._conn.pipeline() as multi:

            try:
                multi.watch(watch_key)
                acquired = multi.setnx(key, value)
                multi.execute()

                if not acquired:
                    return False

            except redis.WatchError:

                # Unable to aquire the lock
                return False

            finally:
                multi.reset()

        return True

    def _loop(self):
        """The application loop"""

        while True:

            # Check to see if the shutdown flag is present for the worker
            # class.
            shutdown_key = self._conn.get(self.get_shutdown_key())
            if shutdown_key in ['__shutdown__', self.node_id]:
                return self.shut_down()

            # Get a lists of workers and the tasks to complete
            workers = self.get_workers()
            node_workers = self.get_node_workers(workers)
            tasks = self.get_tasks()

            # If the worker is no longer registered then the main loop should
            # be exited.
            if self._id not in workers:
                return

            # Population check (growth and culling)
            population_change = self._population_control.population_change(
                workers,
                node_workers,
                tasks
            )
            population_lock_key = self.get_population_lock_key()
            time_idle = time.time() - self._idle_since

            # self._conn.delete(population_lock_key)

            if population_change > 0:

                # Attempt to aquire a lock so we can spawn new workers
                if self._acquire_lock(population_lock_key, self.id):
                    try:
                        self._population_spawner.spawn(population_change)

                        # Wait until the new workers have spawned
                        spawned_at = time.time()
                        while True:
                            added_workers = self.get_workers() - workers

                            if len(added_workers) >= population_change:
                                break

                            spawning = time.time() - spawned_at
                            if spawning > (self._max_spawn_time or math.inf):
                                break

                            time.sleep(1)

                    except Exception as error:

                        # Report the error
                        self.on_spawn_error(error)

                    finally:
                        self._conn.delete(population_lock_key)

            elif population_change < 0 \
                    and time_idle > (self._idle_lifespan or math.inf):

                # Attempt to aquire lock so we can shut this worker down
                if self._acquire_lock(population_lock_key, self.id):
                    try:
                        return self.shut_down()

                    finally:
                        self._conn.delete(population_lock_key)

            # Find any tasks that are currently assigned to workers that are
            # no longer registered and convert them to pending tasks.
            for task in tasks.values():

                # Check for any task that is assigned to a non-existent worker
                if task.assigned_to and task.assigned_to not in workers:

                    if self._conn.setnx(task.reclaim_key, self._id):

                        # Clear the existing lock for the task
                        self._conn.expire(task.reclaim_key, 10)
                        self._conn.delete(task.lock_key)

                        # Unassign the task so it will be handled by this
                        # worker.
                        task.unassign()

                # Check if the task is pending
                if not task.assigned_to and self._conn.exists(task.id):

                    # Acquire a lock to perform the task
                    acquired = self._acquire_lock(
                        task.lock_key,
                        self.id,
                        task.id
                    )
                    if not acquired:
                        continue

                    # Assign the task to this worker
                    task.assign_to(self._id)
                    self._conn.set(task._id, task.dumps())

                    # Update the workers status to busy
                    self._conn.setex(
                        self._id,
                        self._max_status_interval,
                        'busy'
                    )

                    # Process the task
                    try:
                        event_data = self.do_task(task)
                        self.on_complete(task.id, event_data)

                    except Exception as e:
                        self.on_error(task.id, e)

                    finally:

                        # Delete the task and associated lock
                        self._conn.delete(task.id)
                        self._conn.delete(task.lock_key)

                        # Record the time at which the worker became idle
                        self._idle_since = time.time()

            # Update the workers status to idle
            self._conn.setex(self._id, self._max_status_interval, 'idle')

            time.sleep(self._sleep_interval)

    # Class methods

    @classmethod
    def get_id_prefix(cls):
        """Return a prefix applied to worker Ids"""
        return 'swm_worker'

    @classmethod
    def get_node_prefix(cls):
        """Return a prefix applied to worker Ids"""
        return f'{cls.get_id_prefix()}_node'

    @classmethod
    def get_population_lock_key(cls):
        """
        Return the lock key for this class or workers to use when managing the
        population.
        """
        return f'{cls.get_id_prefix()}_population_contol'

    @classmethod
    def get_shutdown_key(cls):
        """Return the key used to flag to workers that they should shutdown"""
        return f'{cls.get_id_prefix()}_shutdown'
