"""Capture and verify an existing trainer's setup without executing the trainer."""
import argparse
import ast
import hashlib
import json
import platform
from pathlib import Path
import struct
import sys

HERE = Path(__file__).resolve().parent
SUPPORTED_TRAINER_SHA256 = '41b40a9ac8cc6b5b53b45f7b2f696f0fe9488b7fd583da1f20be53b293af97a9'
FILES = ('trainer/baseline.py', 'trainer/rice.csv', 'job-config.json', 'requirements-lock.txt')


def sha256(raw):
    return hashlib.sha256(raw).hexdigest()


def identity(value):
    raw = json.dumps(value, sort_keys=True, separators=(',', ':'), allow_nan=False).encode()
    return sha256(raw)


def parse_config(raw):
    def pairs(items):
        result = {}
        for key, value in items:
            if key in result:
                raise ValueError('duplicate configuration key')
            result[key] = value
        return result
    def constant(value):
        raise ValueError('nonstandard JSON constant: ' + value)
    return json.loads(raw, object_pairs_hook=pairs, parse_constant=constant)


def trainer_contract(raw):
    if sha256(raw) != SUPPORTED_TRAINER_SHA256:
        raise ValueError('unsupported trainer revision; review the setup adapter before recapturing')
    tree = ast.parse(raw)
    names = {'EPOCHS', 'LEARNING_RATE', 'SALT', 'POSITIVE', 'DATA_SHA256'}
    constants = {}
    for node in tree.body:
        if isinstance(node, ast.Assign):
            for target in node.targets:
                if isinstance(target, ast.Name) and target.id in names:
                    constants[target.id] = ast.literal_eval(node.value)
    function = next(node for node in tree.body if isinstance(node, ast.FunctionDef)
                    and node.name == 'train_logistic')
    # Read the actual definition-time defaults, rather than assuming later global
    # assignment would change already-bound Python function defaults.
    defaults = {argument.arg: constants[value.id]
                for argument, value in zip(function.args.args[-2:], function.args.defaults)
                if isinstance(value, ast.Name)}
    if defaults != {'learning_rate': constants['LEARNING_RATE'], 'epochs': constants['EPOCHS']}:
        raise ValueError('trainer defaults differ from the reviewed adapter')
    imports = set()
    for node in ast.walk(tree):
        if isinstance(node, ast.Import):
            imports.update(alias.name.split('.')[0] for alias in node.names)
        elif isinstance(node, ast.ImportFrom) and node.module:
            imports.add(node.module.split('.')[0])
    if not imports <= sys.stdlib_module_names:
        raise ValueError('trainer declares imports outside the standard library')
    return constants, defaults, sorted(imports)


def describe_runtime():
    return {'implementation': platform.python_implementation(),
            'python_version': platform.python_version(),
            'compiler': platform.python_compiler(),
            'system': platform.system(), 'machine': platform.machine(),
            'pointer_bits': struct.calcsize('P') * 8}


def capture(root=HERE):
    root = Path(root)
    raw = {name: (root / name).read_bytes() for name in FILES}
    constants, defaults, imports = trainer_contract(raw['trainer/baseline.py'])
    config = parse_config(raw['job-config.json'])
    expected = {'mode': config.get('mode') if isinstance(config, dict) else None,
                'epochs': defaults['epochs'], 'learning_rate': defaults['learning_rate'],
                'split_salt': constants['SALT'], 'positive_class': constants['POSITIVE'],
                'scaler_fitted_on': 'training'}
    if (not isinstance(config, dict) or set(config) != set(expected)
            or config['mode'] not in ('validation', 'final')
            or type(config['epochs']) is not int
            or type(config['learning_rate']) not in (int, float) or config != expected):
        raise ValueError('configuration does not describe the existing trainer behavior')
    if sha256(raw['trainer/rice.csv']) != constants['DATA_SHA256']:
        raise ValueError('dataset differs from the trainer data version')
    argv = ['python', 'trainer/baseline.py']
    if config['mode'] == 'final':
        argv.append('--final')
    return {'format': 1, 'mode': 'setup capture only; trainer not executed',
            'files': {name: sha256(blob) for name, blob in raw.items()},
            'capture_tool_sha256': sha256((HERE / 'capture_job.py').read_bytes()),
            'effective_config': config,
            'dependencies': {'trainer_direct_imports': imports, 'third_party': []},
            'runtime': describe_runtime(), 'planned_argv': argv}


def verify(snapshot, expected_id, root=HERE):
    if identity(snapshot) != expected_id:
        raise ValueError('snapshot differs from the separately expected identity')
    current = capture(root)
    if not isinstance(snapshot, dict) or identity(current) != expected_id:
        differences = sorted(key for key in set(current) | set(snapshot)
                             if current.get(key) != snapshot.get(key)) if isinstance(snapshot, dict) else []
        raise ValueError('setup differs: ' + (', '.join(differences) or 'canonical values'))
    return {'setup_matches': True, 'training_started': False, 'recipe_id': expected_id}


def main():
    parser = argparse.ArgumentParser()
    commands = parser.add_subparsers(dest='command', required=True)
    save = commands.add_parser('capture')
    save.add_argument('--root', type=Path, default=HERE)
    save.add_argument('--out', type=Path, required=True)
    check = commands.add_parser('verify')
    check.add_argument('--root', type=Path, default=HERE)
    check.add_argument('--snapshot', type=Path, required=True)
    check.add_argument('--expected-id', required=True)
    args = parser.parse_args()
    if args.command == 'capture':
        snapshot = capture(args.root)
        with args.out.open('x') as file:
            file.write(json.dumps(snapshot, indent=2) + '\n')
        print(identity(snapshot))
    else:
        result = verify(parse_config(args.snapshot.read_bytes()), args.expected_id, args.root)
        print(json.dumps(result, indent=2))


if __name__ == '__main__':
    main()