#!/usr/bin/env python3
# Copyright lowRISC contributors.
# Licensed under the Apache License, Version 2.0, see LICENSE for details.
# SPDX-License-Identifier: Apache-2.0

import argparse
import collections
import datetime
import io
import logging
import os.path
import re
import subprocess
import sys
import tarfile
import time
import urllib.request
import xml.etree.ElementTree
from pathlib import Path
from typing import Dict, Set

# Default location of the bitstreams cache.
CACHE_DIR = '~/.cache/opentitan-bitstreams'
# Default bucket URL.
BUCKET_URL = 'https://storage.googleapis.com/opentitan-bitstreams/'
# The xml document returned by the bucket is in this namespace.
XMLNS = {'': 'http://doc.s3.amazonaws.com/2006-03-01'}

parser = argparse.ArgumentParser(
    description='Bitstream Downloader & Cache manager')
parser.add_argument('--cache', default=CACHE_DIR, help='Cache directory name')
parser.add_argument('--latest-update',
                    default='latest.txt',
                    help='Last time the cache was updated')
parser.add_argument('--bucket-url', default=BUCKET_URL, help='GCP Bucket URL')
parser.add_argument('--build-file',
                    default='BUILD.bazel',
                    help='Name of the genrated BUILD file')
parser.add_argument('--list',
                    default=False,
                    action=argparse.BooleanOptionalAction,
                    help='List GCP Bucket contents')
parser.add_argument('--offline',
                    default=False,
                    action=argparse.BooleanOptionalAction,
                    help='Operating in an offline environment')
parser.add_argument('--refresh',
                    default=False,
                    action=argparse.BooleanOptionalAction,
                    help='Force a refresh')
parser.add_argument('--refresh-time',
                    default=300,
                    type=int,
                    help='How often to check for new bitstreams')
parser.add_argument('--repo',
                    default='',
                    help="Location of the source git repo")
parser.add_argument(
    'bitstream',
    default='HEAD',
    nargs='?',
    help='Bitstream to retrieve: "latest" or git commit identifier')


