"""Focused local exercises against the recorded candidate; no retraining."""
import argparse
import copy
import html
import json
import tempfile
import threading
import urllib.error
import urllib.request
from pathlib import Path
import pipeline as p
from server import server


def refusal(action):
    try:
        action()
    except ValueError as error:
        return str(error)
    raise AssertionError('expected a refusal')


def exercise(stage):
    data = (p.HERE / 'candidate.json').read_bytes()
    expected = p.decode((p.HERE / 'result.json').read_bytes())['gate']['candidate_sha256']
    if p.digest(data) != expected:
        raise ValueError('saved candidate changed')
    bundle = p.validate(p.decode(data))
    common = {'execution_mode': 'local exercise using a previously trained candidate',
              'candidate_sha256': expected}
    if stage == 'export':
        return {**common, 'contract': bundle['contract'], 'features': bundle['features'],
                'mean': bundle['mean'], 'std': bundle['std'],
                'weights': bundle['weights'], 'bias': bundle['bias'],
                'invalid_dimensions': refusal(lambda: p.validate({**bundle, 'weights': []}))}
    if stage == 'gate':
        bad = copy.deepcopy(bundle)
        bad['weights'] = [0.0] * len(bundle['features'])
        bad['bias'] = 100.0
        constant = p.encode(bad)
        with tempfile.TemporaryDirectory() as root:
            rejected = refusal(lambda: p.register(root, constant, p.digest(constant)))
            files_after = len(list(Path(root).iterdir()))
        return {**common, 'candidate': p.quality_gate(data, expected),
                'constant_class_candidate': p.quality_gate(constant, p.digest(constant)),
                'registration_refusal': rejected, 'files_after_failed_gate': files_after}
    if stage == 'registry':
        with tempfile.TemporaryDirectory() as root:
            p.register(root, data, expected)
            p.register(root, data, expected)
            loaded = p.load(root, expected)
            path = p.registered_file(root, expected)
            path.write_bytes(data + b' ')
            changed = refusal(lambda: p.load(root, expected))
            alias = refusal(lambda: p.load(root, 'latest'))
        return {**common, 'first_registration': 'accepted', 'identical_retry': 'accepted',
                'loaded_contract': loaded['contract'], 'changed_bytes': changed, 'alias': alias}
    if stage != 'serve':
        raise ValueError('unknown exercise')
    request = p.decode((p.HERE / 'result.json').read_bytes())['requests'][0]
    requests = [request, {**request, 'features': dict(reversed(list(request['features'].items())))}]
    offline = p.batch(bundle, requests, expected)
    service = server(bundle, expected)
    worker = threading.Thread(target=service.serve_forever, daemon=True)
    worker.start()
    try:
        def post(value):
            outgoing = urllib.request.Request(
                f'http://127.0.0.1:{service.server_port}/predict', data=p.encode(value),
                headers={'Content-Type': 'application/json'})
            with urllib.request.urlopen(outgoing, timeout=3) as response:
                return response.status, json.load(response)
        responses = [post(value)[1] for value in requests]
        malformed = copy.deepcopy(request)
        del malformed['features']['area']
        try:
            post(malformed)
        except urllib.error.HTTPError as error:
            status, message = error.code, json.load(error)
            error.close()
        else:
            raise AssertionError('missing feature should be refused')
    finally:
        service.shutdown()
        worker.join(timeout=3)
        service.server_close()
    if responses != offline:
        raise AssertionError('batch/HTTP mismatch')
    return {**common, 'request': request, 'response': responses[0],
            'reordered_fields_match': responses[0] == responses[1],
            'batch_http_match': responses == offline, 'missing_feature_status': status,
            'missing_feature_response': message}


def report(stage, result):
    # Render actual program data, without turning it into claims or recreated terminal text.
    rows = ''.join('<tr><th>' + html.escape(key.replace('_', ' ')) + '</th><td><pre>'
                   + html.escape(json.dumps(value, indent=2, allow_nan=False)) + '</pre></td></tr>'
                   for key, value in result.items())
    return ('<!doctype html><html lang="en"><meta charset="utf-8">'
            '<meta name="viewport" content="width=device-width"><title>Local ' + stage + ' result</title>'
            '<style>body{margin:0;padding:28px;font:18px system-ui;color:#192136;background:white}'
            'h1{font-size:28px}p{color:#596375;line-height:1.5}table{border-collapse:collapse;width:100%;table-layout:fixed}'
            'th,td{border:1px solid #dfe4ed;padding:16px;text-align:left;vertical-align:top}th{width:25%}'
            'pre{margin:0;white-space:pre-wrap;overflow-wrap:anywhere;font:16px/1.5 monospace}'
            '@media(max-width:600px){body{padding:16px}th,td{display:block;width:auto}th{background:#f5f7fb}pre{font-size:16px}}'
            '</style><h1>' + stage.capitalize() + ': actual local result</h1>'
            '<p>Previously trained Rice candidate. This exercise runs local Python code; it does not train a new model. '
            'The serving exercise makes actual loopback HTTP requests.</p><table>' + rows + '</table></html>')


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('stage', choices=['export', 'gate', 'registry', 'serve'])
    parser.add_argument('--record', action='store_true')
    args = parser.parse_args()
    result = exercise(args.stage)
    text = json.dumps(result, indent=2, allow_nan=False) + '\n'
    print(text, end='')
    if args.record:
        (p.HERE / (args.stage + '-result.json')).write_text(text)
        (p.HERE / (args.stage + '-output.txt')).write_text(text)
        (p.HERE / (args.stage + '-report.html')).write_text(report(args.stage, result))
        record = {'command': ['python', 'exercise.py', args.stage, '--record'], 'exit_status': 0,
                  'execution_mode': result['execution_mode'],
                  'source_sha256': p.digest(Path(__file__).read_bytes()),
                  'candidate_sha256': result['candidate_sha256'],
                  'result_sha256': p.digest(text.encode())}
        (p.HERE / (args.stage + '-record.json')).write_text(json.dumps(record, indent=2) + '\n')