#!/usr/bin/python3
#
# Copyright (c) 2010-2012 Steffen Kieß
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#

import argparse
import json
import os
import pathlib
import re
import shlex
import subprocess
import sys
import time
import urllib.error
import urllib.parse
import urllib.request

from git_write_wd_tree import write_wd_tree


BRANCH = 'refs/heads/test'
REMOTE = 'origin'
CI_FILE = '.gitlab-ci.yml'
CI_KEY = re.compile(r'^([A-Za-z0-9.][A-Za-z0-9_.-]*):[ \t]*(?:#.*)?$')
CI_OPTION_MARKER = re.compile(
    r'^([ ]+)# test-gitlab-ci: job-option-name=([A-Za-z0-9_.-]+)$')
CI_SCRIPT_KEY = re.compile(r'^[ ]*script:[ \t]*(?:#.*)?$')
CI_SCRIPT_ITEM = re.compile(
    r'^([ ]*-[ ]+)(\'(?:[^\']|\'\')*\'|"(?:[^"\\]|\\.)*")([ \t]*)(\r?\n?)$')
CI_GLOBAL_KEYS = {'stages', 'default', 'variables', 'workflow'}
FINAL_STATUSES = {'success', 'failed', 'canceled', 'skipped', 'manual'}


def runget(command, *, cwd, input=None):
    result = subprocess.run(command, check=True, cwd=cwd, input=input,
                            stdout=subprocess.PIPE)
    return result.stdout.decode('utf-8').rstrip('\n')


def pipeline_url(remote_url, commit):
    host, separator, path = remote_url.partition(':')
    if separator and '@' in host:
        remote_url = 'https://{}/{}'.format(host.split('@', 1)[1], path)
    if remote_url.endswith('.git'):
        remote_url = remote_url[:-4]
    return '{}/commit/{}/pipelines?ref={}'.format(remote_url, commit, BRANCH)


def gitlab_project(remote_url):
    if '://' in remote_url:
        parsed = urllib.parse.urlsplit(remote_url)
        if (parsed.scheme not in ('http', 'https', 'ssh') or
                not parsed.hostname):
            raise ValueError('Unsupported origin URL: ' + remote_url)
        host = parsed.netloc if parsed.scheme != 'ssh' else parsed.hostname
        if '@' in host:
            host = host.rsplit('@', 1)[1]
        base = '{}://{}'.format(
            parsed.scheme if parsed.scheme != 'ssh' else 'https', host)
        path = parsed.path.strip('/')
    else:
        match = re.fullmatch(r'(?:[^@]+@)?([^:]+):/?(.+)', remote_url)
        if match is None:
            raise ValueError('Unsupported origin URL: ' + remote_url)
        base, path = 'https://' + match.group(1), match.group(2)
    if path.endswith('.git'):
        path = path[:-4]
    if not path:
        raise ValueError('Origin URL has no GitLab project path')
    return base, path


class GitLabAPI:
    def __init__(self, remote_url):
        base, path = gitlab_project(remote_url)
        self.project_url = '{}/{}'.format(base, path)
        self.url = '{}/api/v4/projects/{}'.format(
            base, urllib.parse.quote(path, safe=''))
        hostname = urllib.parse.urlsplit(base).hostname
        self.token_path = (pathlib.Path.home() / '.config' / 'gitlab' /
                           hostname / 'token')
        token = os.environ.get('GITLAB_TOKEN')
        if not token:
            try:
                token = self.token_path.read_text(encoding='utf-8').strip()
            except FileNotFoundError:
                token = None
            except (OSError, UnicodeError) as error:
                raise RuntimeError(
                    'Could not read GitLab token file {}: {}'.format(
                        self.token_path, error)) from error
            if token == '':
                raise RuntimeError('GitLab token file {} is empty'.format(
                    self.token_path))
        self.headers = {'PRIVATE-TOKEN': token} if token else {}

    def get(self, path, query=None, allow_missing=False):
        url = self.url + path
        if query:
            url += '?' + urllib.parse.urlencode(query)
        request = urllib.request.Request(url, headers=self.headers)
        try:
            with urllib.request.urlopen(request, timeout=30) as response:
                return response.read(), response.headers
        except urllib.error.HTTPError as error:
            if allow_missing and error.code == 404:
                return None, error.headers
            if error.code in (401, 403, 404):
                raise RuntimeError(
                    'GitLab API access denied (HTTP {}) for {}. Set '
                    'GITLAB_TOKEN or put an access token with read_api or '
                    'api scope and project access in {}.'.format(
                        error.code, url, self.token_path)) from error
            raise RuntimeError('GitLab API returned HTTP {} for {}'.format(
                error.code, url)) from error
        except urllib.error.URLError as error:
            raise RuntimeError('Could not reach GitLab API at {}: {}'.format(
                url, error.reason)) from error

    def get_json(self, path, query=None):
        data, headers = self.get(path, query)
        try:
            return json.loads(data), headers
        except (ValueError, UnicodeDecodeError) as error:
            raise RuntimeError(
                'Invalid GitLab API response for ' + path) from error

    def pipelines(self, commit):
        result, _ = self.get_json('/pipelines', {
            'sha': commit, 'ref': 'test', 'source': 'push',
            'order_by': 'id', 'sort': 'desc', 'per_page': 1})
        return result

    def jobs(self, pipeline_id):
        jobs = []
        page = 1
        while True:
            batch, headers = self.get_json('/pipelines/{}/jobs'.format(
                pipeline_id), {'per_page': 100, 'page': page,
                               'include_retried': 'true'})
            jobs.extend(batch)
            next_page = headers.get('X-Next-Page')
            if next_page:
                page = int(next_page)
            elif len(batch) == 100:
                page += 1
            else:
                return jobs


