"""Base class for ARC monitoring and utilities for state tracking."""

import asyncio
from functools import cached_property
from glob import glob
import os
import shutil
import time
from typing import Dict, List, Optional, Set, Tuple
import jinja2 # type: ignore
from arcnagios import arcclients, nagutils, persistence, vomsutils
from arcnagios.arcutils import jobstate_of_str, ArcClient, JobState, ParseError
from arcnagios.nagutils import NagiosPerflogTimer
from arcnagios.persistence import ObsoletePersistentObject
from arcnagios.reputation import ReputationTracker
from arcnagios.rescheduler import Rescheduler
from arcnagios.utils import host_of_uri
from arcnagios.ce.jobplugin import JobPlugin, StagingSpec
from arcnagios.ce.jobplugins import load_jobplugin

class JobDescription:
    _required_attributes = ["job_name", "application_name", "script_path"]
    def __init__(
            self,
            *,
            job_name: str,
            application_name: str,
            logdir: Optional[str] = None,
            script_path: str,
            script_args: Optional[List[str]] = None,
            output: str = 'stdout.txt',
            error: str = 'stderr.txt',
            wall_time_limit: Optional[int] = None,
            memory_limit: Optional[int] = None,
            staged_inputs: Optional[List[StagingSpec]] = None,
            staged_outputs: Optional[List[StagingSpec]] = None,
            runtime_environments: Optional[List[str]] = None,
            queue_name: Optional[str] = None,
            template: str = 'default.xrsl.j2'):
        self.template = template
        self.job_name = job_name or 'ARC Probe'
        self.application_name = application_name
        self.logdir = logdir
        self.script_path = script_path
        self.script_name = os.path.basename(script_path)
        self.script_args = script_args or []
        self.output = output
        self.error = error
        self.wall_time_limit = wall_time_limit
        self.memory_limit = memory_limit
        self.staged_inputs = staged_inputs or []
        self.staged_outputs = staged_outputs or []
        self.runtime_environments = runtime_environments or []
        self.queue_name = queue_name

    def verify(self):
        for attr in self._required_attributes:
            if getattr(self, attr) is None:
                raise nagutils.ServiceUnknown('Missing %s for job description.'
                                              % attr)

pt_jobstate = persistence.PersistentType(jobstate_of_str)
pt_jobstate_opt = persistence.PersistentType(jobstate_of_str, str, False)

class JobInfo(persistence.PersistentObject):

    persistence_version = 1

    # TODO: Check if we can use introspection of the below declarations instead.
    persistent_attributes = {
        'submission_time':      persistence.pt_float,
        'host':                 persistence.pt_str,
        'job_tag':              persistence.pt_str_opt,
        'progress_service':     persistence.pt_str_opt,
        'termination_service':  persistence.pt_str_opt,
        'job_id':               persistence.pt_str,
        'job_state':            pt_jobstate,
        'job_specific_state':   pt_jobstate_opt,
        'job_state_time':       persistence.pt_float_opt,
        'job_state_alert':      persistence.pt_int_opt,
        'check_time':           persistence.pt_float_opt,
        'check_attempts':       persistence.pt_int_opt,
        'fetch_attempts':       persistence.pt_int_opt,
        'stored_urls':          persistence.pt_str_list,
        'tests':                persistence.pt_str_list,
        'reputation_choices':   persistence.pt_json_opt,
    }

    submission_time: float
    host: str
    job_tag: Optional[str]
    progress_service: Optional[str]
    termination_service: Optional[str]
    job_id: str
    job_state: JobState
    job_specific_state: Optional[str]
    job_state_time: Optional[float]
    job_state_alert: Optional[int]
    check_time: Optional[float]
    check_attempts: Optional[int]
    fetch_attempts: Optional[int]
    stored_urls: List[str]
    tests: List[str]
    reputation_choices: Optional[Dict[str, str]]

    @property
    def host_and_tag(self):
        if self.job_tag:
            return '%s#%s' % (self.host, self.job_tag)
        return self.host

    def __eq__(self, other):
        return self.host_and_tag == other.host_and_tag
    def __lt__(self, other):
        return self.host_and_tag < other.host_and_tag

