#!/usr/bin/python
import argparse
import bz2
import datetime
import io
import logging
import os
import os.path
import re
import socket
import sys
import itertools
import errno
import time

try:
    import boto
except ImportError:
    boto = None

from blueox import store

DEFAULT_LOG_PATH = "/var/log/blueox"
DEFAULT_RETAIN_DAYS = 7
ENV_VAR_BUCKET_NAME = "BLUEOX_BUCKET"
ENV_VAR_LOG_PATH = "BLUEOX_LOG_PATH"

ZIP_COMMAND = "zip"
PRUNE_COMMAND = "prune"
ARCHIVE_COMMAND = "archive"
UPLOAD_COMMAND = "upload"
DOWNLOAD_COMMAND = "download"
CAT_COMMAND = "cat"

log = logging.getLogger(__name__)


def setup_logging(verbose_count):
    level = logging.WARNING
    if verbose_count > 0:
        level = logging.INFO
    if verbose_count > 1:
        level = logging.DEBUG

    logging.basicConfig(level=level, format="%(asctime)-15s %(message)s")


def lock_file(file_name):
    try:
        return os.open(file_name,
                       os.O_WRONLY | os.O_CREAT | os.O_EXLOCK | os.O_NDELAY)
    except OSError as e:
        if e.errno == errno.EAGAIN:
            return None
        elif e.errno == errno.ENOENT:
            log.error("Failed to open %r: %r", file_name, e)
            return None
        else:
            raise


def unlock_file(fd, file_name):
    os.close(fd)
    os.unlink(file_name)


def do_upload(log_path, bucket, zipped_only):
    """Upload available local log files to S3"""
    log_files = store.list_log_files(log_path)

    for lf in log_files:
        log.debug("Examining %s for archive", lf.file_path)
        if zipped_only and not lf.bzip:
            continue

        remote_lf = lf.build_remote(socket.gethostname())

        key = remote_lf.s3_key(bucket)

        log.debug("Checking for key %s", key)
        # TODO: Be nice to check for size differences
        # Might also be a good optimization to somehow download existing keys
        # in bulk, rather than make a request for each one.
        if not key.exists():
            with open(lf.get_local_file_path(log_path), "r") as fp:
                log.info("Uploading %s", key)
                key.set_contents_from_file(fp)
                log.debug("Done Uploading %s", key)
        else:
            log.debug("key %s already exists", key)


def do_download(bucket, type_name, start_dt, end_dt):
    "Run the download action"
    for log_file in store.find_log_files_in_s3(bucket, type_name, start_dt,
                                               end_dt):
        if os.path.exists(log_file.file_name):
            log.info("Skipping %s, exists", log_file.file_name)
        else:
            log.info("Downloading %s", log_file.file_name)
            log_file.s3_key(bucket).get_contents_to_filename(log_file.file_name)


def do_cat_from_s3(bucket, type_name, start_dt, end_dt):
    "Run the cat action against remote logs (s3)"
    log_files = store.find_log_files_in_s3(bucket, type_name, start_dt, end_dt)

    log_files.sort(key=lambda f: f.sort_dt)

    stream = itertools.chain(*(lf.open(bucket) for lf in log_files))
    for data in stream:
        sys.stdout.write(data)


def do_cat_from_local(log_path, type_name, start_dt, end_dt):
    "Run the cat action against local log path"
    log_files = store.find_log_files_in_path(log_path, type_name, start_dt,
                                             end_dt)

    log_files.sort(key=lambda f: f.sort_dt)

    stream = itertools.chain(*(lf.open(log_path) for lf in log_files))
    for data in stream:
        sys.stdout.write(data)


def do_zip(log_path):
    log.debug("Listing log files for %r", log_path)
    log_files = store.list_log_files(log_path)

    for log_file in store.filter_log_files_for_zipping(log_files):
        log.info("Zipping %s", log_file.file_path)

        store.zip_log_file(log_file, log_path)


def do_prune(log_path, retain_days):
    min_archive_date = (
        datetime.datetime.utcnow() - datetime.timedelta(days=retain_days)
    ).date()

    for dirpath, dirnames, filenames in os.walk(log_path):
        try:
            dir_date = datetime.datetime.strptime(os.path.basename(dirpath),
                                                  '%Y%m%d').date()
        except ValueError:
            continue

        if dir_date < min_archive_date:
            for filename in filenames:
                full_path = os.path.join(dirpath, filename)
                log.info("Removing %s", full_path)
                os.unlink(full_path)

            os.rmdir(dirpath)


def parse_date_range_arguments(parser, args):
    if args.start is None:
        start_date = datetime.datetime.utcnow().replace(
            hour=0,
            minute=0,
            second=0,
            microsecond=0)
    else:
        try:
            start_date = store.parse_date_range_argument(args.start)
        except store.InvalidDateError:
            parser.error("Bad format for start date. Try YYYYMMDD")

    if args.end is None:
        end_date = start_date + datetime.timedelta(hours=23,
                                                   minutes=59,
                                                   seconds=59)
    else:
        try:
            end_date = store.parse_date_range_argument(args.end)
        except store.InvalidDateError:
            parser.error("Bad format for end date. Try YYYYMMDD")

    return start_date, end_date


