#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from argparse import ArgumentParser
import json
import logging as log
from pathlib import Path
from pprint import pformat
from random import uniform
from subprocess import run, PIPE, Popen, STDOUT
import sys
from threading import Thread, Event
from typing import List, Optional, Dict
from urllib.parse import urlparse
from urllib.request import urlretrieve
from venv import EnvBuilder

from legiond.config import LegionConfig

RUNNER_FORMAT = """
# FILE AUTO-GENERATED BY LEGIOND, DO NOT EDIT DIRECTLY, CHANGES WILL BE OVERIDDEN
# Exit Codes:
#  0: success
#  1: unknown error
from argparse import ArgumentParser
import json
import sys
from time import sleep, time
from traceback import format_exc

from robotnikmq import Topic, Message

from {script_name} import run

CONFIG_JSON = '''
{config_json}
'''

parser = ArgumentParser(description="Legion Alerting-Monitoring System Reporter Unit: "
                                    "{reporter_name}")
parser.add_argument('--once', action='store_true', help='Run the reporter only once, not on a loop')
args = parser.parse_args()

failure_counter = 0

try:
    while 42:
        try:
            start = int(time())
            print("Running...", flush=True)
            run(json.loads(CONFIG_JSON))
            print("Run complete.", flush=True)
            failure_counter = 0
        except:
            exc_str = format_exc()
            print("Run function failure detected: ", flush=True)
            print(exc_str, flush=True)
            failure_counter += 1
            priority = 'WARNING'
            if failure_counter >= {critical_threshold}:
                priority = 'CRITICAL'
            elif failure_counter >= {error_threshold}:
                priority = 'ERROR'
            Topic(exchange='{meta_queue}',
                  on_conn_error=lambda a : print(f'RobotnikMQ Connection Failure: '
                                                  f'{reporter_name}'))\
                .broadcast(Message(contents={{
                    'name': 'Reporter script failure: {script_name}',
                    'priority': priority,
                    'description': 'The reporter script {script_name} has failed ' + 
                                    str(failure_counter) + 
                                    ' times in a row. The last exception stack trace was: ' +
                                    exc_str
                }})) # TODO: Convert to a real alert using legion-utils
                # TODO: Alert priority should be configurable based on failure counter
                # TODO: 
        if args.once:
            break
        print(f'Sleeping for {delay}s', flush=True)
        sleep({delay})
except KeyboardInterrupt:
    print("Shutdown signal received.")
"""

CONFIGURATION_PATHS = [Path.cwd() / 'legiond.conf',
                       Path.home() / '.config' / 'legiond' / 'legiond.conf',
                       Path('/etc/legiond/legiond.conf',
                       Path('/etc/opt/legiond/legiond.conf'))]
log.basicConfig(stream=sys.stdout, level=log.DEBUG,
                format='%(levelname)s %(asctime)s [%(threadName)s]: %(message)s',
                datefmt='%Y-%m-%d %H:%M:%S')


def script_of(file_url) -> Path:
    return Path(urlparse(file_url).path)


class ReporterThread(Thread):
    def __init__(self,
                 reporter_config: LegionConfig.Reporter,
                 env_exe: str,
                 runner_script: Path,
                 cwd: Path,
                 once: bool = False,
                 timeout: Optional[int] = None):
        super().__init__(name=reporter_config.name)
        self.reporter_config = reporter_config
        self.once = once
        self.timeout = timeout
        self._stopped = Event()
        self.env_exe = env_exe
        self.runner_script = runner_script
        self.cwd = cwd

    def stop(self):
        log.info(f"Stop signal sent to: {self.reporter_config.name}")
        self._stopped.set()

    def stopped(self):
        return self._stopped.is_set()

    @property
    def cmd(self) -> List[str]:
        return [self.env_exe,
                str(self.runner_script)] + (['--once'] if self.once else [])

    def run(self) -> None:
        if not self.once:
            startup_delay = uniform(0, self.reporter_config.delay)
            log.info('Starting reporter thread after {:.2f}s'.format(startup_delay))
            self._stopped.wait(startup_delay)
        with Popen(self.cmd, stdout=PIPE, stderr=STDOUT, cwd=self.cwd) as proc:
            if proc is not None:
                while proc.poll() is None and not self.stopped():
                    stdout = proc.stdout
                    if stdout:
                        output = stdout.readline().decode()
                        if output == '' and proc.poll() is not None:
                            break
                        if output:
                            log.info(output.strip('\n'))
                proc.terminate()
                proc.wait()
                return
        raise RuntimeError("Unable to start reporter subprocess")


