#!/usr/bin/env python3
import sys
import os
import argparse
import tempfile
import shutil
import subprocess
import re
import fileinput
from collections import OrderedDict
import json
import hashlib
import traceback

# Hack to combine two argparse formatters
class CustomFormatter(argparse.ArgumentDefaultsHelpFormatter, argparse.RawDescriptionHelpFormatter):
    pass

parser = argparse.ArgumentParser(
    description='Build a Zotero XPI',
    formatter_class=CustomFormatter,
    epilog='''
Example: build_xpi -b 5.0 -x 5.0.1 -z
  - Builds from the 5.0 branch
  - Builds zotero-build.xpi and update-build.rdf
  - Points update-build.rdf to https://download.zotero.org/extension/zotero-5.0.1.xpi
  - Points install.rdf to https://www.zotero.org/download/update.rdf

Example: build_xpi -b 5.0 -x 5.0b2 -r beta -z
  - Builds from the 5.0 branch
  - Builds zotero-build.xpi and update-build.rdf
  - Points update-build.rdf to https://download.zotero.org/extension/zotero-5.0b2.xpi
  - Points install.rdf to https://www.zotero.org/download/update-beta.rdf

Example: build_xpi -b master -c alpha -x 5.0-alpha -r 5.0-branch --xpi-dir dev -z
  - Builds from the master branch
  - Builds zotero-build.xpi and update-build.rdf
  - Points update-build.rdf to https://download.zotero.org/extension/dev/zotero-5.0-alpha.xpi
  - Points install.rdf to https://zotero.org/download/dev/update-5.0-branch.rdf''')

parser.add_argument('-b', '--branch', default='master', help='Git branch or tag to build from')
parser.add_argument('-c', '--channel', default='source', help='channel to add to dev build version number (e.g., "beta" for "5.0-beta.3+a5f28ca8"), or "release" or "source" to skip')
parser.add_argument('--build-suffix', metavar='SUFFIX', default='build', help='suffix of output XPI')
parser.add_argument('--xpi-suffix', '-x', metavar='SUFFIX', default='', help='suffix of XPI to reference in update.rdf')
parser.add_argument('--rdf-suffix', '-r', metavar='SUFFIX', default='', help='suffix of update.rdf file to reference in install.rdf (e.g., "beta" for "update-beta.rdf")')
parser.add_argument('--xpi-dir', metavar='DIR', default='', help='extra directory to point to when referencing the XPI in update.rdf')
parser.add_argument('--repo-url', metavar='URL', default='https://github.com/zotero/zotero', help='Git URL to pull from')
parser.add_argument('--source-dir', '-s', metavar='DIR', help='Directory to build from (overrides --repo-url and --branch)')
parser.add_argument('--zip', '-z', action='store_true', help="Create XPI instead of leaving files in build/staging")

args = parser.parse_args()

