from typing import Any, Callable, Dict, List, Optional, Tuple
from types import ModuleType

import json
import sys
import tqdm
import re
import logging
import datetime
from pathlib import Path

from solitude import TOOL_NAME, hookspecs
from solitude.cache import Cache
from solitude.commands import CommandBase
import solitude.utils.ssh as ssh
import solitude.plugins.default
import solitude.plugins.csubmit
import solitude.plugins.interactive

from multiprocessing.pool import ThreadPool

import pluggy

from solitude.config import Config
from solitude.slurmjob import SlurmJob

logger = logging.getLogger(f"{TOOL_NAME}")


COMMENT_MARKER = "#"
USER_MACRO = "{user}"
SOL_DOCKER_REPO_MACRO = "{sol-docker-repo}"
# these are all aliases for the same repository
SOL_DOCKER_REPO_ALIASES = ["doduo1.umcn.nl", "doduo2.umcn.nl", "doduo.umcn.nl"]
# this is the repository that will be used at core.run
SOL_DOCKER_REPO = "doduo1.umcn.nl"


def resolve_cmd_files(cmd_files: List[Path]) -> List[str]:
    commands: List[str] = []
    for cmd_file in cmd_files:
        commands = commands + resolve_cmd_file(cmd_file=cmd_file)
    return commands


def resolve_cmd_file(cmd_file: Path) -> List[str]:
    try:
        with open(cmd_file, "r") as f:
            data = f.read()
        # auto-replaces hard-coded SOL_DOCKER_REPO_ALIASES to the SOL_DOCKER_REPO_MACRO
        # this makes started jobs still link to the macro under the hood
        # and facilitates transition to using the macro
        for repo_url in SOL_DOCKER_REPO_ALIASES:
            data = data.replace(repo_url, SOL_DOCKER_REPO_MACRO)
        # extract commands per line
        data_lines = [line.strip() for line in data.splitlines(keepends=False)]
        return [d for d in data_lines if any(d) and d[0] != COMMENT_MARKER]
    except FileNotFoundError:
        logger.error(
            f"ERROR: Could not find command file: {cmd_file}, skipping..."
        )
        return []


def shared_section(
    cmd_files: List[Path], cache_file: Path, cfg: Config
) -> Tuple[List[CommandBase], Cache, Optional[Config.SSHConfig]]:
    logger.debug(f"Starting action at: {datetime.datetime.now()}")
    # load plugins
    plugin_manager = setup_plugins(plugins=cfg.plugins)
    logger.debug(f"Read {len(plugin_manager.get_plugins())} command plugins")
    # load cache file
    cache: Cache = Cache(data_class=SlurmJob, file_name=cache_file)
    # perform cache auto cleanup (removes expired jobs)
    cache.cleanup()
    # load command files
    cmds_strs = resolve_cmd_files(cmd_files=cmd_files)
    cmds: List[CommandBase] = read_commands(
        commands=cmds_strs, cache=cache, plugin_manager=plugin_manager,
    )
    # retrieve ssh credentials
    return cmds, cache, cfg.ssh


class SpecialFormatter(logging.Formatter):
    def __init__(self, *args, **kwargs):
        super(SpecialFormatter, self).__init__(*args, **kwargs)
        self._main_formatter = logging.Formatter(fmt="%(message)s")
        self._ssh_formatter = logging.Formatter(fmt="[SSHClient] %(message)s")
        self._special_formatter = logging.Formatter(
            fmt="%(asctime)s [%(name)-10s] %(message)s",
            datefmt="%y-%m-%d %H:%M:%S",
        )

    def format(self, record: logging.LogRecord) -> str:
        if record.name == TOOL_NAME:
            return self._main_formatter.format(record=record)
        elif record.name == "SSHClient":
            return self._ssh_formatter.format(record=record)
        else:
            return self._special_formatter.format(record=record)