class ReporterVenv:
    UPGRADE_PREFIXES = {'git', '/', 'http'}
    COMMON_REQS = {'pip': 'pip',
                   'robotnikmq': 'robotnikmq'}
                   # 'legion-utils': 'legion-utils'}

    def __init__(self,
                 reporter_dir: Path,
                 reporter_url: str,
                 requirements: List[str],
                 common_reqs_overrides: Optional[Dict[str, str]] = None):
        self.common_reqs = {**self.COMMON_REQS, **(common_reqs_overrides or {})}
        self.reporter_dir = reporter_dir
        self.venv_dir = self.reporter_dir / 'venv'
        self.reporter_url = reporter_url
        builder = EnvBuilder(system_site_packages=True,
                             clear=False,
                             symlinks=True,
                             upgrade=True,
                             with_pip=True,
                             # upgrade_deps=True, # TODO: add with Python 3.9
                             prompt=self.reporter_dir.stem[:-11])
        builder.create(self.venv_dir)
        self.context = builder.ensure_directories(self.venv_dir)
        # Upgrade pip # TODO: Remove in 3.9
        run([self.context.bin_path + '/pip3', 'install', '--upgrade', 'pip'], check=True)
        self.reporter_file = self.download_reporter()
        for req_package in list(self.common_reqs.values()) + requirements:
            self.pip_install(req_package)

    def pip_install(self, package: str) -> 'ReporterVenv':
        # If a package is coming from git, http server, or file we should force a reinstall
        if any(package.startswith(prefix) for prefix in self.UPGRADE_PREFIXES):
            run([self.context.bin_path + '/pip3', 'install', '--upgrade', '--force-reinstall',
                 package], check=True)
        else:
            # Otherwise, a package version would need to change for a reinstall to happen
            run([self.context.bin_path + '/pip3', 'install', package], check=True)
        return self

    @property
    def runner_script(self) -> Path:
        return self.reporter_dir / 'run_reporter.py'

    def create_runner(self,
                      reporter_config: LegionConfig.Reporter,
                      runner_fmt: str = RUNNER_FORMAT) -> 'ReporterVenv':
        with (self.runner_script).open('w+') as script:
            script.write(
                runner_fmt.format(
                    reporter_name=reporter_config.name,
                    script_name=script_of(reporter_config.file_url).stem,
                    config_json=json.dumps(reporter_config.config),
                    delay=reporter_config.delay,
                    meta_queue=reporter_config.meta_queue,
                    critical_threshold=reporter_config.failure_threshold_critical,
                    error_threshold=reporter_config.failure_threshold_error,))
        return self

    def download_reporter(self) -> Path:
        script_path = self.reporter_dir / script_of(self.reporter_url).name
        return Path(urlretrieve(self.reporter_url, script_path)[0])

    def reporter(self, reporter_config: LegionConfig.Reporter,
                 once: bool = False,
                 timeout: Optional[int] = None) -> ReporterThread:
        return ReporterThread(reporter_config=reporter_config,
                              env_exe=self.context.env_exe,
                              runner_script=self.runner_script,
                              cwd=self.reporter_dir,
                              once=once, timeout=timeout)


def create_update_virtualenv(reporter_dir: Path,
                             config: LegionConfig.Reporter,
                             req_overrides: Optional[Dict[str, str]] = None) -> ReporterVenv:
    log.info(f"Creating/Updating Virtualenv: {reporter_dir}")
    venv = ReporterVenv(reporter_dir,
                        config.file_url,
                        config.requirements,
                        req_overrides)
    venv.create_runner(config)
    return venv


def try_to_load_configuration() -> LegionConfig:
    for config_file in CONFIGURATION_PATHS:
        if config_file.exists():
            log.info(f'Loading configuration from {str(config_file)}...')
            config = LegionConfig(config_file)
            log.info('Configuration loaded.')
            return config
    log.critical('Unable to find a configuration file in:')
    for config_file in CONFIGURATION_PATHS:
        log.critical(f'\t{str(config_file)}')
    sys.exit(1)


def parser_of(config: LegionConfig) -> ArgumentParser:
    parser = ArgumentParser(description="Legion Alerting-Monitoring System: Client "
                                        "monitoring service")
    parser.add_argument('--config',
                        action='store_true',
                        help='Display the entire configuration')
    subparsers = parser.add_subparsers(dest='reporter',
                                       help='Individual sub-commands are available for '
                                            'running/viewing configured reporters')
    for reporter in config.reporters:
        reporter_parser = subparsers.add_parser(reporter.name,
                                                description=reporter.description,
                                                help=f'{reporter.name}: {reporter.description}')
        reporter_parser.add_argument('--config',
                                     action='store_true',
                                     help=f'Display the configuration for {reporter.name}')
    return parser


def print_config(config: LegionConfig, reporter: Optional[str] = None) -> None:
    if reporter is None:
        for line in pformat(dict(config.dict)).split('\n'):
            log.info(line)
    else:
        reporter_config = config.reporters_dict[reporter]
        log.info(f'Reporter: {reporter_config.name}')
        for line in pformat(dict(reporter_config._asdict())).split('\n'):
            log.info(line)


def run_reporters(config: LegionConfig) -> None:
    reporter_venvs = [(reporter,
                       create_update_virtualenv(config.reporters_dir / reporter.name,
                                                reporter,
                                                config.package_overrides)) for reporter in config.reporters]
    reporters: List[ReporterThread] = []
    for reporter, venv in reporter_venvs:
        reporters.append(venv.reporter(reporter, once=False))
    for rep in reporters:
        rep.start()
    try:
        for rep in reporters:
            rep.join()
    except KeyboardInterrupt:
        log.info("Waiting for reporters to stop...")
        try:
            for rep in reporters:
                rep.stop()
                rep.join()
        except KeyboardInterrupt as exc:
            log.warning("Another interrupt received, forcing shutdown")
            raise exc


def cli() -> None:
    config = try_to_load_configuration()
    parser = parser_of(config)
    args = parser.parse_args()
    if args.config:
        print_config(config, args.reporter)
    elif args.reporter:
        log.info("Creating/updating virtualenvs...")
        venv = create_update_virtualenv(config.reporters_dir / args.reporter,
                                        config.reporter(args.reporter),
                                        config.package_overrides)
        log.info("Virtualenv created/updated.")
        log.info(f"Executing reporter: {args.reporter}...")
        venv.reporter(config.reporter(args.reporter), once=True).run()
        log.info("Reporter execution complete.")
    else:
        run_reporters(config)


if __name__ == '__main__':
    cli()
