# Copyright 2017 The Kubernetes Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Generates a SQLite DB containing test data downloaded from GCS."""


import argparse
import logging
import json
import os
import random
import re
import signal
import sys
import time
import urllib.parse
from xml.etree import cElementTree as ET

import multiprocessing
import multiprocessing.pool
import requests
import ruamel.yaml as yaml

import model


def pad_numbers(string):
    """Modify a string to make its numbers suitable for natural sorting."""
    return re.sub(r'\d+', lambda m: m.group(0).rjust(16, '0'), string)

WORKER_CLIENT = None  # used for multiprocessing


class GCSClient:
    def __init__(self, jobs_dir, metadata=None):
        self.jobs_dir = jobs_dir
        self.metadata = metadata or {}
        self.session = requests.Session()

    def _request(self, path, params, as_json=True):
        """GETs a JSON resource from GCS, with retries on failure.

        Retries are based on guidance from
        cloud.google.com/storage/docs/gsutil/addlhelp/RetryHandlingStrategy

        """
        url = f'https://www.googleapis.com/storage/v1/b/{path}'
        for retry in range(23):
            try:
                resp = self.session.get(url, params=params, stream=False)
                if 400 <= resp.status_code < 500 and resp.status_code != 429:
                    return None
                resp.raise_for_status()
                if as_json:
                    try:
                        return resp.json()
                    except json.decoder.JSONDecodeError:
                        logging.exception('Failed to decode request for %s', path)
                        return None
                return resp.text
            except requests.exceptions.RequestException:
                logging.exception('request failed %s', url)
            time.sleep(random.random() * min(60, 2 ** retry))

    @staticmethod
    def _parse_uri(path):
        if not path.startswith('gs://'):
            raise ValueError("Bad path: %s, should start with 'gs://'" % path)
        bucket, prefix = path[5:].split('/', 1)
        return bucket, prefix

    def get(self, path, as_json=False):
        """Get an object from GCS."""
        bucket, path = self._parse_uri(path)
        return self._request(f'{bucket}/o/{urllib.parse.quote(path, "")}',
                             {'alt': 'media'}, as_json=as_json)

    def ls(self,
           path,
           dirs=True,
           files=True,
           delim=True,
           item_field='name',
           build_limit=sys.maxsize,):
        """Lists objects under a path on gcs."""
        # pylint: disable=invalid-name

        bucket, path = self._parse_uri(path)
        params = {'prefix': path, 'fields': 'nextPageToken'}
        if delim:
            params['delimiter'] = '/'
            if dirs:
                params['fields'] += ',prefixes'
        if files:
            params['fields'] += f',items({item_field})'
        while build_limit > 0:
            resp = self._request(f'{bucket}/o', params)
            if resp is None:  # nothing under path?
                return
            for prefix in resp.get('prefixes', []):
                build_limit -= 1
                yield f'gs://{bucket}/{prefix}'
            for item in resp.get('items', []):
                if item_field == 'name':
                    build_limit -= 1
                    yield f'gs://{bucket}/{item["name"]}'
                else:
                    build_limit -= 1
                    yield item[item_field]
            if 'nextPageToken' not in resp:
                break
            params['pageToken'] = resp['nextPageToken']

    def ls_dirs(self, path):
        return self.ls(path, dirs=True, files=False)

    def _ls_junit_paths(self, build_dir):
        """Lists the paths of JUnit XML files for a build."""
        url = f'{build_dir}artifacts/'
        for path in self.ls(url):
            if re.match(r'.*/junit.*\.xml$', path):
                yield path

    def get_junits_from_build(self, build_dir):
        """Generates all tests for a build."""
        files = {}
        assert not build_dir.endswith('/')
        for junit_path in self._ls_junit_paths(build_dir + '/'):
            files[junit_path] = self.get(junit_path)
        return files

    def _get_jobs(self):
        """Generates all jobs in the bucket."""
        for job_path in self.ls_dirs(self.jobs_dir):
            yield os.path.basename(os.path.dirname(job_path))

    def _get_builds(self, job, build_limit=sys.maxsize):
        '''Returns whether builds are precise (guarantees existence)'''
        if self.metadata.get('sequential', True):
            try:
                latest_build = int(self.get(f'{self.jobs_dir}{job}/latest-build.txt'))
            except (ValueError, TypeError):
                pass
            else:
                return False, (str(n) for n in range(latest_build, 0, -1)[:build_limit])
        # Invalid latest-build or bucket is using timestamps
        build_paths = self.ls_dirs(f'{self.jobs_dir}{job}/')
        return True, sorted(
            (os.path.basename(os.path.dirname(b)) for b in build_paths),
            key=pad_numbers, reverse=True)[:build_limit]

    def get_started_finished(self, job, build):
        if self.metadata.get('pr'):
            build_dir = self.get(f'{self.jobs_dir}/directory/{job}/{build}.txt').strip()
        else:
            build_dir = f'{self.jobs_dir}{job}/{build}'
        started = self.get(f'{build_dir}/started.json', as_json=True)
        finished = self.get(f'{build_dir}/finished.json', as_json=True)
        return build_dir, started, finished

    def get_builds(self, builds_have, build_limit=sys.maxsize):
        """Generates all (job, build) pairs ever."""
        if self.metadata.get('pr'):
            files = self.ls(self.jobs_dir + '/directory/', delim=False, build_limit=build_limit)
            for fname in files:
                if fname.endswith('.txt') and 'latest-build' not in fname:
                    job, build = fname[:-4].split('/')[-2:]
                    if (job, build) in builds_have:
                        continue
                    yield job, build
            return
        for job in self._get_jobs():
            if job in self.metadata.get('exclude_jobs', []):
                continue
            have = 0
            precise, builds = self._get_builds(job, build_limit)
            for build in builds:
                if (job, build) in builds_have:
                    have += 1
                    if have > 40 and not precise:
                        break
                    continue
                yield job, build


def mp_init_worker(jobs_dir, metadata, client_class, use_signal=True):
    """
    Initialize the environment for multiprocessing-based multithreading.
    """

    if use_signal:
        signal.signal(signal.SIGINT, signal.SIG_IGN)
    # Multiprocessing doesn't allow local variables for each worker, so we need
    # to make a GCSClient global variable.
    global WORKER_CLIENT  # pylint: disable=global-statement
    WORKER_CLIENT = client_class(jobs_dir, metadata)

def get_started_finished(job_info):
    (job, build) = job_info
    try:
        return WORKER_CLIENT.get_started_finished(job, build)
    except: # pylint: disable=W0702
        logging.exception('failed to get tests for %s/%s', job, build)
        return None, None, None

def get_junits(build_info):
    (build_id, gcs_path) = build_info
    try:
        junits = WORKER_CLIENT.get_junits_from_build(gcs_path)
        return build_id, gcs_path, junits
    except:
        logging.exception('failed to get junits for %s', gcs_path)
        raise