from crawler_cluster.utils import get_redis_conn, import_module
from crawler_cluster import BrowserWorker, HTTPWorker
from crawler_cluster.worker import Worker
from file_logger import get_logger
import aioredis
from typing import Dict, Any, Optional
from pprint import pformat
import asyncio
import signal
import pickle
from pickle import PicklingError
from typing import Callable
import inspect
import sys


class QueueWorker:
    """Execute cached functions and cache results."""
    logger = get_logger('crawler-cluster')

    def __init__(self, worker: Worker):
        self.worker = worker
        # key of Redis List where functions waiting to be executed are queued.
        self.function_queue_key = 'waiting_functions'
        # key to Redis List where results of executed functions are cached.
        self.function_results_key = 'function_results'
        # map module name/path to module.
        self._dynamically_imported_modules = {}
        # keep track of task in progress so it can be re-cached in case of failure.
        self._task_in_progres = None
        # state status
        self._launched = False
        self._running = False

    async def run(self) -> None:
        """Launch and execute queued functions.
        If member functions should be called before executing queued functions,
        call launch(), then other member functions, then run_queued_functions().
        """
        await self.launch()
        await self.run_queued_functions()
        self.logger.info("Queue worker is running.")

    async def launch(self) -> None:
        """Initialize connections and open a single browser with a single page."""
        if self._launched:
            self.logger.warning(
                f"Queue worker {self} is already launched!")
            return
        # connect to Redis.
        self.redis_client = await get_redis_conn()
        launch_ok = await self.worker.launch()
        if self.redis_client.closed or not launch_ok:
            await self.shutdown()
        # shutdown on signal interupt.
        self._add_signal_handler()
        self._launched = True

    async def run_queued_functions(self) -> None:
        """Execute all queued functions."""
        if self._running:
            self.logger.warning(f"Queue worker {self} is already running!")
            return
        asyncio.create_task(self._execute_function())
        self._running = True

    @property
    def task_in_progres(self) -> Dict[str, Any]:
        """Return dict of metadata for currently running task."""
        return self._task_in_progres

    async def shutdown(self, sig=None) -> None:
        """Close connections and save incomplete task."""
        if sig is not None:
            self.logger.info(f"Caught signal: {sig.name}")
        self.logger.info("Shutting down.")
        await self.worker.shutdown()
        try:
            if self._task_in_progres is not None:
                self.logger.info(
                    f"Requeuing task: {pformat(self._task_in_progres)}")
                # re-cache currently in progress/incomplete function.
                await self.redis_client.lpush(self.function_queue_key, pickle.dumps(self._task_in_progres))
            self.redis_client.close()
        except Exception as e:
            self.logger.error(f"Error shutting down Redis: {e}")
        # exit so systemd can restart the service.
        sys.exit(1)

    def _get_module(self, module_path: str) -> 'module':
        """Get a dynamically imported module from cache or import it now if not in cache."""
        if module_path in self._dynamically_imported_modules:
            # use already imported module.
            return self._dynamically_imported_modules[module_path]
        module = import_module(module_path)
        if module is not None:
            self._dynamically_imported_modules[module_path] = module
            return module

    def _get_function(self, module_path: str, function_name: str) -> Callable:
        """Use function name to load a function from a module."""
        module = self._get_module(module_path)
        if module is not None:
            try:
                return getattr(module, function_name)
            except AttributeError:
                self.logger.error(
                    f"Could not load function {function_name} from module {module_path}!")

    def _add_signal_handler(self) -> None:
        """Catch signals and shutdown."""
        loop = asyncio.get_running_loop()
        for s in (signal.SIGINT, signal.SIGTERM, signal.SIGHUP):
            loop.add_signal_handler(
                s, lambda s=s: asyncio.create_task(self.shutdown(s)))

    async def _execute_function(self) -> None:
        """Execute earliest-queued function and cache result in Redis."""
        # get earliest-queued function from List.
        task = await self.redis_client.rpop(self.function_queue_key)
        if task is not None:
            task = pickle.loads(task)
            # get function from a dynamically imported module.
            func = self._get_function(task['module'], task['function'])
            self.logger.info(f"Executing function {task}: {func}")
            # mark func_id as in progress so it can be re-queued in case of failure.
            self._task_in_progres = task
            try:
                kwargs = {k: v for k, v in task.items() if k not in (
                    'module', 'function', 'task_id')}
                # check if the function requires reference to the Worker instance
                if 'worker' in inspect.signature(func).parameters:
                    kwargs['worker'] = self.worker
                # execute function.
                task['result'] = await func(**kwargs)
            except Exception as e:
                self.logger.error(
                    f"Error executing user-provided function: {e}. Task: {task}")
                task['result'] = str(e)
                await self.worker.handle_exception(e)
            # check if worker has encountered terminal error during execution of function.
            if not self.worker.ok:
                await self.shutdown()
            try:
                task = pickle.dumps(task)
                # cache function result so it can be retrieved by client.
                await self.redis_client.lpush(self.function_results_key, task)
            except PicklingError as e:
                self.logger.error(
                    f"Failed to pickle function result! Task: {task}. Error: {e}")
            # mark that function is no longer in progress.
            self._task_in_progres = None
        else:
            await asyncio.sleep(0.5)
        # get next function.
        asyncio.create_task(self._execute_function())
