#!/usr/bin/env python3
"""Upload one reviewed file with a private, reusable receipt. Python 3 stdlib only."""
import argparse
import hashlib
import json
import os
from pathlib import Path
import time
import urllib.error
import urllib.parse
import urllib.request


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None


OPENER = urllib.request.build_opener(NoRedirect)


def request(url, method='GET', data=None, headers=None, timeout=30):
    parsed = urllib.parse.urlsplit(url)
    if parsed.scheme != 'https' and not (parsed.scheme == 'http' and parsed.hostname in ('localhost', '127.0.0.1')):
        raise RuntimeError('HTTPS required')
    req = urllib.request.Request(url, data=data, headers=headers or {}, method=method)
    try:
        with OPENER.open(req, timeout=timeout) as response:
            return response.status, response.headers, response.read()
    except urllib.error.HTTPError as error:
        if error.code == 308:
            return error.code, error.headers, b''
        # Do not print URLs, capability tokens or server-provided private values.
        raise RuntimeError(f'HTTP {error.code}; inspect saved state before retrying') from None


def digest(path):
    with path.open('rb') as source:
        value = hashlib.sha256()
        for chunk in iter(lambda: source.read(1024 * 1024), b''):
            value.update(chunk)
        return value.hexdigest()


def save_receipt(path, value):
    fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
    with os.fdopen(fd, 'w') as target:
        json.dump(value, target)
        target.flush()
        os.fsync(target.fileno())


def transfer(upload_url, source, size, policy):
    code, headers, _ = request(upload_url, 'PUT', b'', {'Content-Range': f'bytes */{size}'}, policy['chunkTimeoutMs'] / 1000)
    if code in (200, 201):
        return
    if code != 308:
        raise RuntimeError('Unexpected resumable status')
    stored = headers.get('Range', '')
    offset = int(stored.rsplit('-', 1)[1]) + 1 if stored.startswith('bytes=0-') else 0
    if not 0 <= offset <= size:
        raise RuntimeError('Invalid resumable offset')
    with source.open('rb') as stream:
        stream.seek(offset)
        while offset < size:
            chunk = stream.read(min(policy['chunkBytes'], size - offset))
            if not chunk:
                raise RuntimeError('File changed during upload')
            end = offset + len(chunk)
            code, _, _ = request(upload_url, 'PUT', chunk, {
                'Content-Type': 'application/octet-stream',
                'Content-Range': f'bytes {offset}-{end - 1}/{size}',
            }, policy['chunkTimeoutMs'] / 1000)
            if code not in (200, 201, 308):
                raise RuntimeError('Unexpected transfer status')
            offset = end


def upload(args):
    origin = os.environ['CODEMERCHANT_URL'].rstrip('/')
    parsed = urllib.parse.urlsplit(origin)
    if parsed.path or parsed.query or parsed.fragment or parsed.username or parsed.password:
        raise RuntimeError('CODEMERCHANT_URL must be a site origin')
    token = os.environ['CODEMERCHANT_API_TOKEN']
    _, _, manifest_bytes = request(origin + '/agent/workflow.json')
    policy = json.loads(manifest_bytes)['limits']
    source = Path(args.file)
    identity = {'origin': origin, 'repository': args.repository, 'size': source.stat().st_size, 'sha256': digest(source)}
    if not 0 < identity['size'] <= policy['maxFileBytes']:
        raise RuntimeError('File is empty or exceeds the service limit')

    def action(action_name, **values):
        body = {'action': action_name, 'repository': args.repository, **values}
        code, _, result = request(origin + '/api/v1/context-uploads', 'POST', json.dumps(body).encode(),
                                 {'Authorization': 'Bearer ' + token, 'Content-Type': 'application/json'}, policy['requestTimeoutMs'] / 1000)
        if not 200 <= code < 300:
            raise RuntimeError('Control request did not complete')
        return json.loads(result)

    receipt_path = Path(args.receipt)
    if receipt_path.exists():
        saved = json.loads(receipt_path.read_text())
        if any(saved.get(key) != value for key, value in identity.items()):
            raise RuntimeError('Receipt does not match this origin, repository or file')
        os.chmod(receipt_path, 0o600)
    else:
        saved = {**identity, **action('initiate', filename=source.name, size=identity['size'])}
        save_receipt(receipt_path, saved)
    current = action('status', id=saved['id'])
    if current['status'] == 'uploading':
        transfer(saved['uploadUrl'], source, identity['size'], policy)
        if source.stat().st_size != identity['size'] or digest(source) != identity['sha256']:
            raise RuntimeError('File changed; do not complete this upload')
        current = action('complete', id=saved['id'])
    deadline = time.monotonic() + policy['maxExtractionSeconds']
    while current['status'] in ('uploading', 'queued', 'processing'):
        if time.monotonic() >= deadline:
            raise RuntimeError('Extraction pending; rerun with this receipt to resume')
        time.sleep(policy['pollMs'] / 1000)
        current = action('status', id=saved['id'])
    if current['status'] != 'complete':
        raise RuntimeError('Extraction failed; retain receipt for diagnosis')
    print(json.dumps({key: current[key] for key in ('id', 'filename', 'size', 'status')}))


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('repository')
    parser.add_argument('file')
    parser.add_argument('--receipt', required=True)
    args = parser.parse_args()
    try:
        upload(args)
    except (RuntimeError, OSError, KeyError, ValueError) as error:
        # Exception text from network libraries may contain a capability URL.
        message = str(error) if isinstance(error, RuntimeError) else type(error).__name__
        parser.exit(1, f'Upload incomplete: {message}. Keep the receipt private.\n')


if __name__ == '__main__':
    main()