"""Run one supported local validation job; select only digest-bound completed reports."""
import argparse
from datetime import datetime, timezone
import hashlib
import json
import math
from pathlib import Path
import re
import shutil
import subprocess
import sys
import time
import uuid
from vendor import capture_job

HERE = Path(__file__).resolve().parent
RUN_ID = re.compile(r'run-[0-9a-f]{32}\Z')
SHA = re.compile(r'[0-9a-f]{64}\Z')
ARTIFACT_KIND = 'training_evaluation_report'


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


def now():
    return datetime.now(timezone.utc).isoformat(timespec='microseconds')


def write_state(folder, record):
    # Atomic replacement of this run's state file, not a crash-recovery service.
    path = folder / 'run.json.next'
    with path.open('x') as file:
        file.write(json.dumps(record, indent=2, allow_nan=False) + '\n')
    path.replace(folder / 'run.json')


def launch(argv, folder, timeout, stdout, stderr):
    return subprocess.run(argv, cwd=folder, timeout=timeout,
                          stdout=stdout, stderr=stderr, check=False)


def validate_report(raw, snapshot):
    report = capture_job.parse_config(raw)
    if not isinstance(report, dict) or not {'data', 'settings', 'model', 'validation'} <= set(report):
        raise ValueError('missing training-report fields')
    expected = snapshot['effective_config']
    try:
        if not all(isinstance(report[name], dict) for name in ('data', 'settings', 'model', 'validation')):
            raise ValueError('report sections must be objects')
        settings = report['settings']
        if ('test' in report or report['data'].get('sha256') != snapshot['files']['trainer/rice.csv']
                or type(settings.get('epochs')) is not int
                or settings['epochs'] != expected['epochs']
                or type(settings.get('learning_rate')) not in (int, float)
                or settings['learning_rate'] != expected['learning_rate']
                or settings.get('positive_class') != expected['positive_class']
                or settings.get('scaler_fitted_on') != 'training'
                or report['split']['salt'] != expected['split_salt']):
            raise ValueError('report does not match validation setup')
        metric = report['validation']['logistic']
        if (not isinstance(metric, dict) or type(metric['rows']) is not int or metric['rows'] < 1
                or type(metric['accuracy']) not in (int, float)
                or not math.isfinite(metric['accuracy']) or not 0 <= metric['accuracy'] <= 1):
            raise ValueError('invalid validation metric fields')
    except (KeyError, TypeError) as error:
        raise ValueError('malformed training report') from error
    return report


def execute(runs, snapshot, expected_recipe, root=HERE, timeout=30):
    if type(timeout) is not int or not 1 <= timeout <= 30:
        raise ValueError('timeout must be an integer from 1 to 30 seconds')
    capture_job.verify(snapshot, expected_recipe, root)
    if snapshot['effective_config']['mode'] != 'validation':
        raise ValueError('this job wrapper runs validation only, never --final')
    runs = Path(runs)
    runs.mkdir(parents=True, exist_ok=True)
    run_id = 'run-' + uuid.uuid4().hex
    folder = runs / run_id
    folder.mkdir()  # Exclusive allocation: there is no reuse/overwrite path.
    workspace = folder / 'workspace'
    workspace.mkdir()
    setup_raw = (json.dumps(snapshot, indent=2) + '\n').encode()
    (folder / 'setup.json').write_bytes(setup_raw)
    record = {'format': 1, 'run_id': run_id, 'state': 'planned',
              'setup_id': expected_recipe, 'setup_sha256': digest(setup_raw),
              'transitions': ['planned'],
              'created_at': now(), 'started_at': None, 'finished_at': None,
              'returncode': None, 'timeout_seconds': timeout, 'reason': '', 'artifact': None,
              'launcher_sha256': digest((HERE / 'run_jobs.py').read_bytes()),
              'argv': ['<sys.executable>', 'workspace/trainer/baseline.py', '--out', 'evaluation.partial.json']}
    write_state(folder, record)
    start = time.monotonic()
    try:
        for name in capture_job.FILES:
            target = workspace / name
            target.parent.mkdir(parents=True, exist_ok=True)
            shutil.copyfile(Path(root) / name, target)
        capture_job.verify(snapshot, expected_recipe, workspace)
        record.update(state='running', started_at=now())
        record['transitions'].append('running')
        write_state(folder, record)
        argv = [sys.executable, *record['argv'][1:]]
        with (folder/'stdout.txt').open('xb') as stdout, (folder/'stderr.txt').open('xb') as stderr:
            result = launch(argv, folder, timeout, stdout, stderr)
        record['returncode'] = result.returncode
        if result.returncode != 0:
            raise ValueError('trainer exited with a nonzero status')
        raw = (folder/'evaluation.partial.json').read_bytes()
        validate_report(raw, snapshot)
        final = folder/'evaluation.json'
        with final.open('xb') as file:
            file.write(raw)
        (folder/'evaluation.partial.json').unlink()
        record.update(state='succeeded', artifact={'kind': ARTIFACT_KIND,
                      'file': 'evaluation.json', 'sha256': digest(raw), 'bytes': len(raw)})
    except subprocess.TimeoutExpired:
        record.update(state='timed_out', reason='trainer exceeded the process timeout')
    except (OSError, ValueError, KeyError, TypeError) as error:
        record.update(state='failed', reason=type(error).__name__ + ': ' + str(error))
    record['transitions'].append(record['state'])
    record.update(finished_at=now(), elapsed_seconds=round(time.monotonic()-start, 6))
    record['logs'] = {name: digest((folder/name).read_bytes())
                      for name in ('stdout.txt','stderr.txt') if (folder/name).is_file()}
    write_state(folder, record)
    return {'run_id': run_id, 'record_sha256': digest((folder/'run.json').read_bytes()),
            'state': record['state'], 'artifact': record['artifact']}


