from copy import deepcopy
import threading
import uuid
from datetime import datetime
from typing import (
    Any, Dict, Optional, List, Union, Set, Callable, Type
)


import dbt.exceptions
import dbt.flags as flags
from dbt.adapters.factory import reset_adapters, register_adapter
from dbt.contracts.graph.manifest import Manifest
from dbt.contracts.rpc import (
    LastParse,
    ManifestStatus,
    GCSettings,
    GCResult,
    TaskRow,
    TaskID,
)
from dbt.logger import LogMessage, list_handler
from dbt.parser.manifest import ManifestLoader
from dbt.rpc.error import dbt_error
from dbt.rpc.gc import GarbageCollector
from dbt.rpc.task_handler_protocol import TaskHandlerProtocol, TaskHandlerMap
from dbt.rpc.task_handler import set_parse_state_with
from dbt.rpc.method import (
    RemoteMethod, RemoteManifestMethod, RemoteBuiltinMethod, TaskTypes,
)
# pick up our builtin methods
import dbt.rpc.builtins  # noqa


# import this to make sure our timedelta encoder is registered
from dbt import helper_types  # noqa


WrappedHandler = Callable[..., Dict[str, Any]]


class UnconditionalError:
    def __init__(self, exception: dbt.exceptions.Exception):
        self.exception = dbt_error(exception)

    def __call__(self, *args, **kwargs):
        raise self.exception


class ParseError(UnconditionalError):
    def __init__(self, parse_error):
        exception = dbt.exceptions.RPCLoadException(parse_error)
        super().__init__(exception)


class CurrentlyCompiling(UnconditionalError):
    def __init__(self):
        exception = dbt.exceptions.RPCCompiling('compile in progress')
        super().__init__(exception)


class ManifestReloader(threading.Thread):
    def __init__(self, task_manager: 'TaskManager') -> None:
        super().__init__()
        self.task_manager = task_manager

    def reload_manifest(self):
        logs: List[LogMessage] = []
        with set_parse_state_with(self.task_manager, lambda: logs):
            with list_handler(logs):
                self.task_manager.parse_manifest()

    def run(self) -> None:
        try:
            self.reload_manifest()
        except Exception:
            # ignore ugly thread-death error messages to stderr
            pass