def wait_for_pipeline(api, commit, previous_id):
    deadline = time.monotonic() + 120
    while True:
        pipelines = api.pipelines(commit)
        pipeline = next((item for item in pipelines
                         if item['id'] > previous_id and
                         item['sha'] == commit and item['ref'] == 'test' and
                         item['source'] == 'push'), None)
        if pipeline is not None:
            break
        if time.monotonic() >= deadline:
            raise RuntimeError('No new push pipeline appeared for {} on test '
                               'within two minutes'.format(commit))
        time.sleep(min(2, max(0, deadline - time.monotonic())))

    print('Pipeline: {}'.format(
        pipeline.get('web_url') or '{}/-/pipelines/{}'.format(
            api.project_url, pipeline['id'])), flush=True)
    while True:
        pipeline, _ = api.get_json('/pipelines/{}'.format(pipeline['id']))
        if pipeline['status'] in FINAL_STATUSES:
            return pipeline
        time.sleep(2)


def print_job(api, job):
    url = job.get('web_url') or '{}/-/jobs/{}'.format(api.project_url,
                                                      job['id'])
    print('Job {}: {} ({})'.format(job['name'], job['status'], url))
    log, _ = api.get('/jobs/{}/trace'.format(job['id']), allow_missing=True)
    if log is None or not log:
        print('(No job log available)')
    else:
        decoded = log.decode('utf-8', errors='replace')
        print(decoded, end='' if decoded.endswith('\n') else '\n')


def print_result(api, pipeline, job_names):
    jobs = api.jobs(pipeline['id'])
    if job_names:
        selected = []
        for name in job_names:
            attempts = [job for job in jobs if job['name'] == name]
            if not attempts:
                raise RuntimeError(
                    'Job {} was not found in pipeline {}'.format(
                        name, pipeline['id']))
            selected.append(max(attempts, key=lambda job: job['id']))
        jobs = selected
    for job in jobs:
        print_job(api, job)
    print('Pipeline result: {} ({})'.format(
        pipeline['status'], pipeline.get('web_url') or
        '{}/-/pipelines/{}'.format(api.project_url, pipeline['id'])))
    if job_names:
        return 0 if all(job['status'] == 'success' for job in jobs) else 1
    return 0 if pipeline['status'] == 'success' else 1


def select_ci_jobs(config, job_names):
    lines = config.decode('utf-8').splitlines(keepends=True)
    headings = []
    for index, line in enumerate(lines):
        stripped = line.rstrip('\r\n')
        if (not stripped or stripped.lstrip().startswith('#') or
                stripped[0].isspace()):
            continue
        match = CI_KEY.fullmatch(stripped)
        if match is None:
            raise ValueError('Unsupported top-level CI configuration: ' +
                             stripped)
        headings.append((index, match.group(1)))

    sections = {}
    for number, (start, name) in enumerate(headings):
        if name in sections:
            raise ValueError('Duplicate CI configuration key: ' + name)
        end = (headings[number + 1][0] if number + 1 < len(headings)
               else len(lines))
        sections[name] = ''.join(lines[start:end])

    available = sorted(name for name in sections
                       if name not in CI_GLOBAL_KEYS and
                       not name.startswith('.'))
    unknown = [name for name in job_names if name not in available]
    if unknown:
        raise ValueError('Unknown job(s): {} (available: {})'.format(
            ', '.join(unknown), ', '.join(available)))
    if 'include' in sections:
        raise ValueError('Cannot select jobs from a CI configuration '
                         'with includes')
    if 'stages' not in sections:
        raise ValueError('Missing stages in CI configuration')

    keep = CI_GLOBAL_KEYS | {name for name in sections if name.startswith('.')}
    keep.update(job_names)
    preamble = ''.join(lines[:headings[0][0]])
    filtered = preamble + ''.join(sections[name] for _, name in headings
                                  if name in keep)
    return filtered.encode('utf-8')


