#!/usr/bin/env python
# ***** BEGIN LICENSE BLOCK *****
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this file,
# You can obtain one at http://mozilla.org/MPL/2.0/.
# ***** END LICENSE BLOCK *****


import copy
import os
import pprint
import sys
import urllib

# load modules from parent dir
sys.path.insert(1, os.path.dirname(os.path.dirname(sys.path[0])))

from mozharness.base.python import PreScriptAction
from mozharness.mozilla.buildbot import TBPL_SUCCESS, TBPL_WARNING, EXIT_STATUS_DICT
from mozharness.mozilla.testing.firefox_ui_tests import (
    FirefoxUIUpdateTests,
    firefox_ui_update_config_options
)


# Command line arguments for release update tests
firefox_ui_update_release_config_options = [
    [['--build-number'], {
        'dest': 'build_number',
        'help': 'Build number of release, eg: 2',
    }],
    [['--limit-locales'], {
        'dest': 'limit_locales',
        'default': -1,
        'type': int,
        'help': 'Limit the number of locales to run.',
    }],
    [['--release-update-config'], {
        'dest': 'release_update_config',
        'help': 'Name of the release update verification config file to use.',
    }],
    [['--this-chunk'], {
        'dest': 'this_chunk',
        'default': 1,
        'help': 'What chunk of locales to process.',
    }],
    [['--tools-repo'], {
        'dest': 'tools_repo',
        'default': 'http://hg.mozilla.org/build/tools',
        'help': 'Which tools repo to check out',
    }],
    [['--tools-tag'], {
        'dest': 'tools_tag',
        'help': 'Which revision/tag to use for the tools repository.',
    }],
    [['--total-chunks'], {
        'dest': 'total_chunks',
        'default': 1,
        'help': 'Total chunks to dive the locales into.',
    }],
] + copy.deepcopy(firefox_ui_update_config_options)


