#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Download, unzip, clean and load the latest CAL-ACCESS database ZIP.
"""
import os
import logging
from sys import exit
from time import sleep
from zipfile import ZIP_DEFLATED, ZIP_STORED, ZipFile
from django.core.files import File
from django.conf import settings
from django.contrib.humanize.templatetags.humanize import naturaltime
from clint.textui import progress
from django.utils.timezone import now
from django.core.management import call_command
from calaccess_raw.management import handle_command
from django.core.management.base import CommandError
from django.template.loader import render_to_string
from calaccess_raw.models.tracking import RawDataVersion
from calaccess_raw.management.commands import CalAccessCommand
from calaccess_raw.management.commands.extractcalaccessrawfiles import TestCommand as TestExtractCommand
from calaccess_raw import (
    get_download_directory,
    get_test_download_directory,
    get_model_list
)
logger = logging.getLogger(__name__)


class Command(CalAccessCommand):
    """
    Download, unzip, clean and load the latest CAL-ACCESS database ZIP.
    """
    help = "Download, unzip, clean and load the latest CAL-ACCESS database ZIP"

    def add_arguments(self, parser):
        """
        Adds custom arguments specific to this command.
        """
        super(Command, self).add_arguments(parser)
        parser.add_argument(
            "--keep-files",
            action="store_true",
            dest="keep_files",
            default=False,
            help="Keep zip, unzipped, TSV and CSV files"
        )
        parser.add_argument(
            "--noinput",
            action="store_true",
            dest="noinput",
            default=False,
            help="Update or resume previous update without asking permission"
        )
        parser.add_argument(
            "--test",
            "--use-test-data",
            action="store_true",
            dest="test_data",
            default=False,
            help="Use sampled test data (skips download, clean a load)"
        )
        parser.add_argument(
            "-a",
            "--app-name",
            dest="app_name",
            default="calaccess_raw",
            help="Name of Django app with models into which data will "
                 "be imported (if other not calaccess_raw)"
        )

    def handle(self, *args, **options):
        """
        Make it happen.
        """
        super(Command, self).handle(*args, **options)

        # set / compute any attributes that multiple class methods need
        self.app_name = options["app_name"]
        self.keep_files = options["keep_files"]
        self.test_mode = options['test_data']
        self.noinput = options['noinput']

        if self.test_mode:
            # and always keep files when running test data
            self.keep_files = True
            self.data_dir = get_test_download_directory()
            # need to set this app-wide because cleancalaccessrawfile
            #   also calls get_download_directory
            settings.CALACCESS_DOWNLOAD_DIR = self.data_dir
            self.noinput = True
        else:
            self.data_dir = get_download_directory()

        # set the download zip file path
        self.zip_path = os.path.join(self.data_dir, self.url.split('/')[-1])
        # make the data dir if necessary
        os.path.exists(self.data_dir) or os.makedirs(self.data_dir)
        self.tsv_dir = os.path.join(self.data_dir, "tsv/")
        # Immediately check that the tsv directory exists when using test data,
        #   so we can stop immediately.
        if self.test_mode:
            if not os.path.exists(self.tsv_dir):
                raise CommandError("Data tsv directory does not exist "
                                   "at %s" % self.tsv_dir)
            elif self.verbosity:
                self.log("Using test data")

        # set up the dir for clean files
        self.csv_dir = os.path.join(self.data_dir, "csv/")
        os.path.exists(self.csv_dir) or os.makedirs(self.csv_dir)

        # in test mode, get download metadata from local file
        if self.test_mode:
            with open(self.data_dir + "/sampled_version.txt", "r") as f:
                # read expected_size from sampled_version.txt
                try:
                    # long int type is big enough for double the current size of the zip
                    expected_size = long(f.readline())
                except NameError:
                    # in py3, no long(), instead int will suffice
                    expected_size = int(f.readline())
                # and release_datetime
                release_datetime = self.parse_imf_datetime_str(f.readline())
                # get or create the RawDataVersion
                latest_version, created = self.get_or_create_version(
                    expected_size,
                    release_datetime,
                )
        # else request it
        else:
            download_metadata = self.get_download_metadata()
            logger.debug('Server: %s' % download_metadata['server'])
            logger.debug('ETag: %s' % download_metadata['etag'])
            logger.debug('Last-Modified: %s' % download_metadata['last-modified'])
            logger.debug('Content-Length: %s' % download_metadata['content-length'])
            # get or create the RawDataVersion
            latest_version, created = self.get_or_create_version(
                download_metadata['content-length'],
                self.parse_imf_datetime_str(
                    download_metadata['last-modified']
                ),
            )
        # log if latest version is new
        if created:
            logger.info('New CAL-ACCESS version available.')

        previous_version = None
        since_previous_release = None
        can_resume_previous = False
        # if update to latest version has not yet started, check to see if previous
        #   version can resume
        if not latest_version.update_start_datetime:
            # get previous versions
            prev_downloaded_versions = RawDataVersion.objects.filter(
                download_start_datetime__isnull=False
            ).exclude(id=latest_version.id)
            # if there are any
            if prev_downloaded_versions:
                # get the last one started
                previous_version = prev_downloaded_versions.latest(
                    'download_start_datetime'
                )
                # for prompt, get time since previous release
                since_previous_release = naturaltime(
                    previous_version.release_datetime
                )
                # if the previous version finished downloading and didn't finish updating
                if (
                    previous_version.download_completed and
                    not previous_version.update_completed
                ):
                    # can resume if zip file has not been deleted
                    #   or if the raw files have been extracted
                    if (
                        os.path.exists(self.zip_path) or
                        previous_version.extract_completed
                    ):
                        can_resume_previous = True

        # if not accepting input
        if self.noinput:
            # skip if up-to-date
            if latest_version.update_completed:
                self.log('Your database is already up-to-date.')
                exit(0)
            else:
                # set to resume if update to latest is stalled or can resume previous
                self.resume = latest_version.update_stalled or can_resume_previous
        else:
            # set up prompt
            prompt_context = dict(
                latest_version=latest_version,
                previous_version=previous_version,
                can_resume_previous=can_resume_previous,
                since_previous_release=since_previous_release,
            )
            # render prompt string
            prompt = render_to_string(
                'calaccess_raw/updatecalaccessrawdata.txt',
                prompt_context,
            )
            # if the user can resume (latest or previous)
            if latest_version.update_stalled or can_resume_previous:
                # and initially confirms
                if self.confirm_proceed(prompt):
                    # go into resume mode
                    self.resume = True
                # if user does not initially confirm
                else:
                    self.resume = False
                    # if can resume latest (but doesn't) prompt restart
                    if latest_version.update_stalled:
                        confirm_restart = self.confirm_proceed(
                            'Do you want re-start your update?\n'
                        )
                    # if can resume previous (but doesn't) prompt for update to latest
                    elif can_resume_previous:
                        confirm_restart = self.confirm_proceed(
                            'Do you want to update to the latest version?\n'
                        )
                    # cancel if user doesn't confirm restart
                    if not confirm_restart:
                        raise CommandError("Update cancelled")
            # if can't resume
            else:
                self.resume = False
                # cancel if user doesn't initially confirm
                if not self.confirm_proceed(prompt):
                    raise CommandError("Update cancelled")

        # set to force restart if the user could have resumed but didn't
        # or no new version is available
        self.force_restart = (
            (
                latest_version.update_stalled or
                can_resume_previous or
                not created
            ) and not self.resume
        )

        # set the version to be updated
        if not self.force_restart and can_resume_previous:
            self.version = previous_version
        else:
            self.version = latest_version

        if self.verbosity and not self.test_mode:
            if self.resume:
                self.header(
                    "Resuming update to {:%m-%d-%Y %H:%M:%S} snapshot".format(
                        self.version.release_datetime
                    )
                )
            else:
                self.header(
                    "Updating to {:%m-%d-%Y %H:%M:%S} snapshot".format(
                        self.version.release_datetime
                    )
                )

        # if not resuming, reset the update start datetime
        if not self.resume:
            self.version.update_start_datetime = now()
        # if restarting, also erase all the other datetimes
        if self.force_restart:
            self.version.update_finish_datetime = None
            self.version.download_start_datetime = None
            self.version.download_finish_datetime = None
            self.version.extract_start_datetime = None
            self.version.extract_finish_datetime = None

        # save here in case the command doesn't finish
        self.version.save()

        # handle download
        if self.test_mode:
            pass
        # can only download the latest version:
        elif self.version == latest_version:
            # if download completed and not forcing restart
            if self.version.download_completed and not self.force_restart:
                self.log('Already downloaded.')
            # otherwise try downloading
            else:
                self.download()
                if self.verbosity:
                    self.duration()

        # handle file extraction
        if self.test_mode:
            handle_command(TestExtractCommand, verbosity=self.verbosity)
        # if the last extract of the version completed and not forcing restart
        elif self.version.extract_completed and not self.force_restart:
            self.log('Already extracted.')
        else:
            # if the zip isn't there
            if not os.path.exists(self.zip_path):
                # if updating to the lastest
                if self.version == latest_version:
                    self.log(
                        '%s not found. Re-downloading before extracting.' % self.zip_path
                    )
                    self.download()
                    if self.verbosity:
                        self.duration()
                else:
                    raise CommandError(
                        'Incomplete file extraction and %s not found.' % self.zip_path
                    )
            # now extract
            call_command(
                'extractcalaccessrawfiles',
                keep_files=self.keep_files
            )
            if self.verbosity:
                self.duration()

        # refresh the version (to get timestamp field values)
        self.version.refresh_from_db()

        self.clean()
        if self.verbosity:
            self.duration()

        # if archive setting is enabled, zip up all of the csv and error logs
        if getattr(settings, 'CALACCESS_STORE_ARCHIVE', False):
            # skip if resuming and the clean zip file is already saved
            if self.resume and bool(self.version.clean_zip_archive):
                self.log('File already zipped.')
            else:
                self.archive()
                if self.verbosity:
                    self.duration()

        self.load()
        if self.verbosity:
            self.duration()

        # store update finish time
        self.version.update_finish_datetime = now()
        # and save the RawDataVersion
        self.version.save()

        if self.verbosity:
            self.success("Done!")

    def download(self):
        """
        Try downloading the zip. Wait and re-try if certain CommandErrors are raised.
        """
        try:
            call_command(
                "downloadcalaccessrawdata",
                verbosity=self.verbosity,
                noinput=True,
                restart=self.force_restart,
            )
        except CommandError as e:
            # if the expected and actual zip size are not the same
            if (
                'expected' in e.message.lower() or
                'modified' in e.message.lower() or
                'version' in e.message.lower()
            ):
                logger.debug('Waiting five minutes before re-trying')
                # wait five minutes
                sleep(300)
                # then try again
                call_command(
                    "downloadcalaccessrawdata",
                    verbosity=self.verbosity,
                    noinput=True,
                    # force a restart on second try
                    restart=True,
                )

    def clean(self):
        """
        Clean up the raw data files from the state so they are ready to get loaded into the database.
        """
        if self.verbosity:
            self.header("Cleaning data files")

        tsv_list = [
            f for f in os.listdir(self.tsv_dir) if '.TSV' in f.upper()
        ]

        if self.resume:
            # get finished clean command logs of last update
            prev_cleaned = [
                x.file_name + '.TSV'
                for x in self.version.files.filter(
                    clean_finish_datetime__isnull=False
                )
            ]
            self.log("{} files already cleaned.".format(len(prev_cleaned)))
            # remove these from tsv_list
            tsv_list = [x for x in tsv_list if x not in prev_cleaned]

        # Loop through all the files in the source directory
        if self.verbosity:
            tsv_list = progress.bar(tsv_list)
        for name in tsv_list:
            call_command(
                "cleancalaccessrawfile",
                name,
                verbosity=self.verbosity,
                keep_file=self.keep_files,
            )

    def archive(self):
        """
        Zip up and archive the .csv and error log files.
        """
        if self.verbosity:
            self.header("Zipping cleaned files")
        # Remove previous zip file
        self.version.clean_zip_archive.delete()
        clean_zip_path = os.path.join(self.data_dir, 'clean.zip')

        # enable zipfile compression
        compression = ZIP_DEFLATED

        try:
            zf = ZipFile(clean_zip_path, 'w', compression, allowZip64=True)
        except RuntimeError:
            logger.error('Zip file cannot be compressed (check zlib module).')
            compression = ZIP_STORED
            zf = ZipFile(clean_zip_path, 'w', compression, allowZip64=True)

        # loop over and save files in csv dir
        for f in os.listdir(self.csv_dir):
            if self.verbosity > 2:
                self.log(" Adding %s to zip" % f)
            csv_path = os.path.join(self.csv_dir, f)
            zf.write(csv_path, f)

        # same for errors dir
        errors_dir = os.path.join(self.data_dir, 'log')
        for f in os.listdir(errors_dir):
            if self.verbosity > 2:
                self.log(" Adding %s to zip" % f)
            error_path = os.path.join(errors_dir, f)
            zf.write(error_path, f)

        # close the zip file
        zf.close()
        if self.verbosity > 2:
            self.log(" All files zipped")

        if not self.test_mode:
            # save the clean zip size
            self.version.clean_zip_size = os.path.getsize(clean_zip_path)
            with open(clean_zip_path, 'rb') as zf:
                # Save the zip on the raw data version
                if self.verbosity > 2:
                    self.log(" Archiving zip")
                self.version.clean_zip_archive.save(
                    os.path.basename(clean_zip_path), File(zf)
                )
            if self.verbosity > 2:
                self.log(" Zip archived.")

    def load(self):
        """
        Loads the cleaned up csv files into the database.
        """
        if self.verbosity:
            self.header("Loading data files")

        model_list = [
            x for x in get_model_list() if os.path.exists(x.objects.get_csv_path())
        ]

        if self.resume:
            # get finished load command logs of last update
            prev_loaded = [
                x.file_name
                for x in self.version.files.filter(
                    load_finish_datetime__isnull=False
                )
            ]
            self.log("{} models already loaded.".format(len(prev_loaded)))
            # remove these from model_list
            model_list = [x for x in model_list if x._meta.db_table not in prev_loaded]

        if self.verbosity:
            model_list = progress.bar(model_list)
        for model in model_list:
            call_command(
                "loadcalaccessrawfile",
                model.__name__,
                verbosity=self.verbosity,
                keep_file=self.keep_files,
                app_name=self.app_name,
            )
