import re

import boto3
import logging
import requests
from datetime import datetime, timezone, timedelta

from botocore import UNSIGNED
from botocore.client import Config

from time import sleep
from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter

from indra import get_config, has_config
from indra.util.nested_dict import NestedDict

logger = logging.getLogger(__name__)


def get_ids(job_list):
    if job_list is None:
        return None
    return [job['jobId'] for job in job_list]


def kill_all(job_queue, reason='None given', states=None, kill_list=None):
    """Terminates/cancels all jobs on the specified queue.

    Parameters
    ----------
    job_queue : str
        The name of the Batch job queue on which you wish to terminate/cancel
        jobs.
    reason : str
        Provide a reason for the kill that will be recorded with the job's
        record on AWS.
    states : None or list[str]
        A list of job states to remove. Possible states are 'STARTING',
        'RUNNABLE', and 'RUNNING'. If None, all jobs in all states will be
        ended (modulo the `kill_list` below).
    kill_list : None or list[dict]
        A list of job dictionaries (as returned by the submit function) that
        you specifically wish to kill. All other jobs on the queue will be
        ignored. If None, all jobs on the queue will be ended (modulo the
        above).

    Returns
    -------
    killed_ids : list[str]
        A list of the job ids for jobs that were killed.
    """
    # Default is all states.
    if states is None:
        states = ['STARTING', 'RUNNABLE', 'RUNNING']

    # Get batch client
    batch = boto3.client('batch')

    # Get all other jobs, and terminate them.
    killed_ids = []
    for status in states:
        running = batch.list_jobs(jobQueue=job_queue, jobStatus=status)
        active_job_list = running.get('jobSummaryList')
        if active_job_list is None:
            continue

        for job in active_job_list:
            # Check if this is one of the specified jobs, if any specified.
            ids_to_kill = get_ids(kill_list)
            if ids_to_kill is not None and job['jobId'] not in ids_to_kill:
                continue

            # End the job.
            if status == 'RUNNING':
                logger.info('Terminating {jobName} ({jobId})'.format(**job))
                res = batch.terminate_job(jobId=job['jobId'], reason=reason)
            else:
                logger.info('Canceling {jobName} ({jobId})'.format(**job))
                res = batch.cancel_job(jobId=job['jobId'], reason=reason)

            # Record the result of the kill
            killed_ids.append(res)

    return killed_ids


def tag_instance(instance_id, **tags):
    """Tag a single ec2 instance."""
    logger.debug("Got request to add tags %s to instance %s."
                 % (str(tags), instance_id))
    ec2 = boto3.resource('ec2')
    instance = ec2.Instance(instance_id)

    # Remove None's from `tags` and reformat to the list format that
    # boto3 expects
    tag_list = [{'Key': k, 'Value': v} for k, v in tags.items() if k and v]
    if tag_list:
        logger.info('Adding project tags "%s" to instance %s'
                    % (str(tag_list), instance_id))
        instance.create_tags(Tags=tag_list)
        vols = instance.volumes.all()
        for page in vols.pages():
            for vol in page:
                vol.create_tags(Tags=tag_list)
    else:
        logger.info('No new tags from: %s' % str(tags))
    instance_tags = {tag.get('Key'): tag.get('Value') for tag in instance.tags}
    logger.info('Updated instance tags: %s' % instance_tags)
    return


def tag_myself(project='aske', **other_tags):
    """Function run when indra is used in an EC2 instance to apply tags."""
    base_url = "http://169.254.169.254"
    try:
        resp = requests.get(base_url + "/latest/meta-data/instance-id")
    except requests.exceptions.ConnectionError:
        logger.warning("Could not connect to service. Note this should only "
                       "be run from within a batch job.")
        return
    instance_id = resp.text
    tag_instance(instance_id, project=project, **other_tags)
    return


def get_batch_command(command_list, project=None, purpose=None):
    """Get the command appropriate for running something on batch."""
    command_str = ' '.join(command_list)
    ret = ['python3', '-m', 'indra.util.aws', 'run_in_batch', command_str]
    if not project and has_config('DEFAULT_AWS_PROJECT'):
        project = get_config('DEFAULT_AWS_PROJECT')
    if project:
        ret += ['--project', project]
    if purpose:
        ret += ['--purpose', purpose]
    return ret


def run_in_batch(command_list, project, purpose):
    from subprocess import call
    tag_myself(project, purpose=purpose)
    logger.info("Running command list: %s" % str(command_list))
    logger.info('\n'+20*'='+' Begin Primary Command Output '+20*'='+'\n')
    ret_code = call(command_list)
    logger.info('\n'+21*'='+' End Primary Command Output '+21*'='+'\n')
    return ret_code


def get_jobs(job_queue='run_reach_queue', job_status='RUNNING'):
    """Returns a list of dicts with jobName and jobId for each job with the
    given status."""
    batch = boto3.client('batch')
    jobs = batch.list_jobs(jobQueue=job_queue, jobStatus=job_status)
    return jobs.get('jobSummaryList')


s3_path_patt = re.compile('^s3:([-a-zA-Z0-9_]+)/(.*?)$')