class BitstreamCache(object):

    def __init__(self, bucket_url, cachedir, latest_update, offline=False):
        """Initialize the Bitstream Cache Manager."""
        if bucket_url[-1] != '/':
            bucket_url += '/'
        self.bucket_url = bucket_url
        cachedir = os.path.expanduser(cachedir)
        self.cachedir = os.path.join(cachedir, 'cache')
        latest_update = os.path.join(cachedir,
                                     os.path.expanduser(latest_update))
        self.latest_update = latest_update
        self.offline = offline
        self.available = {}

    def InitRepository(self):
        """Create the cache directory and symlink it into the bazel repository dir."""
        os.makedirs(self.cachedir, exist_ok=True)
        os.symlink(self.cachedir, 'cache')

    def Touch(self, key):
        """Set the latest known bitstream.

        Args:
            key: str; The git hash of the latest bitstream.
        """
        with open(self.latest_update, 'wt') as f:
            print(key, file=f)

    def NeedRefresh(self, interval):
        """Determine if the cache needs a refresh.

        Args:
            interval: int; Desired interval between refresh.
        Returns:
            bool: whether a refresh is needed.
        """
        try:
            st = os.stat(self.latest_update)
            return time.time() - st.st_mtime > interval
        except FileNotFoundError:
            return True

    def Get(self, file):
        """Perform an HTTP GET from the GCP bitstream bucket.

        Args:
            file: Filename in the bucket to retrieve.
        Returns:
            bytes
        """
        response = urllib.request.urlopen(self.bucket_url + file)
        return response.read()

    def GetBitstreamsAvailable(self, refresh):
        """Inventory which bitstreams are available.

        Args:
            refresh: bool; whether to refresh from the network.
        """
        if not refresh:
            for (_, dirnames, _) in os.walk('cache'):
                for d in dirnames:
                    self.available[d] = 'local'
            try:
                with open(self.latest_update, 'rt') as f:
                    self.available['latest'] = f.read().strip()
            except FileNotFoundError:
                if self.offline:
                    logging.error(
                        'Must pre-initialize bitstream cache in offline mode.')
                else:
                    logging.error(
                        f'Bitstream cache missing {self.latest_update}.')
                sys.exit(1)
            return
        document = self.Get('').decode('utf-8')
        et = xml.etree.ElementTree.fromstring(document)
        for content in et.findall('Contents', XMLNS):
            for key in content.findall('Key', XMLNS):
                m = re.search(r'bitstream-([0-9A-Fa-f]+).tar.gz', key.text)
                if m:
                    self.available[m.group(1)] = key.text
        latest = self.Get('master/latest.txt').decode('utf-8').split('\n')
        self.available['latest'] = latest[1]

    def GetClosest(self, repodir, key):
        """Get the best match for a bitstream (exact or next older commit).

        Args:
            repodir: path; Path to the repo from which bitstreams are built.
            key: str; A git hash or identifier of the desired bitstream.
        Returns:
            str or None: git hash of the closest bitstream.
        """
        if key in self.available:
            return key
        commits = []
        lines = subprocess.check_output(
            ['git', 'log', '--oneline', '--no-abbrev-commit', key],
            universal_newlines=True,
            cwd=repodir)
        for line in lines.split('\n'):
            commits.append(line.split(' ')[0])

        for commit in commits:
            if commit in self.available:
                return commit
        return None

    def _GetFromLocal(self, key):
        """Get the bitstream files from the local filesystem.

        Args:
            key: str; A git hash or the string 'latest'.
        Returns:
            list[str]: The requested bitstream files or empty list.
        """
        if key == 'latest':
            key = self.available['latest']
        files = []
        local_dir = os.path.join('cache', key)
        for (dirname, _, filenames) in os.walk(local_dir):
            files.extend(os.path.join(dirname, f) for f in filenames)
        return files

    def _GetFromRemote(self, key):
        """Get the bitstream files from GCP bucket.
        The retrieved files are extracted into the cache directory.

        Args:
            key: str; A git hash or the string 'latest'.
        """
        if self.offline:
            return
        if key == 'latest':
            latest = self.available['latest']
            key = latest
        else:
            latest = None

        remote_filename = self.available[key]
        local_dir = os.path.join(self.cachedir, key)
        archive = io.BytesIO(self.Get(remote_filename))
        tar = tarfile.open(fileobj=archive, mode='r|*')
        tar.extractall(local_dir)
        if latest:
            self.Touch(latest)

    def GetFromCache(self, key: str) -> Dict[str, Set[str]]:
        """Get the requested bitstream files.

        Args:
            key: A git hash or the string 'latest'.
        Returns:
            A dictionary that maps file extensions to the set of bitstream files
            with that extension.
        """
        files = self._GetFromLocal(key)
        if not files:
            self._GetFromRemote(key)
            files = self._GetFromLocal(key)

        extension_map: Dict[str, Set[str]] = collections.defaultdict(set)
        for file_name in files:
            _, extension = os.path.splitext(file_name)
            if extension[0] != '.':
                logging.error("Bitstream file has no extension: " +
                              repr(file_name))
                sys.exit(1)
            key = extension[1:]
            extension_map[key] |= {file_name}
        return extension_map

    @staticmethod
    def _GetDateTimeStr():
        return datetime.datetime.now().isoformat()

    def _ConstructBazelString(self, build_file: Path, key: str) -> str:
        files_by_extension = self.GetFromCache(key)

        if len(files_by_extension.get('orig', set())) != 1 or \
           len(files_by_extension.get('splice', set())) != 1:
            error_msg_lines = [
                "Could not find the bitstreams to generate a BUILD file:" +
                repr(build_file),
                "in files_by_extension:" + repr(files_by_extension),
                "using key:" + repr(key),
            ]
            logging.error('\n'.join(error_msg_lines))
            sys.exit(1)

        (test_rom_file, ) = files_by_extension['orig']
        (rom_file, ) = files_by_extension['splice']

        def filegroup_lines(name, src):
            return [
                'filegroup(',
                '    name = "{}",'.format(name),
                '    srcs = ["{}"],'.format(src),
                ')',
                '',
            ]

        bazel_lines = [
            '# This file was autogenerated. Do not edit!',
            '# Built at {}.'.format(BitstreamCache._GetDateTimeStr()),
            '# Configured for bitstream: {}'.format(key),
            '',
            'package(default_visibility = ["//visibility:public"])',
            '',
            'exports_files(glob(["cache/**"]))',
            '',
        ] + filegroup_lines('bitstream_test_rom', test_rom_file) \
          + filegroup_lines('bitstream_rom', rom_file)

        used_target_names: Set[str] = set()

        for mmi_file in sorted(files_by_extension.get('mmi', set())):
            target_name = os.path.basename(mmi_file).replace('.', '_')

            # Only use the top-level files for now.
            target_directory = "cache/{}".format(key)
            if (os.path.dirname(mmi_file) != target_directory):
              continue

            if target_name in used_target_names:
                logging.error(
                    "Target name {} for file {} would collide with another target"
                    .format(repr(target_name), repr(mmi_file)))
                sys.exit(1)
            used_target_names.add(target_name)

            bazel_lines += filegroup_lines(target_name, mmi_file)

        return '\n'.join(bazel_lines)

    def WriteBuildFile(self, build_file: Path, key: str) -> str:
        """Write a BUILD file for the requested bitstream files.

        Args:
            build: path; Filename of the BUILD file to write.
            key: str; A git hash or the string 'latest'.
        Returns:
            Either `key` or the corresponding commit hash if `key` is 'latest'.
        """
        with open(build_file, 'wt') as f:
            f.write(self._ConstructBazelString(build_file, key))

        if key != 'latest':
            return key
        return self.available[key]