def main():
    default_log_path = os.environ.get(ENV_VAR_LOG_PATH, DEFAULT_LOG_PATH)
    default_bucket = os.environ.get(ENV_VAR_BUCKET_NAME, None)

    parser = argparse.ArgumentParser(description='Manage blueox logs')
    parser.add_argument('--verbose', '-v', action='count')
    parser.add_argument(
        '--lock-file',
        action='store',
        help=
        "Create a lock file to prevent another process from running concurrently")

    subparsers = parser.add_subparsers(
        dest='command',
        help='available commands')

    parser_zip = subparsers.add_parser(ZIP_COMMAND, help='Zip log files')
    parser_zip.add_argument(
        '--log-path', '-p',
        action='store',
        default=default_log_path)

    parser_prune = subparsers.add_parser(PRUNE_COMMAND, help='Prune log files')
    parser_prune.add_argument(
        '--log-path', '-p',
        action='store',
        default=default_log_path)
    parser_prune.add_argument(
        '--retain-days', '-d',
        action='store',
        type=int,
        default=DEFAULT_RETAIN_DAYS)

    parser_upload = subparsers.add_parser(
        UPLOAD_COMMAND,
        help='Upload local logs to S3')
    parser_upload.add_argument('--bucket', '-b',
                               action='store',
                               default=default_bucket)
    parser_upload.add_argument(
        '--log-path', '-p',
        action='store',
        default=default_log_path)
    parser_upload.add_argument(
        '--zipped-only',
        action='store_true',
        default=False)

    parser_download = subparsers.add_parser(
        DOWNLOAD_COMMAND,
        help='Download logs from S3')
    parser_download.add_argument('--bucket', '-b',
                                 action='store',
                                 default=default_bucket)
    parser_download.add_argument(
        '--start', '-s',
        action='store',
        help="start date (YYYYMMDD [HH:MM])")
    parser_download.add_argument(
        '--end', '-e',
        action='store',
        required=False,
        help="end date")
    parser_download.add_argument(
        'log_type',
        action='store',
        nargs=1,
        help="log type to download")

    parser_cat = subparsers.add_parser(
        CAT_COMMAND,
        help='output specified log to stdout')

    parser_cat.add_argument(
        'log_type',
        action='store',
        nargs=1,
        help="log type")
    parser_cat.add_argument('--local',
                            action='store_true',
                            help="Check local file system for log files")
    parser_cat.add_argument('--bucket', '-b',
                            action='store',
                            default=default_bucket)
    parser_cat.add_argument(
        '--log-path', '-p',
        action='store',
        default=default_log_path)
    parser_cat.add_argument(
        '--start', '-s',
        action='store',
        help="start date (YYYYMMDD [HH:MM])")
    parser_cat.add_argument('--end', '-e', action='store', help="end date")

    parser_archive = subparsers.add_parser(
        ARCHIVE_COMMAND,
        help='Zip, Upload and Prune')
    parser_archive.add_argument('--bucket', '-b',
                                action='store',
                                default=default_bucket)
    parser_archive.add_argument(
        '--log-path', '-p',
        action='store',
        default=default_log_path)
    parser_archive.add_argument(
        '--retain-days', '-d',
        action='store',
        type=int,
        default=DEFAULT_RETAIN_DAYS)

    args = parser.parse_args()

    setup_logging(args.verbose)

    lock_fd = None
    if args.lock_file:
        lock_fd = lock_file(args.lock_file)
        if not lock_fd:
            parser.error("Failed to lock on {}".format(args.lock_file))

    bucket = None
    if hasattr(args, 'bucket'):
        if boto is None:
            parser.error("boto library not available")

        bucket_name = args.bucket or os.environ.get(ENV_VAR_BUCKET_NAME, None)
        if not bucket_name:
            parser.error("Bucket name not provided")

        try:
            bucket = store.open_bucket(bucket_name)
        except boto.exception.NoAuthHandlerFound:
            parser.error(
                "AWS Credentials not found, check your AWS environment variables")

        if bucket is None:
            parser.error("Bucket not found")

    if args.command == ZIP_COMMAND:
        do_zip(args.log_path)
    elif args.command == PRUNE_COMMAND:
        do_prune(args.log_path, args.retain_days)
    elif args.command == UPLOAD_COMMAND:
        do_upload(args.log_path, bucket, args.zipped_only)
    elif args.command == DOWNLOAD_COMMAND:
        start_dt, end_dt = parse_date_range_arguments(parser, args)

        do_download(bucket, args.log_type[0], start_dt, end_dt)
    elif args.command == CAT_COMMAND:
        start_dt, end_dt = parse_date_range_arguments(parser, args)

        if args.local:
            do_cat_from_local(args.log_path, args.log_type[0], start_dt, end_dt)
        elif bucket:
            do_cat_from_s3(bucket, args.log_type[0], start_dt, end_dt)
        else:
            parser.error("Bucket or --local not specified")

    elif args.command == ARCHIVE_COMMAND:

        do_zip(args.log_path)

        do_upload(args.log_path, bucket, True)

        do_prune(args.log_path, args.retain_days)

    else:
        parser.error("Unknown command")

    if lock_fd:
        unlock_file(lock_fd, args.lock_file)

    log.info("done")


if __name__ == '__main__':
    main()