def add_job_options(config, job_options):
    lines = config.decode('utf-8').splitlines(keepends=True)
    option_names = {name for name, _ in job_options}
    options = {name: [] for name in option_names}
    for name, value in job_options:
        options[name].append(value)
    matched = set()
    script_indent = None
    pending = None

    for index, line in enumerate(lines):
        content = line.rstrip('\r\n')
        stripped = content.strip()
        indent = len(content) - len(content.lstrip(' '))
        if pending is not None:
            name, marker_indent = pending
            match = CI_SCRIPT_ITEM.fullmatch(line)
            if match is None or indent != marker_indent:
                raise ValueError('Expected a quoted script command after '
                                 'job-option-name=' + name)
            extra = ' ' + ' '.join(shlex.quote(value) for value in options[name])
            scalar = match.group(2)
            if scalar.startswith("'"):
                extra = extra.replace("'", "''")
            else:
                extra = json.dumps(extra)[1:-1]
            lines[index] = (match.group(1) + scalar[:-1] + extra +
                            scalar[-1] + match.group(3) + match.group(4))
            matched.add(name)
            pending = None
            continue
        if stripped and not stripped.startswith('#'):
            if script_indent is not None and indent <= script_indent:
                script_indent = None
            if CI_SCRIPT_KEY.fullmatch(content):
                script_indent = indent
        marker = CI_OPTION_MARKER.fullmatch(content)
        if marker is not None and marker.group(2) in option_names:
            if script_indent is None or indent <= script_indent:
                raise ValueError('job-option-name={} is outside a script list'
                                 .format(marker.group(2)))
            pending = (marker.group(2), indent)

    if pending is not None:
        raise ValueError('Expected a quoted script command after '
                         'job-option-name=' + pending[0])
    missing = option_names - matched
    if missing:
        raise ValueError('No marked command for job option name(s): ' +
                         ', '.join(sorted(missing)))
    return ''.join(lines).encode('utf-8')


def main(argv):
    parser = argparse.ArgumentParser(description='Start a GitLab CI test run')
    parser.add_argument('--job', metavar='NAME', action='append', default=[],
                        help='Run a named GitLab CI job (repeatable)')
    parser.add_argument('--job-option', metavar='NAME:ARG', action='append',
                        default=[], help='Append ARG to commands marked NAME '
                        '(repeatable)')
    parser.add_argument('--wait-for-result', action='store_true',
                        help='Wait for GitLab CI and print job logs')
    args = parser.parse_args(argv)

    tools_dir = pathlib.Path(__file__).resolve().parent
    repo_root = tools_dir.parent
    job_names = list(dict.fromkeys(args.job))
    job_options = []
    for entry in args.job_option:
        name, separator, value = entry.partition(':')
        if (not separator or not re.fullmatch(r'[A-Za-z0-9_.-]+', name)
                or not value):
            parser.error('--job-option must be NAME:ARG with a nonempty ARG')
        if any(character in value for character in '\r\n\0'):
            parser.error('--job-option ARG must be a single line')
        job_options.append((name, value))
    file_overrides = None
    if job_names or job_options:
        try:
            config = (repo_root / CI_FILE).read_bytes()
            if job_names:
                config = select_ci_jobs(config, job_names)
            if job_options:
                config = add_job_options(config, job_options)
        except ValueError as error:
            parser.error(str(error))
        file_overrides = {CI_FILE: config}

    try:
        tree = write_wd_tree(file_overrides=file_overrides)
        commit = runget(['git', 'commit-tree', tree, '-p', 'HEAD'],
                        cwd=repo_root, input=b'Gitlab CI test run\n')
        remote_url = None
        if args.wait_for_result:
            remote_url = runget(['git', 'remote', 'get-url', REMOTE],
                                cwd=repo_root)
            api = GitLabAPI(remote_url)
            existing = api.pipelines(commit)
            previous_id = existing[0]['id'] if existing else 0
        subprocess.run(
            ['git', 'push', REMOTE, '+{}:{}'.format(commit, BRANCH)],
            check=True, cwd=repo_root)
        if remote_url is None:
            remote_url = runget(['git', 'remote', 'get-url', REMOTE],
                                cwd=repo_root)
    except subprocess.CalledProcessError as error:
        return error.returncode
    except (RuntimeError, ValueError) as error:
        print('Error: {}'.format(error), file=sys.stderr)
        return 1

    print('See status on: {}'.format(pipeline_url(remote_url, commit)))
    if args.wait_for_result:
        try:
            pipeline = wait_for_pipeline(api, commit, previous_id)
            return print_result(api, pipeline, job_names)
        except (RuntimeError, ValueError) as error:
            print('Error: {}'.format(error), file=sys.stderr)
            return 1
    return 0


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