def log_setup(log_level: int = logging.DEBUG) -> logging.Logger:
    """setup some basic logging"""
    log = logging.getLogger("")
    log.setLevel(log_level)
    fmt = SpecialFormatter()
    ch = logging.StreamHandler(sys.stdout)
    ch.setLevel(log_level)
    ch.setFormatter(fmt)
    log.addHandler(ch)
    logging.getLogger("paramiko.transport").setLevel(logging.ERROR)
    logging.getLogger("urllib3.connectionpool").setLevel(logging.ERROR)
    return logging.getLogger(TOOL_NAME)


def get_attr_from_module(pclass: str) -> Any:
    splitted: List[str] = pclass.rsplit(".", 1)
    mod = __import__(splitted[0], fromlist=[str(splitted[1])])
    return getattr(mod, splitted[1])


def parse_plugins(plugins: List[str]) -> List[ModuleType]:
    results = []
    for plugin_name in plugins:
        if len(plugin_name.strip()) > 0:
            try:
                results.append(get_attr_from_module(plugin_name.strip()))
            except ModuleNotFoundError:
                logger.error(
                    f"!WARNING! Could not find plugin listed in config: {plugin_name}"
                )
    return results


def match_plugin(pm: pluggy.PluginManager, cmd: str) -> ModuleType:
    try:
        idx = pm.hook.matches_command(cmd=cmd).index(True)
    except ValueError:
        raise ValueError(
            f"No plugin found for handling this type of command: {cmd}"
        )
    return pm.hook.matches_command.get_hookimpls()[::-1][idx].plugin


def read_commands(
    commands: List[str], cache: Cache, plugin_manager: pluggy.PluginManager
) -> List[CommandBase]:
    return [
        CommandBase(
            cmd=x, plugin=match_plugin(pm=plugin_manager, cmd=x), cache=cache
        )
        for x in commands
    ]


def execute_commands(
    commands: List[CommandBase],
    map_fn: Callable,
    workers: int,
    show_progress_bar: bool = True,
):
    if workers > 0:
        tpex = ThreadPool(processes=workers)
        list(
            tqdm.tqdm(
                tpex.imap(map_fn, commands),
                total=len(commands),
                desc="Processing cmds ({} threads)".format(workers),
                disable=not show_progress_bar,
            )
        )
    else:
        for cmd in tqdm.tqdm(
            commands,
            desc="Processing cmds (no threads)",
            disable=not show_progress_bar,
        ):
            map_fn(cmd)
    if show_progress_bar:
        logger.info("")


def read_cache(file_name: Path) -> Dict:
    if not file_name.resolve().is_file():
        return {}
    with open(file_name, "r") as f:
        data = json.load(f)
    return data


def extend(
    cmd_files: List[Path],
    jobids: List[int],
    workers: int,
    cfg: Config,
    cache_file: Path,
):
    _run_command_on_running_job(
        cmd_files,
        jobids,
        workers,
        cfg,
        cache_file,
        command="c-extend",
        action_str="EXTENDING",
    )


def stop(
    cmd_files: List[Path],
    jobids: List[int],
    workers: int,
    cfg: Config,
    cache_file: Path,
):
    _run_command_on_running_job(
        cmd_files,
        jobids,
        workers,
        cfg,
        cache_file,
        command="c-stop",
        action_str="STOPPING",
    )


def link(
    cmd_files: List[Path],
    jobid: int,
    slurm_jobid: int,
    cfg: Config,
    cache_file: Path,
) -> bool:
    commands, cache, _ = shared_section(
        cmd_files=cmd_files, cache_file=cache_file, cfg=cfg
    )
    cmd = commands[jobid - 1]
    # link jobid to slurm_jobid
    if SlurmJob.check_if_job_exists(jobid=slurm_jobid):
        job = SlurmJob(jobid=slurm_jobid)
        job.update()
        cache.update(key=cmd.hash, value=job)
        return True
    return False


