#!/opt/anaconda1anaconda2anaconda3/bin/python
# -*- coding: utf-8

import os
import sys
import glob
import stat
import uuid
import shutil
import zipfile
import argparse
import requests
import subprocess
from distutils.spawn import find_executable

import anvio
import anvio.utils as utils
import anvio.terminal as terminal
import anvio.filesnpaths as filesnpaths

from anvio.errors import UpgradeError

run = terminal.Run()
progress = terminal.Progress()

class Upgrader(object):
    def __init__(self, args, run=run, progress=progress):
        A = lambda x: args.__dict__[x] if x in args.__dict__ else None
        self.run = run
        self.progress = progress
        self.args = args
        self.repository = A('repository')


    def upgrade(self):
        self.version = self.get_version()
        self.releases = self.get_releases()
        self.target_version = self.get_target_version()

        if self.version == self.target_version:
            self.run.info_single('You are already using the latest version.', nl_before=1, nl_after=1, mc='green')
            return

        self.dist = self.get_dist_info()
        self.working_dir = filesnpaths.get_temp_directory_path()
        self.run.info('Temporary working directory', self.working_dir)
        self.download()
        self.delete_dist_files()
        self.move_files()

        self.run.info_single('Upgrade was succesful. Testing new version with "anvi-profile --version", \
            you may also want to use "anvi-self-test" to test it further.', nl_before=1, nl_after=1, mc='green')
        subprocess.call(['anvi-profile', '--version'])


    def get_releases(self):
        releases = {}
        url = 'https://api.github.com/repos/%s/tags' % self.repository

        if 'ANVIO_GITHUB_API_TOKEN' in os.environ:
            url = url + '?access_token=' + os.environ['ANVIO_GITHUB_API_TOKEN']

        response = requests.get(url).json()

        if 'message' in response and 'documentation_url' in response:
            raise UpgradeError('Github API returned error: %s, %s' % (response['message'], response['documentation_url']))

        for release in response:
            releases[release['name']] = release

        return releases


    def get_version(self):
        version = anvio.__version__

        if '-master' in version:
            raise UpgradeError("You are already using the master repository. To upgrade, run `git pull`.")

        self.run.info('Installed version', '%s' % (version))

        return version


    def get_dist_info(self):
        dist = {}

        bin_file = find_executable('anvi-profile')

        with open(bin_file, 'r') as f:
            shebang = f.readline().strip()
            dist['shebang'] = shebang

        dist['bin_path'] = os.path.dirname(bin_file)
        dist['bin_stat'] = os.stat(bin_file)

        dist['module_path'] = os.path.dirname(anvio.__file__)
        dist['module_stat'] = os.stat(anvio.__file__)
        dist['module_dir_stat'] = os.stat(dist['module_path'])

        script_file = find_executable('anvi-script-reformat-fasta')
        dist['sandbox_path'] = os.path.dirname(script_file)
        dist['sandbox_stat'] = os.stat(script_file)

        self.run.info('Program location found', dist['bin_path'])
        self.run.info('Module location found', dist['module_path'])
        self.run.info('Scripts location found', dist['sandbox_path'])

        return dist


    def check_if_dist_writable(self):
        rand_filename = str(uuid.uuid4())

        for target in ['bin', 'module', 'sandbox']:
            target_path = os.path.join(self.dist['%s_path' % target], rand_filename)

            try:
                with open(target_path, 'w') as f:
                    f.write('test')

                self.restore_file_stats(target_path, target)

                os.remove(target_path)
            except:
                raise ConfigError("This script can not write files in anvi'o installation directories, try running with 'sudo'")


    def get_target_version(self):
        version_parts = self.version.split('.')

        if len(version_parts) > 1:
            major, minor = int(version_parts[0]), int(version_parts[1])
        else:
            major, minor = int(version_parts[0]), 0

        while ('v%d.%d' % (major, minor+1)) in self.releases:
            minor += 1

        self.run.info('Target version', '%d.%d' % (major, minor))

        return '%d.%d' % (major, minor)


    def download(self):
        self.get_dist_info()

        current_url = self.releases['v' + self.version]['zipball_url']
        current_zip_path = os.path.join(self.working_dir, self.version + '.zip')

        target_url = self.releases['v' + self.target_version]['zipball_url']
        target_zip_path = os.path.join(self.working_dir, self.target_version + '.zip')

        utils.download_file(current_url, current_zip_path, progress=self.progress, run=self.run)
        utils.download_file(target_url, target_zip_path, progress=self.progress, run=self.run)

        self.progress.new('Extracting Archives')

        with zipfile.ZipFile(current_zip_path, "r") as zip_ref:
            self.current_version_dir = current_zip_path[:-4] + '/' + zip_ref.namelist()[0]
            zip_ref.extractall(current_zip_path[:-4])
            self.progress.update('v' + self.version)

        with zipfile.ZipFile(target_zip_path, "r") as zip_ref:
            self.target_version_dir = target_zip_path[:-4] + '/' + zip_ref.namelist()[0]
            zip_ref.extractall(target_zip_path[:-4])
            self.progress.update('v' + self.target_version)

        self.progress.end()


    def delete_dist_files(self):
        self.progress.new('Deleting current version files')
        dist_files = []

        for script_type in ['bin', 'sandbox']:
            for file_path in glob.glob(os.path.join(self.current_version_dir, script_type, 'anvi-*')):
                file_name = os.path.basename(file_path)
                dist_file_path = os.path.join(self.dist['%s_path' % script_type], file_name)
                dist_files.append(dist_file_path)
                self.progress.update(file_name)
                
        for path_name, sub_directories, file_list in os.walk(os.path.join(self.current_version_dir, 'anvio')):
            relative_path = path_name[len(os.path.join(self.current_version_dir, 'anvio')) + 1:]

            for file_name in file_list + sub_directories:
                file_path = os.path.join(path_name, file_name)
                dist_file_path = os.path.join(self.dist['module_path'], relative_path, file_name)
                dist_files.append(dist_file_path)

        for dist_file_path in dist_files:
            if not filesnpaths.is_file_exists(dist_file_path, dont_raise=True):
                raise UpgradeError("File %s does not exists" % dist_file_path)

        for dist_file_path in dist_files:
            self.progress.update(os.path.basename(dist_file_path))
            if os.path.isfile(dist_file_path):
                os.remove(dist_file_path)
            else:
                try:
                    os.rmdir(dist_file_path)
                except:
                    # Directory not empty
                    pass

        self.progress.end()


    def restore_file_stats(self, target_path, target):
        os.chmod(target_path, self.dist['%s_stat' % target].st_mode)
        os.chown(target_path, self.dist['%s_stat' % target].st_uid, self.dist['%s_stat' % target].st_gid)


    def move_files(self):
        self.progress.new('Moving target version files')

        for script_type in ['bin', 'sandbox']:
            for file_path in glob.glob(os.path.join(self.target_version_dir, script_type, 'anvi-*')):
                file_name = os.path.basename(file_path)
                dist_file_path = os.path.join(self.dist['%s_path' % script_type], file_name)

                with open(file_path, 'r') as f_orig:
                    lines = f_orig.readlines()

                    # delete shebang if any
                    if lines[0].startswith('#!'):
                        lines.pop(0)

                    with open(dist_file_path, 'w') as f_dest:
                        # write distribution specific shebang
                        f_dest.write(self.dist['shebang'] + '\n')
                        f_dest.write(''.join(lines))

                self.progress.update(file_name)
                self.restore_file_stats(dist_file_path, script_type)
        
        for path_name, sub_directories, file_list in os.walk(os.path.join(self.target_version_dir, 'anvio')):
            relative_path = path_name[len(os.path.join(self.target_version_dir, 'anvio')) + 1:]
            absolute_dist_path = os.path.join(self.dist['module_path'], relative_path)

            if os.path.isdir(path_name) and not os.path.exists(absolute_dist_path):
                os.makedirs(absolute_dist_path)
                self.restore_file_stats(absolute_dist_path, 'module_dir')

            for file_name in file_list:
                self.progress.update(file_name)
                file_path = os.path.join(path_name, file_name)
                dist_file_path = os.path.join(absolute_dist_path, file_name)

                shutil.copy(file_path, dist_file_path)
                self.restore_file_stats(dist_file_path, 'module')

        self.progress.end()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Download and install minor releases of anvi'o from a Github repository.")
    parser.add_argument(*anvio.A('repository'), **anvio.K('repository'))
    args, unknown = parser.parse_known_args()

    try:
        upgrader = Upgrader(args)
        upgrader.upgrade()
    except UpgradeError as e:
        print(e)
        sys.exit(-1)