def main():
    try:
        if args.xpi_suffix:
            args.xpi_suffix = "-" + args.xpi_suffix
        if args.rdf_suffix:
            args.rdf_suffix = "-" + args.rdf_suffix
        if args.build_suffix:
            args.build_suffix = "-" + args.build_suffix
        
        root_dir = os.path.dirname(os.path.realpath(__file__))
        
        # Use BUILD_DIR environmental variable if present
        build_dir = os.environ.get('BUILD_DIR', os.path.join(root_dir, 'build'))
        tmp_dir = os.path.join(build_dir, 'tmp')
        
        if not os.path.isdir(build_dir):
            raise Exception(build_dir + " is not a directory")
        
        src_dir = args.source_dir
        if src_dir:
            if not os.path.isdir(src_dir):
                raise Exception(src_dir + " is not a directory")
            
            log("Using existing source directory of " + src_dir)
            os.chdir(src_dir)
        # Create 'zotero' directory inside build directory if necessary
        # TEMP: For now, use 'zotero-5.0' to keep separate from 4.0 checkout
        else:
            src_dir = os.path.join(build_dir, 'zotero')
            if not os.path.isdir(src_dir):
                log(src_dir + " does not exist -- creating")
                os.mkdir(src_dir)
            
            os.chdir(src_dir)
            
            # If 'zotero' directory doesn't contain Zotero code, clone the repo into it
            if not os.path.exists(os.path.join(src_dir, '.git')):
                log("Cloning Zotero into " + src_dir)
                subprocess.check_call([
                    'git', 'clone', '--recursive', args.repo_url, '.'
                ])
            
            subprocess.check_call(['git', 'fetch', 'origin'])
            # Determine if 'branch' is a reference or a commit
            code = subprocess.call(['git', 'show-ref', '--verify', '-q',
                                    'refs/heads/{0}'.format(args.branch)])
            # If a reference, use git checkout origin/{branch} instead of just
            # the commit in order to get a detached HEAD with the latest
            # fetched changes. (There's probably a better way to do this.)
            if code == 0:
                checkout = 'origin/' + args.branch
            else:
                checkout = args.branch
            subprocess.check_call(['git', 'checkout', checkout])
            
            subprocess.check_call('git submodule update --init', shell=True)
            # Uncomment to ignore fixed submodule versions and always pull latest versions
            #subprocess.check_call('git submodule foreach git pull origin master', shell=True)
        
        if not os.path.exists('install.rdf'):
            raise FileNotFoundError("install.rdf not found in {0}".format(src_dir))
        
        # Extract version number from install.rdf
        with open('install.rdf') as f:
            rdf = f.read()
        m = re.search('version>([0-9].+)\\.SOURCE</', rdf)
        if not m:
            raise Exception("Version number not found in install.rdf")
        version = m.group(1)
        commit_hash = subprocess.check_output([
            'git', 'rev-parse', '--short', 'HEAD'
        ]).decode(encoding="utf-8").strip()
        
        # Determine build targets
        target_xpi_file = os.path.join(
            build_dir, 'zotero' + args.build_suffix + '.xpi'
        )
        target_update_file = os.path.join(
            build_dir, 'update' + args.build_suffix + '.rdf'
        )
        staging_dir = os.path.join(build_dir, 'staging')
        
        # Delete any existing build targets
        try:
            os.remove(target_xpi_file)
        except OSError:
            pass
        try:
            os.remove(target_update_file)
        except OSError:
            pass
        if os.path.exists(staging_dir):
            shutil.rmtree(staging_dir)
        
        # Remove tmp build directory if it already exists
        if os.path.exists(tmp_dir):
            shutil.rmtree(tmp_dir)
        os.mkdir(tmp_dir)
        tmp_src_dir = os.path.join(tmp_dir, 'zotero')
        
        # Export a clean copy of the source tree
        subprocess.check_call([
            'rsync',
            '-a',
            # Exclude hidden files
            '--exclude', '.*',
            '.' + os.sep,
            tmp_src_dir + os.sep
        ])
        
        # Make sure rsync worked
        d = os.path.join(tmp_src_dir, 'chrome')
        if not os.path.isdir(d):
            raise FileNotFoundError(d + " not found")
        
        log("Deleting CSL locale support files")
        subprocess.check_call([
            'find',
            os.path.normpath(tmp_src_dir + '/chrome/content/zotero/locale/csl/'),
            '-mindepth', '1',
            '!', '-name', '*.xml',
            '!', '-name', 'locales.json',
            '-print',
            '-delete'
        ])
        
        # Delete test files
        shutil.rmtree(os.path.join(tmp_src_dir, 'test'))
        
        # Delete styles build script
        os.remove(os.path.join(tmp_src_dir, 'styles', 'update'))
        
        translators_dir = os.path.join(tmp_src_dir, 'translators')
        
        # Move deleted.txt out of translators directory
        f = os.path.join(translators_dir, 'deleted.txt')
        if os.path.exists(f):
            shutil.move(f, tmp_src_dir)
        
        # Build translator index
        index = OrderedDict()
        for fn in sorted((fn for fn in os.listdir(translators_dir)), key=str.lower):
            with open(os.path.join(translators_dir, fn), 'r') as f:
                contents = f.read()
            # Parse out the JSON metadata block
            m = re.match('^\s*{[\S\s]*?}\s*?[\r\n]', contents)
            if not m:
                raise Exception("Metadata block not found in " + f)
            metadata = json.loads(m.group(0))
            index[metadata["translatorID"]] = {
                "fileName": fn,
                "label": metadata["label"],
                "lastUpdated": metadata["lastUpdated"]
            }
        
        # Unminify translator framework lines
        with open(os.path.join(root_dir, 'zotero-transfw', 'framework.js')) as f:
            framework_contents = f.read()
        with open(os.path.join(translators_dir, 'fw.js'), 'w') as f:
            f.write(
                "/*********************** BEGIN FRAMEWORK ***********************/\n"
                + framework_contents
                + "\n/*********************** END FRAMEWORK ***********************/\n"
            )
        os.chdir(translators_dir)
        subprocess.check_call(
            "perl -pe 's/.+FW LINE.+/`cat fw.js`/ge' -i *.js", shell=True
        )
        os.remove('fw.js')
        
        # Write translator index as JSON file
        with open(os.path.join(tmp_src_dir, 'translators.json'), 'w') as f:
            json.dump(index, f, indent=True, ensure_ascii=False)
        
        install_file = os.path.join(tmp_src_dir, 'install.rdf')
        update_file = os.path.join(tmp_src_dir, 'update.rdf')
        
        log_line()
        log('Original install.rdf:')
        dump_file(install_file)
        if args.zip:
            log('Original update.rdf:\n')
            dump_file(update_file)
            log_line()
        
        # Modify install.rdf and update.rdf as necessary
        
        # The dev build revision number is stored in build/lastrev-{version}-{channel}.
        #
        # If we're including it, get the current version number and increment it.
        if args.channel not in ["release", "source"]:
            lastrev_file = os.path.join(
                build_dir, 'lastrev-{0}-{1}'.format(version, args.channel)
            )
            if not os.path.exists(lastrev_file):
                with open(lastrev_file, 'w') as f:
                    f.write("0")
                    rev = 1
            else:
                with open(lastrev_file, 'r') as f:
                    rev = int(f.read()) + 1
        
        if args.channel == "release":
            rev_sub_str = ""
        elif args.channel == "source":
            rev_sub_str = ".SOURCE.{0}".format(commit_hash)
        else:
            rev_sub_str = "-{0}.{1}+{2}".format(args.channel, str(rev), commit_hash)
        if args.xpi_dir:
            xpi_dir = args.xpi_dir + '/'
        else:
            xpi_dir = ''
        # Update install.rdf and update.rdf
        for line in fileinput.FileInput(install_file, inplace=1):
            line = line.replace('.SOURCE', rev_sub_str)
            line = line.replace(
                'update-source.rdf',
                xpi_dir + 'update' + args.rdf_suffix + '.rdf'
            )
            print(line, file=sys.stdout, end='')
        for line in fileinput.FileInput(update_file, inplace=1):
            line = line.replace(".SOURCE", rev_sub_str)
            line = line.replace(
                'zotero.xpi',
                xpi_dir + 'zotero' + args.xpi_suffix + '.xpi'
            )
            print(line, file=sys.stdout, end='')
        
        log_line()
        log('Modified install.rdf:\n')
        dump_file(install_file)
        
        # Create XPI
        if args.zip:
            # Move update.rdf out of code root
            shutil.move(update_file, tmp_dir)
            tmp_update_file = os.path.join(tmp_dir, 'update.rdf')
            
            os.chdir(tmp_src_dir)
            tmp_xpi_file = os.path.join(tmp_dir, 'zotero' + args.build_suffix + '.xpi')
            subprocess.check_call(['zip', '-r', tmp_xpi_file, '.'])
            
            # Add SHA1 of XPI to update.rdf
            sha1 = sha1file(tmp_xpi_file)
            for line in fileinput.FileInput(tmp_update_file, inplace=1):
                line = line.replace("sha1:", "sha1:" + sha1)
                print(line, file=sys.stdout, end='')
            
            log('Modified update.rdf:\n')
            dump_file(tmp_update_file)
            log_line()
            
            # Move files to build directory
            os.rename(tmp_xpi_file, target_xpi_file)
            os.rename(tmp_update_file, target_update_file)
            
            log("")
            log("zotero{0}.xpi and update{0}.rdf saved to {1}".format(args.build_suffix, build_dir))
            log("")
        # Leave unzipped in staging directory
        else:
            # Don't create update.rdf
            os.remove(update_file)
            
            # Move source files to staging
            shutil.move(tmp_src_dir, staging_dir)
            
            log("")
            log("Build files saved to {0}".format(staging_dir))
            log("")
        
        # Update lastrev file with new revision number
        if args.channel not in ["release", "source"]:
            with open(lastrev_file, 'w') as f:
                f.write(str(rev))
        
        return 0
    
    except Exception as err:
        sys.stderr.write("\n" + traceback.format_exc())
        return 1
    
    # Clean up
    finally:
        if os.path.exists(tmp_dir):
            shutil.rmtree(tmp_dir)


def dump_file(f):
    with open(f, 'r') as f:
        log(f.read())


def log(msg):
    print(msg, file=sys.stdout)


def log_line():
    log('======================================================\n\n')

def sha1file(f):
    sha1 = hashlib.sha1()
    with open(f, 'rb') as f:
        sha1.update(f.read())
    return sha1.hexdigest()


if __name__ == '__main__':
    sys.exit(main())