class TaskManager:
    def __init__(self, args, config, task_types: TaskTypes) -> None:
        self.args = args
        self.config = config
        self.manifest: Optional[Manifest] = None
        self._task_types: TaskTypes = task_types
        self.active_tasks: TaskHandlerMap = {}
        self.gc = GarbageCollector(active_tasks=self.active_tasks)
        self.last_parse: LastParse = LastParse(state=ManifestStatus.Init)
        self._lock: flags.MP_CONTEXT.Lock = flags.MP_CONTEXT.Lock()
        self._reloader: Optional[ManifestReloader] = None
        self.reload_manifest()

    def single_threaded(self):
        return flags.SINGLE_THREADED_WEBSERVER or self.args.single_threaded

    def _reload_task_manager_thread(self, reloader: ManifestReloader):
        """This function can only be running once at a time, as it runs in the
        signal handler we replace
        """
        # compile in a thread that will fix up the tag manager when it's done
        reloader.start()
        # only assign to _reloader here, to avoid calling join() before start()
        self._reloader = reloader

    def _reload_task_manager_fg(self, reloader: ManifestReloader):
        """Override for single-threaded mode to run in the foreground"""
        # just reload directly
        reloader.reload_manifest()

    def reload_manifest(self) -> bool:
        """Reload the manifest using a manifest reloader. Returns False if the
        reload was not started because it was already running.
        """
        if not self.set_parsing():
            return False
        if self._reloader is not None:
            # join() the existing reloader
            self._reloader.join()
        # perform the reload
        reloader = ManifestReloader(self)
        if self.single_threaded():
            self._reload_task_manager_fg(reloader)
        else:
            self._reload_task_manager_thread(reloader)
        return True

    def reload_config(self):
        config = self.config.from_args(self.args)
        self.config = config
        reset_adapters()
        register_adapter(config)
        return config

    def add_request(self, request_handler: TaskHandlerProtocol):
        self.active_tasks[request_handler.task_id] = request_handler

    def get_request(self, task_id: TaskID) -> TaskHandlerProtocol:
        try:
            return self.active_tasks[task_id]
        except KeyError:
            # We don't recognize that ID.
            raise dbt.exceptions.UnknownAsyncIDException(task_id) from None

    def _get_manifest_callable(
        self, task: Type[RemoteManifestMethod]
    ) -> Union[UnconditionalError, RemoteManifestMethod]:
        state = self.last_parse.state
        if state == ManifestStatus.Compiling:
            return CurrentlyCompiling()
        elif state == ManifestStatus.Error:
            return ParseError(self.last_parse.error)
        else:
            if self.manifest is None:
                raise dbt.exceptions.InternalException(
                    f'Manifest should not be None if the last parse state is '
                    f'{state}'
                )
            return task(deepcopy(self.args), self.config, self.manifest)

    def rpc_task(
        self, method_name: str
    ) -> Union[UnconditionalError, RemoteMethod]:
        with self._lock:
            task = self._task_types[method_name]
            if issubclass(task, RemoteBuiltinMethod):
                return task(self)
            elif issubclass(task, RemoteManifestMethod):
                return self._get_manifest_callable(task)
            elif issubclass(task, RemoteMethod):
                return task(deepcopy(self.args), self.config)
            else:
                raise dbt.exceptions.InternalException(
                    f'Got a task with an invalid type! {task} with method '
                    f'name {method_name} has a type of {task.__class__}, '
                    f'should be a RemoteMethod'
                )

    def ready(self) -> bool:
        with self._lock:
            return self.last_parse.state == ManifestStatus.Ready

    def set_parsing(self) -> bool:
        with self._lock:
            if self.last_parse.state == ManifestStatus.Compiling:
                return False
            self.last_parse = LastParse(state=ManifestStatus.Compiling)
        return True

    def parse_manifest(self) -> None:
        self.manifest = ManifestLoader.get_full_manifest(self.config, reset=True)

    def set_compile_exception(self, exc, logs=List[LogMessage]) -> None:
        assert self.last_parse.state == ManifestStatus.Compiling, \
            f'invalid state {self.last_parse.state}'
        self.last_parse = LastParse(
            error={'message': str(exc)},
            state=ManifestStatus.Error,
            logs=logs
        )

    def set_ready(self, logs=List[LogMessage]) -> None:
        assert self.last_parse.state == ManifestStatus.Compiling, \
            f'invalid state {self.last_parse.state}'
        self.last_parse = LastParse(
            state=ManifestStatus.Ready,
            logs=logs
        )

    def methods(self) -> Set[str]:
        with self._lock:
            return set(self._task_types)

    def currently_compiling(self, *args, **kwargs):
        """Raise an RPC exception to trigger the error handler."""
        raise dbt_error(dbt.exceptions.RPCCompiling('compile in progress'))

    def compilation_error(self, *args, **kwargs):
        """Raise an RPC exception to trigger the error handler."""
        raise dbt_error(
            dbt.exceptions.RPCLoadException(self.last_parse.error)
        )

    def get_handler(
        self, method, http_request, json_rpc_request
    ) -> Optional[Union[WrappedHandler, RemoteMethod]]:
        # get_handler triggers a GC check. TODO: does this go somewhere else?
        self.gc_as_required()

        if method not in self._task_types:
            return None

        task = self.rpc_task(method)

        return task

    def task_table(self) -> List[TaskRow]:
        rows: List[TaskRow] = []
        now = datetime.utcnow()
        with self._lock:
            for task in self.active_tasks.values():
                rows.append(task.make_task_row(now))
        return rows

    def gc_as_required(self) -> None:
        with self._lock:
            return self.gc.collect_as_required()

    def gc_safe(
        self,
        task_ids: Optional[List[uuid.UUID]] = None,
        before: Optional[datetime] = None,
        settings: Optional[GCSettings] = None,
    ) -> GCResult:
        with self._lock:
            return self.gc.collect_selected(
                task_ids=task_ids, before=before, settings=settings,
            )