class ReleaseFirefoxUIUpdateTests(FirefoxUIUpdateTests):

    def __init__(self):
        all_actions = [
            'clobber',
            'checkout',
            'create-virtualenv',
            'query_minidump_stackwalk',
            'read-release-update-config',
            'run-tests',
        ]

        super(ReleaseFirefoxUIUpdateTests, self).__init__(
            all_actions=all_actions,
            default_actions=all_actions,
            config_options=firefox_ui_update_release_config_options,
            append_env_variables_from_configs=True,
        )

        self.tools_repo = self.config.get('tools_repo')
        self.tools_tag = self.config.get('tools_tag')

        assert self.tools_repo and self.tools_tag, \
            'Without the "--tools-tag" we can\'t clone the releng\'s tools repository.'

        self.limit_locales = int(self.config.get('limit_locales'))

        # This will be a list containing one item per release based on configs
        # from tools/release/updates/*cfg
        self.releases = None

    def checkout(self):
        """
        We checkout the tools repository and update to the right branch
        for it.
        """
        dirs = self.query_abs_dirs()

        super(ReleaseFirefoxUIUpdateTests, self).checkout()

        self.vcs_checkout(
            repo=self.tools_repo,
            dest=dirs['abs_tools_dir'],
            branch=self.tools_tag,
            vcs='hg'
        )

    def query_abs_dirs(self):
        if self.abs_dirs:
            return self.abs_dirs

        abs_dirs = super(ReleaseFirefoxUIUpdateTests, self).query_abs_dirs()
        dirs = {
            'abs_tools_dir': os.path.join(abs_dirs['abs_work_dir'], 'tools'),
        }

        for key in dirs:
            if key not in abs_dirs:
                abs_dirs[key] = dirs[key]
        self.abs_dirs = abs_dirs

        return self.abs_dirs

    def read_release_update_config(self):
        '''
        Builds a testing matrix based on an update verification configuration
        file under the tools repository (release/updates/*.cfg).

        Each release info line of the update verification files look similar to the following.

        NOTE: This shows each pair of information as a new line but in reality
        there is one white space separting them. We only show the values we care for.

            release="38.0"
            platform="Linux_x86_64-gcc3"
            build_id="20150429135941"
            locales="ach af ... zh-TW"
            channel="beta-localtest"
            from="/firefox/releases/38.0b9/linux-x86_64/%locale%/firefox-38.0b9.tar.bz2"
            ftp_server_from="http://archive.mozilla.org/pub"

        We will store this information in self.releases as a list of releases.

        NOTE: We will talk of full and quick releases. Full release info normally contains a subset
        of all locales (except for the most recent releases). A quick release has all locales,
        however, it misses the fields 'from' and 'ftp_server_from'.
        Both pairs of information complement each other but differ in such manner.
        '''
        dirs = self.query_abs_dirs()
        assert os.path.exists(dirs['abs_tools_dir']), \
            'Without the tools/ checkout we can\'t use releng\'s config parser.'

        if self.config.get('release_update_config'):
            # The config file is part of the tools repository. Make sure that if specified
            # we force a revision of that repository to be set.
            if self.tools_tag is None:
                self.fatal('Make sure to specify the --tools-tag')

            self.release_update_config = self.config['release_update_config']

        # Import the config parser
        sys.path.insert(1, os.path.join(dirs['abs_tools_dir'], 'lib', 'python'))
        from release.updates.verify import UpdateVerifyConfig

        uvc = UpdateVerifyConfig()
        config_file = os.path.join(dirs['abs_tools_dir'], 'release', 'updates',
                                   self.config['release_update_config'])
        uvc.read(config_file)
        if not hasattr(self, 'update_channel'):
            self.update_channel = uvc.channel

        # Filter out any releases that are less than Gecko 38
        uvc.releases = [r for r in uvc.releases
                        if int(r['release'].split('.')[0]) >= 38]

        temp_releases = []
        for rel_info in uvc.releases:
            # This is the full release info
            if 'from' in rel_info and rel_info['from'] is not None:
                # Let's find the associated quick release which contains the remaining locales
                # for all releases except for the most recent release which contain all locales
                quick_release = uvc.getRelease(build_id=rel_info['build_id'], from_path=None)
                if quick_release != {}:
                    rel_info['locales'] = sorted(rel_info['locales'] + quick_release['locales'])
                temp_releases.append(rel_info)

        uvc.releases = temp_releases
        chunked_config = uvc.getChunk(
            chunks=int(self.config['total_chunks']),
            thisChunk=int(self.config['this_chunk'])
        )

        self.releases = chunked_config.releases

    @PreScriptAction('run-tests')
    def _pre_run_tests(self, action):
        assert ('release_update_config' in self.config or
                self.installer_url or self.installer_path), \
            'Either specify --update-verify-config, --installer-url or --installer-path.'

    def run_tests(self):
        dirs = self.query_abs_dirs()

        # We don't want multiple outputs of the same environment information. To prevent
        # that, we can't make it an argument of run_command and have to print it on our own.
        self.info('Using env: {}'.format(pprint.pformat(self.query_env())))

        results = {}

        locales_counter = 0
        for rel_info in sorted(self.releases, key=lambda release: release['build_id']):
            build_id = rel_info['build_id']
            results[build_id] = {}

            self.info('About to run {buildid} {path} - {num_locales} locales'.format(
                buildid=build_id,
                path=rel_info['from'],
                num_locales=len(rel_info['locales'])
            ))

            # Each locale gets a fresh port to avoid address in use errors in case of
            # tests that time out unexpectedly.
            marionette_port = 2827
            for locale in rel_info['locales']:
                locales_counter += 1
                self.info('Running {buildid} {locale}'.format(buildid=build_id,
                                                              locale=locale))

                if self.limit_locales > -1 and locales_counter > self.limit_locales:
                    self.info('We have reached the limit of locales we were intending to run')
                    break

                if self.config['dry_run']:
                    continue

                # Determine from where to download the file
                installer_url = '{server}/{fragment}'.format(
                    server=rel_info['ftp_server_from'],
                    fragment=urllib.quote(rel_info['from'].replace('%locale%', locale))
                )
                installer_path = self.download_file(
                    url=installer_url,
                    parent_dir=dirs['abs_work_dir']
                )

                binary_path = self.install_app(app=self.config.get('application'),
                                               installer_path=installer_path)

                marionette_port += 1

                retcode = self.run_test(
                    binary_path=binary_path,
                    env=self.query_env(avoid_host_env=True),
                    marionette_port=marionette_port,
                )

                self.uninstall_app()

                # Remove installer which is not needed anymore
                self.info('Removing {}'.format(installer_path))
                os.remove(installer_path)

                if retcode:
                    self.warning('FAIL: {} has failed.'.format(sys.argv[0]))

                    base_cmd = 'python {command} --firefox-ui-branch {branch} ' \
                        '--release-update-config {config} --tools-tag {tag}'.format(
                            command=sys.argv[0],
                            branch=self.firefox_ui_branch,
                            config=self.release_update_config,
                            tag=self.tools_tag
                        )

                    for config in self.config['config_files']:
                        base_cmd += ' --cfg {}'.format(config)

                    if self.symbols_url:
                        base_cmd += ' --symbols-path {}'.format(self.symbols_url)

                    base_cmd += ' --installer-url {}'.format(installer_url)

                    self.info('You can run the *specific* locale on the same machine with:')
                    self.info(base_cmd)

                    self.info('You can run the *specific* locale on *your* machine with:')
                    self.info('{} --cfg developer_config.py'.format(base_cmd))

                results[build_id][locale] = retcode

                self.info('Completed {buildid} {locale} with return code: {retcode}'.format(
                    buildid=build_id,
                    locale=locale,
                    retcode=retcode))

            if self.limit_locales > -1 and locales_counter > self.limit_locales:
                break

        # Determine which locales have failed and set scripts exit code
        exit_status = TBPL_SUCCESS
        for build_id in sorted(results.keys()):
            failed_locales = []
            for locale in sorted(results[build_id].keys()):
                if results[build_id][locale] != 0:
                    failed_locales.append(locale)

            if failed_locales:
                if exit_status == TBPL_SUCCESS:
                    self.info('\nSUMMARY - Failed locales for {}:'.format(self.cli_script))
                    self.info('====================================================')
                    exit_status = TBPL_WARNING

                self.info(build_id)
                self.info('  {}'.format(', '.join(failed_locales)))

        self.return_code = EXIT_STATUS_DICT[exit_status]


if __name__ == '__main__':
    myScript = ReleaseFirefoxUIUpdateTests()
    myScript.run_and_exit()