def unlink(cmd_files: List[Path], jobid: int, cfg, cache_file: Path) -> bool:
    commands, cache, _ = shared_section(
        cmd_files=cmd_files, cache_file=cache_file, cfg=cfg
    )
    cmd = commands[jobid - 1]
    if cmd.hash in cache:
        cache.delete(key=cmd.hash)
        return True
    return False


def _run_command_on_running_job(
    cmd_files: List[Path],
    jobids: List[int],
    workers: int,
    cfg: Config,
    cache_file: Path,
    command: str,
    action_str: str,
):
    ssh_client = ssh.SSHClient()
    assert cfg.ssh is not None
    ssh_client.connect(*cfg.ssh)

    commands, _, credentials = shared_section(
        cmd_files=cmd_files, cache_file=cache_file, cfg=cfg
    )

    # update selected commands only (use workers)
    subcmds = [cmd for idx, cmd in enumerate(commands) if idx + 1 in jobids]

    def _update_cmd(cmd):
        cmd.update_job_info()

    execute_commands(
        subcmds, map_fn=_update_cmd, workers=workers, show_progress_bar=True
    )

    for idx, cmd in zip(jobids, subcmds):
        if cmd.is_running():
            assert cmd.job_info is not None
            logger.info(f"ATTEMPT {action_str} Job: {idx}")
            ssh_client.exec_command(f"./{command} {cmd.job_info.id}")
        else:
            logger.info(f"Job {idx} is no longer running. Skipping...")
    log_action_finished()


def list_jobs(
    cmd_files: List[Path],
    jobids: List[int],
    workers: int,
    cfg: Config,
    cache_file: Path,
    selected_only: bool = False,
    show_progress_bar: bool = True,
):
    commands, _, _ = shared_section(
        cmd_files=cmd_files, cache_file=cache_file, cfg=cfg
    )

    # update selected commands only (use workers)
    selected_ids = jobids
    if not selected_only:
        jobids = [e for e in range(1, len(commands) + 1)]
    subcmds = [
        cmd
        for idx, cmd in enumerate(commands)
        if (idx + 1 in jobids) or not selected_only
    ]

    def _update_cmd(cmd):
        cmd.update()

    execute_commands(
        subcmds,
        map_fn=_update_cmd,
        workers=workers,
        show_progress_bar=show_progress_bar,
    )
    logger.info("======Commands======")
    logger.info(CommandBase.header_str())
    for idx, cmd in zip(jobids, subcmds):
        logger.info(f'{idx:2d}{"*" if idx in selected_ids else " "} {cmd}')
    log_action_finished()


def log_action_finished():
    logger.debug(f"Finished action at: {datetime.datetime.now()}")


