# The ZeekControl interactive shell.

import logging
import os
import sys

from ZeekControl import cmdresult, config, control, execute, lock, pluginreg, version
from ZeekControl import node as node_mod
from ZeekControl.exceptions import InvalidNodeError, LockError, RuntimeEnvironmentError


class TermUI:
    def info(self, txt):
        print(txt)

    def error(self, txt):
        print(txt, file=sys.stderr)

    warn = error


class NullHandler(logging.Handler):
    def emit(self, record):
        pass


def expose(func):
    func.api_exposed = True
    return func


def lock_required(func):
    def wrapper(self, *args, **kwargs):
        self.lock()
        try:
            return func(self, *args, **kwargs)
        finally:
            self.unlock()

    wrapper.lock_required = True
    return wrapper


def lock_required_silent(func):
    def wrapper(self, *args, **kwargs):
        self.lock(showwait=False)
        try:
            return func(self, *args, **kwargs)
        finally:
            self.unlock()

    wrapper.lock_required = True
    return wrapper


def check_config(func):
    def wrapper(self, *args, **kwargs):
        if config.Config.is_cfg_changed():
            self.ui.warn('Configuration has changed. Run the "deploy" command.')
        return func(self, *args, **kwargs)

    return wrapper


class ZeekCtl:
    def __init__(
        self,
        basedir=version.ZEEKBASE,
        libdir=version.LIBDIR,
        libdirinternal=version.LIBDIRINTERNAL,
        cfgfile=version.CFGFILE,
        zeekscriptdir=version.ZEEKSCRIPTDIR,
        ui=TermUI(),
        state=None,
    ):
        self.ui = ui
        self.zeekbase = basedir
        self.libdir = libdir
        self.libdirinternal = libdirinternal

        self.config = config.Configuration(
            self.zeekbase,
            self.libdir,
            self.libdirinternal,
            cfgfile,
            zeekscriptdir,
            self.ui,
            state,
        )

        # Remove all log handlers (set by any previous calls to logging.*)
        logging.getLogger().handlers = []

        if self.config.debug:
            # Add a log handler that logs to a file.
            try:
                logging.basicConfig(
                    filename=self.config.debuglog,
                    format="%(asctime)s [%(module)s] %(message)s",
                    datefmt=self.config.timefmt,
                    level=logging.DEBUG,
                )
            except OSError as err:
                raise RuntimeEnvironmentError(
                    f"{err}\nCheck if the user running ZeekControl has write access to the debug log file."
                )
        else:
            # Add a log handler that does nothing.
            h = NullHandler()
            logging.getLogger().addHandler(h)

        self.executor = execute.Executor(self.config)
        self.plugins = pluginreg.PluginRegistry()
        self.setup()
        self.controller = control.Controller(
            self.config, self.ui, self.executor, self.plugins
        )

    def setup(self):
        plugindirs = self.config.sitepluginpath.split(":")
        plugindirs.append(self.config.plugindir)
        plugindirs.append(self.config.pluginzeekdir)

        for pdir in plugindirs:
            if pdir:
                self.plugins.addDir(pdir)

        self.plugins.loadPlugins(self.ui, self.executor)
        self.plugins.initPluginOptions()
        self.plugins.addNodeKeys()
        self.config.initPostPlugins()
        self.plugins.initPlugins(self.ui)
        self.plugins.initPluginCmds()
        os.chdir(self.config.zeekbase)
        if self.config.get_state("cronenabled") is None:
            self.config.set_state("cronenabled", True)

    def reload_cfg(self):
        self.config.reload_cfg()

        if self.config.debug:
            if isinstance(logging.getLogger().handlers[0], NullHandler):
                # Remove the null handler and configure logging to a file.
                logging.getLogger().handlers = []
                logging.basicConfig(
                    filename=self.config.debuglog,
                    format="%(asctime)s [%(module)s] %(message)s",
                    datefmt=self.config.timefmt,
                    level=logging.DEBUG,
                )

            # Re-enable all log levels.
            logging.disable(logging.NOTSET)
        else:
            # Disable logging to all log levels that we use.
            logging.disable(logging.CRITICAL)

        self.executor.finish()
        self.plugins.initPluginOptions()
        self.config.initPostPlugins()
        self.plugins.initPlugins(self.ui)
        self.plugins.initPluginCmds()

    def finish(self):
        self.executor.finish()
        self.plugins.finishPlugins()

    def warn_zeekctl_install(self):
        self.config.warn_zeekctl_install()

    # Turns node name arguments into a list of nodes.  If "get_hosts" is True,
    # then only one node per host is chosen.  If "get_types" is True, then
    # only one node per node type (manager, proxy, etc.) is chosen.
    def node_args(self, args=None, get_hosts=False, get_types=False):
        nodes = []

        if args:
            for arg in args.split():
                nodelist = self.config.nodes(arg)
                if not nodelist:
                    raise InvalidNodeError(f"unknown node '{arg}'")

                nodes += nodelist

            # Remove duplicate nodes
            newlist = list(set(nodes))
            if len(newlist) != len(nodes):
                nodes = newlist
        else:
            # Get all nodes.
            nodes = self.config.nodes()

        # Sort the list so that it doesn't depend on initial order of arguments
        nodes.sort(key=node_mod.sortnode)

        if get_hosts:
            hosts = {}
            hostnodes = []
            for node in nodes:
                if node.host not in hosts:
                    hosts[node.host] = 1
                    hostnodes.append(node)

            nodes = hostnodes

        if get_types:
            types = {}
            typenodes = []
            for node in nodes:
                if node.type not in types:
                    types[node.type] = 1
                    typenodes.append(node)

            nodes = typenodes

        return nodes

    def lock(self, showwait=True):
        lockstatus = lock.lock(self.ui, showwait)
        if not lockstatus:
            raise LockError("Unable to get lock")

        self.config.read_state()

    def unlock(self):
        lock.unlock(self.ui)

    def node_names(self):
        return [n.name for n in self.config.nodes()]

    def node_groups(self):
        return node_mod.node_groups()

    @expose
    @check_config
    def nodes(self):
        results = cmdresult.CmdResult()

        if self.plugins.cmdPre("nodes"):
            for n in self.config.nodes():
                results.set_node_data(n, True, n.to_dict())
        else:
            results.ok = False

        self.plugins.cmdPost("nodes")

        return results

    @expose
    @check_config
    def get_config(self):
        results = cmdresult.CmdResult()

        if self.plugins.cmdPre("config"):
            results.keyval = self.config.options()
        else:
            results.ok = False

        self.plugins.cmdPost("config")
        return results

    @expose
    @check_config
    @lock_required
    def install(self, local=False):
        if self.plugins.cmdPre("install"):
            results = self.controller.install(local)
        else:
            results = cmdresult.CmdResult(ok=False)

        self.plugins.cmdPost("install")
        return results

    @expose
    @check_config
    @lock_required
    def start(self, node_list=None):
        nodes = self.node_args(node_list)

        nodes = self.plugins.cmdPreWithNodes("start", nodes)
        results = self.controller.start(nodes)
        self.plugins.cmdPostWithResults("start", results.get_node_data())

        return results

    @expose
    @check_config
    @lock_required
    def stop(self, node_list=None):
        nodes = self.node_args(node_list)

        nodes = self.plugins.cmdPreWithNodes("stop", nodes)
        results = self.controller.stop(nodes)
        self.plugins.cmdPostWithResults("stop", results.get_node_data())

        return results

    @expose
    @check_config
    @lock_required
    def restart(self, clean=False, node_list=None):
        nodes = self.node_args(node_list)

        nodes = self.plugins.cmdPreWithNodes("restart", nodes, clean)

        self.ui.info("stopping ...")
        results = self.stop(node_list)
        if not results.ok:
            return results

        if clean:
            self.ui.info("cleaning up ...")
            results = self.cleanup(node_list=node_list)
            if not results.ok:
                return results

            self.ui.info("checking configurations ...")
            results = self.check(node_list)
            if not results.ok:
                return results

            self.ui.info("installing ...")
            results = self.install()
            if not results.ok:
                return results

        self.ui.info("starting ...")
        results = self.start(node_list)

        self.plugins.cmdPostWithNodes("restart", nodes)
        return results

    @expose
    @lock_required
    def deploy(self):
        if not self.plugins.cmdPre("deploy"):
            results = cmdresult.CmdResult(ok=False)
            return results

        if self.config.is_cfg_changed():
            self.ui.info("Reloading zeekctl configuration ...")
            self.reload_cfg()

        self.ui.info("checking configurations ...")
        results = self.check(check_node_types=True)
        if not results.ok:
            for node, success, output in results.get_node_output():
                if not success:
                    self.ui.info(f"{node} scripts failed.")
                    self.ui.info(output)

            return results

        self.ui.info("installing ...")
        results = self.install()
        if not results.ok:
            return results

        self.ui.info("stopping ...")
        results = self.stop()
        if not results.ok:
            return results

        self.ui.info("starting ...")
        results = self.start()

        self.plugins.cmdPost("deploy")
        return results

    @expose
    @check_config
    @lock_required
    def status(self, node_list=None):
        nodes = self.node_args(node_list)

        nodes = self.plugins.cmdPreWithNodes("status", nodes)
        results = self.controller.status(nodes)
        self.plugins.cmdPostWithNodes("status", nodes)
        return results

    @expose
    @lock_required
    def top(self, node_list=None):
        nodes = self.node_args(node_list)

        nodes = self.plugins.cmdPreWithNodes("top", nodes)
        results = self.controller.top(nodes)
        self.plugins.cmdPostWithNodes("top", nodes)

        return results

    @expose
    @check_config
    @lock_required
    def diag(self, node_list=None):
        nodes = self.node_args(node_list)

        nodes = self.plugins.cmdPreWithNodes("diag", nodes)
        results = self.controller.diag(nodes)
        self.plugins.cmdPostWithNodes("diag", nodes)

        return results

    @expose
    @lock_required_silent
    def cron(self, watch=True):
        if self.plugins.cmdPre("cron", "", watch):
            self.controller.cron(watch)
        self.plugins.cmdPost("cron", "", watch)

        return True

    @expose
    @check_config
    @lock_required
    def cronenabled(self):
        results = False
        if self.plugins.cmdPre("cron", "?", False):
            results = self.config.cronenabled
        self.plugins.cmdPost("cron", "?", False)
        return results

    @expose
    @check_config
    @lock_required
    def setcronenabled(self, enable=True):
        if enable:
            if self.plugins.cmdPre("cron", "enable", False):
                self.config.set_state("cronenabled", True)
                self.ui.info("cron enabled")
            self.plugins.cmdPost("cron", "enable", False)
        else:
            if self.plugins.cmdPre("cron", "disable", False):
                self.config.set_state("cronenabled", False)
                self.ui.info("cron disabled")
            self.plugins.cmdPost("cron", "disable", False)

        return True

    @expose
    @check_config
    @lock_required
    def check(self, node_list=None, check_node_types=False):
        nodes = self.node_args(node_list, get_types=check_node_types)

        nodes = self.plugins.cmdPreWithNodes("check", nodes)
        results = self.controller.check(nodes)
        self.plugins.cmdPostWithResults("check", results.get_node_data())

        return results

    @expose
    @check_config
    @lock_required
    def cleanup(self, cleantmp=False, node_list=None):
        nodes = self.node_args(node_list)

        nodes = self.plugins.cmdPreWithNodes("cleanup", nodes, cleantmp)
        results = self.controller.cleanup(nodes, cleantmp)
        self.plugins.cmdPostWithNodes("cleanup", nodes, cleantmp)

        return results

    @expose
    @check_config
    @lock_required
    def capstats(self, interval=10, node_list=None):
        nodes = self.node_args(node_list)
        nodes = self.plugins.cmdPreWithNodes("capstats", nodes, interval)
        results = self.controller.capstats(nodes, interval)
        self.plugins.cmdPostWithNodes("capstats", nodes, interval)

        return results

    @expose
    @check_config
    @lock_required
    def df(self, node_list=None):
        nodes = self.node_args(node_list, get_hosts=True)
        nodes = self.plugins.cmdPreWithNodes("df", nodes)
        results = self.controller.df(nodes)
        self.plugins.cmdPostWithNodes("df", nodes)

        return results

    @expose
    @check_config
    @lock_required
    def print_id(self, id, node_list=None):
        nodes = self.node_args(node_list)
        nodes = self.plugins.cmdPreWithNodes("print", nodes, id)
        results = self.controller.print_id(nodes, id)
        self.plugins.cmdPostWithNodes("print", nodes, id)

        return results

    @expose
    @check_config
    @lock_required
    def peerstatus(self, node_list=None):
        nodes = self.node_args(node_list)
        nodes = self.plugins.cmdPreWithNodes("peerstatus", nodes)
        results = self.controller.peerstatus(nodes)
        self.plugins.cmdPostWithNodes("peerstatus", nodes)

        return results

    @expose
    @check_config
    @lock_required
    def netstats(self, node_list=None):
        if not node_list:
            node_list = None
            if not self.config.standalone:
                node_list = node_mod.worker_group()

        nodes = self.node_args(node_list)
        nodes = self.plugins.cmdPreWithNodes("netstats", nodes)
        results = self.controller.netstats(nodes)
        self.plugins.cmdPostWithNodes("netstats", nodes)

        return results

    @expose
    @check_config
    def execute(self, cmd):
        nodes = self.node_args(get_hosts=True)

        if self.plugins.cmdPre("exec", cmd):
            results = self.controller.execute_cmd(nodes, cmd)
        else:
            results = cmdresult.CmdResult(ok=False)

        self.plugins.cmdPost("exec", cmd)

        return results

    @expose
    @check_config
    @lock_required
    def scripts(self, check=False, node_list=None):
        nodes = self.node_args(node_list)

        nodes = self.plugins.cmdPreWithNodes("scripts", nodes, check)
        results = self.controller.scripts(nodes, check)
        self.plugins.cmdPostWithNodes("scripts", nodes, check)

        return results

    @expose
    @check_config
    @lock_required
    def process(self, trace, options, scripts):
        if self.plugins.cmdPre("process", trace, options, scripts):
            results = self.controller.process(trace, options, scripts)
        else:
            results = cmdresult.CmdResult(ok=False)

        self.plugins.cmdPost("process", trace, options, scripts, results.ok)

        return results

    @expose
    @check_config
    @lock_required
    def plugincmd(self, cmd, args):
        return self.plugins.runCustomCommand(cmd, args, self.ui)