def key_value(s: str) -> Tuple[str, str]:
    kv = s.split('=', 1)
    if len(kv) != 2:
        raise ValueError('Expecting an argument of the form KEY=VALUE.')
    return (kv[0], kv[1])

class JobNagiosPlugin(nagutils.NagiosPlugin, vomsutils.NagiosPluginVomsMixin):
    """Nagios probe to test ARC CEs.  The probe has two sub-commands
    implemented by `check_submit` and `check_monitor`.  The former is run on
    all CEs, while the latter is run to collect submitted jobs."""

    # pylint: disable=abstract-method,super-init-not-called

    probe_name = 'ARCCE'
    main_config_section = ['arcce']

    JOB_DESCRIPTION_FILENAME = 'job.xrsl'
    JOB_SCRIPT_FILENAME = 'job.sh'
    JOBID_FILENAME = 'active.jobid'
    ACTIVE_JOB_FILENAME = 'active.map'
    JOB_OUTPUT_DIRNAME = 'job_output'

    _archive_filenames = [
        JOBID_FILENAME, JOB_DESCRIPTION_FILENAME, JOB_SCRIPT_FILENAME, ACTIVE_JOB_FILENAME
    ]

    prev_status = None

    # Timeout in seconds for cleaner tasks.
    cleaner_arcrm_timeout = 5
    cleaner_arcclean_timeout = 5
    cleaner_arckill_timeout = 5

    arcclient: ArcClient
    _reputation_tracker: ReputationTracker

    def _disconnect_reputation_tracker(self):
        if self._reputation_tracker is not None:
            self._reputation_tracker.disconnect()

    def __init__(
            self, use_host: bool = False, use_port: bool = False,
            default_port: Optional[int] = None):
        default_template_dirs = glob(os.path.join(self.config_dir(),'*.d'))
        default_template_dirs.sort()
        nagutils.NagiosPlugin.__init__(
                self, use_host=use_host, use_port=use_port,
                default_port=default_port)

        argp = self.argparser
        argp.add_argument('--fqan', dest = 'fqan')
        argp.add_argument('-O', dest = 'jobtest_options',
                action = 'append', type = key_value, default = [],
                help = 'Given a value of the form VAR=VALUE, binds VAR to '
                       'VALUE in the environment of the job tests.')
        argp.add_argument('--template-dir', dest = 'template_dirs',
                action = 'append', default = default_template_dirs,
                help = 'Add a directory from which job description templates '
                       'can be loaded.')
        argp.add_argument('--granular-perfdata', dest = 'granular_perfdata',
                default = False, action = 'store_true',
                help = 'Report ARC command timing performance data per host '
                       'using labels of the form ARCCMD[HOST]. By default '
                       'report the aggretate time across hosts.')
        argp.add_argument('--concurrency', type = int, default = 20,
                help = 'The maximum number of ARC client processes to run '
                        'concurrently.')
        self.at_exit(self._disconnect_reputation_tracker)

    def parse_args(self, args: List[str]) -> None:
        nagutils.NagiosPlugin.parse_args(self, args)
        self.opts.template_dirs.reverse()
        self.arcclient = ArcClient(self.perflog)
        self._reputation_tracker = ReputationTracker(
                self.config,
                os.path.join(self.opts.arcnagios_spooldir, 'reputation.db'))

    @cached_property
    def _arcclients_semaphore(self):
        return asyncio.Semaphore(self.opts.concurrency or 20)

    def top_vopath(self, suffix: str) -> str:
        return os.path.join(self.opts.arcnagios_spooldir,
                            self.voms_suffixed('ce') + '-' + suffix)

    @cached_property
    def top_workdir(self) -> str:
        return self.top_vopath('state')

    def workdir_for(self, host: str, job_tag: Optional[str]):
        if job_tag:
            return os.path.join(self.top_workdir, host + '#' + job_tag)
        else:
            return os.path.join(self.top_workdir, host)

    def archivedir_for(self, host: str, job_tag: Optional[str]):
        return os.path.join(self.top_vopath('archive'),
                            time.strftime('%Y-%m/%Y-%m-%d/%H:%M:%S-')
                            + host + (job_tag and ('#' + job_tag) or ''))

    @cached_property
    def template_environment(self) -> jinja2.Environment:
        return jinja2.Environment(
            loader=jinja2.FileSystemLoader(self.opts.template_dirs),
            autoescape=False)

    def write_job_description(self, out_path: str, jobdesc: JobDescription) \
            -> None:
        jobdesc.verify()
        try:
            templ = self.template_environment.get_template(jobdesc.template)
        except jinja2.TemplateNotFound as exn:
            raise nagutils.ServiceUnknown('%s. The searched included %s.'
                    % (exn, ', '.join(self.opts.template_dirs)))
        content = templ.render(jd=jobdesc)
        with open(out_path, 'w', encoding='utf-8') as fh:
            fh.write(content)

    async def _cleaner_arcrm(self, url: str, n_attempts: int) -> bool:
        se_host = host_of_uri(url)
        with NagiosPerflogTimer(self.perflog, "arcrm_time", se_host):
            result = await arcclients.arcrm(
                    url, force = n_attempts > 8,
                    timeout = self.cleaner_arcrm_timeout,
                    log = self.log)
        if result.is_ok():
            self.log.info('Removed test file %s.', url)
            return True
        else:
            self.log.warning('Failed to remove %s.', url)
            return False

    async def _cleaner_arcclean(self, job_id: str, n_attempts: int) -> bool:
        ce_host = host_of_uri(job_id)
        with NagiosPerflogTimer(self.perflog, "arcrm_time", ce_host):
            result = await arcclients.arcclean(
                    job_id, force = n_attempts > 8,
                    timeout = self.cleaner_arcclean_timeout,
                    log = self.log)
        if result.is_ok():
            self.log.info('Removed job %s', job_id)
            return True
        else:
            self.log.warning('Failed to clean %s', job_id)
            return False

    async def _cleaner_arckill(self, job_id: str, n_attempts: int) -> bool:
        ce_host = host_of_uri(job_id)
        with NagiosPerflogTimer(self.perflog, "arckill_time", ce_host):
            result = await arcclients.arckill(
                    job_id, timeout = self.cleaner_arckill_timeout,
                    log = self.log)
        if result.is_ok():
            return True

        result = await arcclients.arcclean(
                job_id, force =  n_attempts > 8,
                timeout = self.cleaner_arcclean_timeout,
                log = self.log)
        if result.is_ok():
            self.log.info('Killed %s', job_id)
            return True
        else:
            self.log.warning('Failed to kill %s', job_id)
            return False

    @cached_property
    def cleaner(self) -> Rescheduler:
        cleaner = Rescheduler(self.top_vopath('state.db'), 'cleaner',
                              nagios_report = self.nagios_report)
        cleaner.register('arcrm', self._cleaner_arcrm)
        cleaner.register('arcclean', self._cleaner_arcclean)
        cleaner.register('arckill', self._cleaner_arckill)
        self.at_exit(cleaner.close)
        return cleaner

    def cleanup_job_files(self, host: str, job_tag: Optional[str],
                          archive: bool = False) -> None:
        self.log.debug('Cleaning up job files for %s.', host)
        workdir = self.workdir_for(host, job_tag)