def run(
    cmd_files: List[Path],
    jobids: List[int],
    workers: int,
    user: str,
    priority: str,
    reservation: str,
    duration: int,
    ignore_errors: bool,
    cfg: Config,
    cache_file: Path,
):
    ssh_client = ssh.SSHClient()
    assert cfg.ssh is not None
    ssh_client.connect(*cfg.ssh)

    commands, cache, credentials = shared_section(
        cmd_files=cmd_files, cache_file=cache_file, cfg=cfg
    )

    def _update_cmd(cmd):
        cmd.update_job_info()

    execute_commands(
        commands, map_fn=_update_cmd, workers=workers, show_progress_bar=False
    )

    high_priority_job_available = not any(
        [
            c.job_info is not None
            and c.is_running()
            and c.job_info.priority == "high"
            for c in commands
            if c.job_info is not None and c.job_info.user == user
        ]
    )
    if high_priority_job_available:
        logger.info(f"No high priority job running for user {user}!")
    elif priority == "high":
        logger.info(
            f"No high priority job available for user {user}! Switching to low priority..."
        )
        priority = "low"

    # update selected commands only (use workers)
    subcmds = [cmd for idx, cmd in enumerate(commands) if idx + 1 in jobids]

    def _update_cmd2(cmd):
        cmd.update_state()
        cmd.update_errors()

    execute_commands(
        subcmds, map_fn=_update_cmd2, workers=workers, show_progress_bar=True
    )

    for idx, cmd in zip(jobids, subcmds):
        logger.info(f"ATTEMPT EXECUTE Job {idx}")
        if cmd.is_running():
            logger.info("Job is already running. skipping...")
            logger.info(str(cmd))
        elif cmd.is_finished():
            logger.info("Job has already finished. skipping...")
            logger.info(str(cmd))
        elif cmd.is_erroneous() and not ignore_errors:
            logger.info(
                "Job has halted with errors. Use --ignore-errors to run anyway. skipping..."
            )
            logger.info(str(cmd))
        else:
            try:
                # set correct priority level in command
                cmdstr = cmd.cmd
                if "c-submit " in cmdstr:
                    # remove any priority optional flags
                    cmdstr = re.sub(r"--priority=[a-zA-Z0-9]*\s", "", cmdstr)
                    # add desired priority flag after c-submit command
                    cmdstr = re.sub(
                        r"c-submit",
                        "c-submit --priority={}".format(priority),
                        cmdstr,
                    )
                    if duration is not None:
                        cmdstr = re.sub(
                            r"\s(\d+)\soni:11500",
                            " {} oni:11500".format(str(duration)),
                            cmdstr,
                        )
                    if reservation is not None:
                        cmdstr = re.sub(
                            r"c-submit",
                            "c-submit --reservation={}".format(reservation),
                            cmdstr,
                        )
                cmdstr = cmdstr.replace(USER_MACRO, user)
                cmdstr = cmdstr.replace(SOL_DOCKER_REPO_MACRO, SOL_DOCKER_REPO)

                # attempt schedule command
                result = ssh_client.exec_command(cmdstr)
                if result is not None and ("c-submit " in cmdstr):
                    link_id = int(result[0].strip().split(" ")[-1])
                    # can have only one high priority job so switch to low priority afterwards
                    cache.update(
                        key=cmd.hash,
                        value=SlurmJob(
                            jobid=link_id, user=user, priority=priority
                        ),
                    )
                    if priority == "high":
                        logger.info(
                            "Assigned a high priority job switching to low priority now"
                        )
                        priority = "low"
                    logger.info("http://oni:11080/show/{}".format(link_id))
            except Exception as e:
                raise e

    log_action_finished()


def setup_plugins(
    plugins: List[str], debug: bool = False
) -> pluggy.PluginManager:
    plugin_manager = pluggy.PluginManager(TOOL_NAME)
    plugin_manager.add_hookspecs(hookspecs)
    # Add default fallback plugins
    plugin_manager.register(solitude.plugins.default)
    plugin_manager.register(solitude.plugins.csubmit)
    plugin_manager.register(solitude.plugins.interactive)
    # Load externally packaged plugins
    plugin_manager.load_setuptools_entrypoints(group=TOOL_NAME)
    # Add additional plugins from cfg
    plugins_mods: List[ModuleType] = parse_plugins(plugins=plugins)
    for plugin in reversed(plugins_mods):
        plugin_manager.register(plugin)
    # validate plugins (must strictly implement all plugin methods)
    plugin_manager.check_pending()
    for plugin in plugin_manager.get_plugins():
        validate_plugin(plugin)
    if debug:
        plugin_manager.trace.root.setwriter(print)
        plugin_manager.enable_tracing()
    return plugin_manager


def validate_plugin(plugin: ModuleType):
    required_fns = [
        e for e in dir(solitude.plugins.default) if "__" not in e and "_" in e
    ]
    for fn in required_fns:
        if not hasattr(plugin, fn):
            raise NotImplementedError(
                f"Plugin {plugin} has no hook_implementation for {fn}"
            )