def main(argv):
    # The user can override some basic behaviors with the BITSTREAM
    # environment variable.
    env = os.environ.get('BITSTREAM')
    if env:
        argv.extend(env.split(' '))
    args = parser.parse_args(args=argv[1:])
    desired_bitstream = args.bitstream

    # We need to know the location of the main git repo, since this script
    # will have its CWD in a bazel-managed 'external' directory.
    # If not provided, we assume this script itself is located in the main
    # git repository.
    if args.repo:
        if os.path.isdir(args.repo):
            repo = args.repo
        else:
            repo = os.path.dirname(args.repo)
    else:
        repo = os.path.dirname(argv[0])

    cache = BitstreamCache(args.bucket_url, args.cache, args.latest_update,
                           args.offline)
    cache.InitRepository()

    # Do we need a refresh?
    need_refresh = (args.refresh or desired_bitstream != 'latest' or
                    cache.NeedRefresh(args.refresh_time))
    cache.GetBitstreamsAvailable(need_refresh and not args.offline)

    # If commanded to print bitstream availability, do so.
    if args.list:
        for k, v in cache.available.items():
            print('{}: {}'.format(k, v))

    # If we aren't getting the latest bitstream, resolve the hash to the closest
    # bitstream we can find.
    if desired_bitstream != 'latest':
        closest_bitstream = cache.GetClosest(repo, desired_bitstream)
        if closest_bitstream is None:
            logging.error('Cannot find a bitstream close to {}'.format(
                desired_bitstream))
            return 1
        if closest_bitstream != desired_bitstream:
            logging.warning('Closest bitstream to {} is {}.'.format(
                desired_bitstream, closest_bitstream))
            desired_bitstream = closest_bitstream

    # Write a build file which allows tests to reference the bitstreams with
    # the labels:
    #   @bitstreams//:bitstream_test_rom
    #   @bitstreams//:bitstream_rom
    configured_bitream = cache.WriteBuildFile(args.build_file,
                                              desired_bitstream)
    if desired_bitstream != 'latest' and configured_bitream != desired_bitstream:
        logging.error(
            'Configured bitstream {} does not match desired bitstream {}.'.
            format(configured_bitream, desired_bitstream))
        return 1
    logging.info(
        'Configured latest bitstream as {}.'.format(configured_bitream))

    return 0


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