#       archdir = os.path.join(workdir,
#                              time.strftime('archive/%Y-%m-%d/%H:%M:%S'))
        archdir = self.archivedir_for(host, job_tag)
        archdir_created = False
        for filename in self._archive_filenames:
            try:
                if archive:
                    if os.path.exists(os.path.join(workdir, filename)):
                        if not archdir_created:
                            os.makedirs(archdir)
                            archdir_created = True
                        os.rename(os.path.join(workdir, filename),
                                  os.path.join(archdir, filename))
                else:
                    os.unlink(os.path.join(workdir, filename))
            except OSError:
                pass
        try:
            job_output_dir = os.path.join(workdir, self.JOB_OUTPUT_DIRNAME)
            if os.path.exists(job_output_dir) and os.listdir(job_output_dir):
                if archive:
                    if not archdir_created:
                        os.makedirs(archdir)
                        archdir_created = True
                    os.rename(job_output_dir,
                              os.path.join(archdir, self.JOB_OUTPUT_DIRNAME))
                else:
                    last_dir = job_output_dir + '.LAST'
                    shutil.rmtree(last_dir, ignore_errors = True)
                    os.rename(job_output_dir, last_dir)
        except OSError as exn:
            self.log.warning('Error clearing %s: %s', job_output_dir, exn)

    def load_active_job(self, host: str, job_tag: Optional[str]) \
            -> Optional[JobInfo]:
        """Load information about the current job on `host : str` tagged with
        `job_tag : str`, or `None` if no information is found."""

        workdir = self.workdir_for(host, job_tag)
        ajf = os.path.join(workdir, self.ACTIVE_JOB_FILENAME)
        if os.path.exists(ajf):
            self.log.debug('Loading job info from %s.', ajf)
            jobinfo = JobInfo()
            try:
                jobinfo.persistent_load(
                    ajf, log=self.log,
                    persistence_version=jobinfo.persistence_version)
                return jobinfo
            except ObsoletePersistentObject:
                self.log.warning('Ignoring outdated job file %s.')
                return None
            except (OSError, ParseError) as exn:
                self.log.error('Cannot load job file %s: %s', ajf, exn)
                return None
        return None

    def save_active_job(self, jobinfo: JobInfo,
                        host: str, job_tag: Optional[str]) -> None:
        """Save information about the current job running on `host : str`
        tagged with `job_tag : str`."""

        workdir = self.workdir_for(host, job_tag)
        ajf = os.path.join(workdir, self.ACTIVE_JOB_FILENAME)
        self.log.debug('Saving active job info.')
        jobinfo.persistent_save(
            ajf, log=self.log,
            persistence_version=jobinfo.persistence_version)

    def collect_active_jobids(self) -> Set[str]:
        jobids = set()
        for path in glob(os.path.join(self.top_workdir, '*', 'active.jobid')):
            try:
                with open(path, encoding='utf-8') as fh:
                    jobids.add(fh.readline().strip())
            except IOError as exn:
                self.log.error('Failed to read %s: %s', path, exn)
        return jobids

    async def _discard_stored_urls(self, jobinfo: JobInfo) -> None:
        for url in jobinfo.stored_urls:
            if not url.startswith('file:'):
                await self.cleaner.call('arcrm', url)

    def cleanup_job_tests(self, jobinfo: JobInfo) -> None:
        for test_name in jobinfo.tests:
            try:
                test = self.load_jobtest(test_name, hostname = jobinfo.host)
                test.cleanup(jobinfo.job_state)
            except Exception as exn: # pylint: disable=broad-except
                self.log.error('Error in cleanup %s for %s: %s',
                               test_name, jobinfo.job_id, exn)

    async def cleanup_job(
            self, jobinfo: JobInfo, archive: bool = False
        ) -> None:
        """Clean up job state from a fetched job."""

        await self._discard_stored_urls(jobinfo)
        self.cleanup_job_tests(jobinfo)
        self.cleanup_job_files(jobinfo.host, jobinfo.job_tag, archive = archive)

    async def discard_job(
            self, jobinfo: JobInfo, archive: bool = False
        ) -> None:
        """Discard the job described by `jobinfo : JobInfo`."""

        if jobinfo.job_state.is_final():
            await self.cleaner.call('arcclean', jobinfo.job_id)
        else:
            await self.cleaner.call('arckill', jobinfo.job_id)
        await self.cleanup_job(jobinfo, archive = archive)

    def load_jobtest(self, jobtest_name: str, **env) -> JobPlugin:
        """Load a plugin-based job-test from the section of the configuration
        specified by `jobtest_name`.  The result is an instance of `JobPlugin`
        subclass specified by the ``jobplugin`` variable of the given
        section."""

        env['config_dir'] = self.config_dir()
        if self.voms:
            env['voms'] = self.voms
        env.update(self.opts.jobtest_options)

        jobplugin_section = 'arcce.%s' % jobtest_name
        if not self.config.has_section(jobplugin_section):
            raise nagutils.ServiceUnknown(
                    'Missing configuration section %s for '
                    'job-plugin test.' % jobplugin_section)
        jobplugin_name = self.config.get(jobplugin_section, 'jobplugin')
        jobplugin_cls = load_jobplugin(jobplugin_name)
        return jobplugin_cls(jobplugin_name, self.config, jobplugin_section,
                             self._reputation_tracker,
                             self.log, self.arcclient, env)