class JobLog(object):
    """Gets the Cloudwatch log associated with the given job.

    Parameters
    ----------
    job_info : dict
        dict containing entries for 'jobName' and 'jobId', e.g., as returned
        by get_jobs()
    log_group_name : string
        Name of the log group; defaults to '/aws/batch/job'

    Returns
    -------
    list of strings
        The event messages in the log, with the earliest events listed first.
    """
    _suffix_base = '/part_'

    def __init__(self, job_info, log_group_name='/aws/batch/job',
                 verbose=False, append_dumps=True):
        self.job_name = job_info['jobName']
        self.job_id = job_info['jobId']
        self.logs_client = boto3.client('logs')
        self.verbose = verbose
        self.log_group_name = log_group_name
        batch = boto3.client('batch')
        resp = batch.describe_jobs(jobs=[self.job_id])
        job_desc = resp['jobs'][0]
        job_def_name = job_desc['jobDefinition'].split('/')[-1].split(':')[0]
        task_arn_id = job_desc['container']['taskArn'].split('/')[-1]
        self.log_stream_name = '%s/default/%s' % (job_def_name, task_arn_id)
        self.latest_timestamp = None
        self.lines = []
        self.nextToken = None
        self.__len = 0
        self.append = append_dumps
        return

    def __len__(self):
        return self.__len

    def clear_lines(self):
        self.lines = []

    def dump(self, out_file, append=None):
        """Dump the logs in their entirety to the specified file."""
        if append is None:
            append = self.append
        elif append != self.append:
            logger.info("Overriding default append behavior. This could muddy "
                        "future loads.")
        m = s3_path_patt.match(out_file)
        if m is not None:
            # If the user wants the files on s3...
            bucket, prefix = m.groups()
            s3 = boto3.client('s3')

            # Find the largest part number among the current suffixes
            if append:
                max_num = 0
                for key in iter_s3_keys(s3, bucket, prefix, do_retry=False):
                    if key[len(prefix):].startswith(self._suffix_base):
                        num = int(key[len(prefix + self._suffix_base):])
                        if max_num > num:
                            max_num = num

                # Create the new suffix, and dump the lines to s3.
                new_suffix = self._suffix_base + str(max_num + 1)
                key = prefix + new_suffix
            else:
                key = prefix
            s3.put_object(Bucket=bucket, Key=key, Body=self.dumps())
        else:
            # Otherwise, if they want them locally...
            with open(out_file, 'wt' if append else 'w') as f:
                for line in self.lines:
                    f.write(line)
        return

    def load(self, out_file):
        """Load the log lines from the cached files."""
        m = s3_path_patt.match(out_file)
        if m is not None:
            bucket, prefix = m.groups()
            s3 = boto3.client('s3')

            if self.append:
                prior_line_bytes = []
                for key in sorted(iter_s3_keys(s3, bucket, prefix)):
                    if key[len(prefix):].startswith(self._suffix_base):
                        res = s3.get_object(Bucket=bucket, Key=key)
                        prior_line_bytes += res['Body'].read().splitlines()
            else:
                res = s3.get_object(Bucket=bucket, Key=prefix)
                prior_line_bytes = res['Body'].read().splitlines()

            prior_lines = [s.decode('utf-8') + '\n'
                           for s in prior_line_bytes]
        else:
            with open(out_file, 'r') as f:
                prior_lines = f.readlines()
        self.lines = prior_lines + self.lines
        return

    def dumps(self):
        return ''.join(self.lines)

    def get_lines(self):
        kwargs = {'logGroupName': self.log_group_name,
                  'logStreamName': self.log_stream_name,
                  'startFromHead': True}
        while True:
            if self.nextToken is not None:
                kwargs['nextToken'] = self.nextToken
            response = self.logs_client.get_log_events(**kwargs)
            # If we've gotten all the events already, the nextForwardToken for
            # this call will be the same as the last one
            if response.get('nextForwardToken') == self.nextToken:
                break
            else:
                events = response.get('events')
                if events:
                    for evt in events:
                        line = '%s: %s\n' % (evt['timestamp'], evt['message'])
                        self.lines.append(line)
                        self.latest_timestamp = \
                            (datetime.fromtimestamp(evt['timestamp']/1000)
                                     .astimezone(timezone.utc)
                                     .replace(tzinfo=None))
                        self.__len += 1
                        if self.verbose:
                            logger.info('%d %s' % (len(self.lines), line))
                self.nextToken = response.get('nextForwardToken')
        return


def dump_logs(job_queue='run_reach_queue', job_status='RUNNING'):
    """Write logs for all jobs with given the status to files."""
    jobs = get_jobs(job_queue, job_status)
    for job in jobs:
        log = JobLog(job)
        log.get_lines()
        log.dump('{jobName}_{jobId}.log'.format(**job))


def get_date_from_str(date_str):
    """Get a utc datetime object from a string of format %Y-%m-%d-%H-%M-%S

    Parameters
    ----------
    date_str : str
        A string of the format %Y(-%m-%d-%H-%M-%S). The string is assumed
        to represent a UTC time.

    Returns
    -------
    datetime.datetime
    """
    date_format = '%Y-%m-%d-%H-%M-%S'