def regular_file(path):
    if path.is_symlink() or not path.is_file():
        raise ValueError('selection needs a regular file inside the run directory')
    return path.read_bytes()


def select(runs, run_id, expected_record_sha256):
    if not isinstance(run_id, str) or not RUN_ID.fullmatch(run_id):
        raise ValueError('select an explicit run ID; aliases and paths are not accepted')
    if not isinstance(expected_record_sha256, str) or not SHA.fullmatch(expected_record_sha256):
        raise ValueError('selection requires the expected run-record SHA-256')
    folder = Path(runs) / run_id
    if folder.is_symlink() or not folder.is_dir():
        raise ValueError('run directory is missing or substituted')
    raw = regular_file(folder/'run.json')
    if digest(raw) != expected_record_sha256:
        raise ValueError('run record differs from the selected identity')
    record = capture_job.parse_config(raw)
    if not isinstance(record, dict) or record.get('run_id') != run_id or record.get('state') != 'succeeded':
        raise ValueError('only the selected succeeded run has an eligible artifact')
    setup_raw = regular_file(folder/'setup.json')
    snapshot = capture_job.parse_config(setup_raw)
    if digest(setup_raw) != record['setup_sha256'] or capture_job.identity(snapshot) != record['setup_id']:
        raise ValueError('setup snapshot binding differs')
    for name, expected in snapshot['files'].items():
        if name not in capture_job.FILES or digest(regular_file(folder/'workspace'/name)) != expected:
            raise ValueError('archived input bytes differ from setup')
    if set(snapshot['files']) != set(capture_job.FILES):
        raise ValueError('archived input inventory differs')
    for name, expected in record['logs'].items():
        if name not in ('stdout.txt', 'stderr.txt') or digest(regular_file(folder/name)) != expected:
            raise ValueError('archived log bytes differ from run record')
    artifact = record['artifact']
    if (not isinstance(artifact, dict) or artifact.get('kind') != ARTIFACT_KIND
            or artifact.get('file') != 'evaluation.json'):
        raise ValueError('unexpected artifact type or filename')
    report_raw = regular_file(folder/'evaluation.json')
    if digest(report_raw) != artifact['sha256'] or len(report_raw) != artifact['bytes']:
        raise ValueError('artifact bytes differ from the selected run record')
    report = validate_report(report_raw, snapshot)
    return {'run_id': run_id, 'record_sha256': expected_record_sha256,
            'artifact': artifact, 'setup_id': record['setup_id'],
            'validation': report['validation']['logistic']}


def main():
    parser = argparse.ArgumentParser()
    commands = parser.add_subparsers(dest='command', required=True)
    run = commands.add_parser('run')
    run.add_argument('--runs', type=Path, required=True)
    run.add_argument('--timeout', type=int, default=30)
    choose = commands.add_parser('select')
    choose.add_argument('--runs', type=Path, required=True)
    choose.add_argument('--run-id', required=True)
    choose.add_argument('--record-sha256', required=True)
    args = parser.parse_args()
    if args.command == 'run':
        setup = capture_job.capture(HERE)
        result = execute(args.runs, setup, capture_job.identity(setup), timeout=args.timeout)
    else:
        result = select(args.runs, args.run_id, args.record_sha256)
    print(json.dumps(result, indent=2))
    if args.command == 'run' and result['state'] != 'succeeded':
        raise SystemExit(1)


if __name__ == '__main__':
    main()