"""Finite, no-training execution of the separately registered 32-call gate.

All files are exclusive creations. A partial/failed row is preserved and is
not retried. Runtime outputs are sealed before external reports are written.
The preparation module owns source/input authority; no summary Boolean alone
is accepted as proof that a subprocess or a complete engineering gate passed.
"""
from __future__ import annotations

import argparse
from datetime import datetime, timezone
import importlib
import json
import os
from pathlib import Path
import shutil
import subprocess
import sys
import time
import traceback
import urllib.request
import zipfile

import numpy as np

import engineering_inputs as I
import engineering_plan as PLAN
import runtime_light as R
import runtime_continuation as C


GPU = 'GPU-79f3a85d-adce-7833-7003-6488140ca596'
ROOT = Path('/data/digital-drosophila')
MODEL = Path(__file__).resolve().parent
EXP = ROOT / 'experiments/phase4an-engineering-validation'
require = R.require


def canonical(value):
    return json.dumps(value, sort_keys=True, separators=(',', ':'), allow_nan=False)


def same_json(left, right):
    return canonical(left) == canonical(right)


def preparation():
    return importlib.import_module('engineering_prepare')


def row_from(registration, row_id):
    plan = PLAN.validate_plan(registration['plan'])
    rows = [row for row in plan['rows'] if row['row_id'] == row_id]
    require(len(rows) == 1, 'Only an exact registered row may execute')
    return rows[0]


def locations(row_id):
    require(type(row_id) is str and row_id and all(ch.isalnum() or ch in '-_' for ch in row_id),
            'Unsafe row identity')
    return dict(output=EXP / 'runs' / row_id,
                runtime=EXP / 'receipts' / (row_id + '.runtime.json'),
                execution=EXP / 'receipts' / (row_id + '.json'),
                derivation=EXP / 'receipts' / (row_id + '.derivation.json'),
                stdout=EXP / 'logs' / (row_id + '.stdout.log'),
                stderr=EXP / 'logs' / (row_id + '.stderr.log'))


def disk_budget(registration):
    limit = registration['budget']['output_budget_bytes']
    reserve = registration['budget']['minimum_free_reserve_bytes']
    require(type(limit) is type(reserve) is int and limit > 0 and reserve > 0, 'Invalid registered budget')
    files = [p for p in EXP.rglob('*') if p.is_file()]
    require(not any(p.is_symlink() for p in EXP.rglob('*')), 'Experiment contains a linked output')
    used = sum(p.stat().st_size for p in files)
    free = shutil.disk_usage(EXP).free
    require(used <= limit and free >= reserve + max(0, limit - used), 'Engineering budget/reserve exceeded; preserve partials')
    return dict(used_bytes=used, free_bytes=free, hard_cap_bytes=limit, reserve_bytes=reserve)


def preflight():
    """Never evict an existing service; configure only this child process."""
    os.environ['CUDA_VISIBLE_DEVICES'] = GPU
    os.environ['CUDA_HOME'] = '/usr/local/cuda'
    os.environ['PATH'] = str(ROOT / 'models/male-cns-lif/.venv/bin') + ':/usr/local/cuda/bin:' + os.environ['PATH']
    os.environ['MAX_JOBS'] = '2'
    os.environ['TORCH_CUDA_ARCH_LIST'] = '8.9'
    for port in (8188, 8189, 8190):
        with urllib.request.urlopen(f'http://192.168.10.51:{port}/queue', timeout=5) as response:
            queue = json.load(response)
        require(queue['queue_running'] == [] and queue['queue_pending'] == [],
                'ComfyUI busy: defer rather than evict')


def derive_case(template, row, parent_sha):
    require(type(template) is dict and set(template) == {'schema', 'row_id', 'original_spec', 'start', 'stop',
            'budget_root', 'output_budget_bytes'}, 'Unexpected registered case template fields')
    require(template['schema'] == 'phase4an-engineering-case-template-v1' and template['row_id'] == row['row_id'] and
            type(template['start']) is int and type(template['stop']) is int and
            (template['start'], template['stop']) == (row['start'], row['stop']), 'Template differs from exact plan')
    require((row['parent_row_id'] is None) == (parent_sha is None), 'Parent linkage differs from registered plan')
    return I.derive_spec(template['original_spec'], start=row['start'], stop=row['stop'],
        budget_root=template['budget_root'], output_budget_bytes=template['output_budget_bytes'], parent_sha=parent_sha)


def npz_storage(path):
    with zipfile.ZipFile(path) as archive:
        rows = archive.infolist()
        require(rows and all(r.compress_type == zipfile.ZIP_DEFLATED and r.filename.endswith('.npy') for r in rows),
                'Declared lossless compressed state format differs from actual archive')
        return dict(file_bytes=Path(path).stat().st_size,
                    uncompressed_entry_bytes=sum(r.file_size for r in rows),
                    compressed_entry_bytes=sum(r.compress_size for r in rows),
                    entries=[dict(name=r.filename, bytes=r.file_size, compressed_bytes=r.compress_size,
                                  compression=r.compress_type) for r in rows])


def runtime_environment():
    """Read settings after the runtime imported Torch; change none of them."""
    torch = sys.modules.get('torch')
    require(torch is not None and os.environ.get('CUDA_VISIBLE_DEVICES') == GPU and
            torch.cuda.device_count() == 1 and '4090' in torch.cuda.get_device_name(0),
            'Actual runtime is not the registered single RTX4090')
    properties = torch.cuda.get_device_properties(0)
    inventory = subprocess.run(['nvidia-smi', '-i', GPU,
        '--query-gpu=uuid,name,memory.total,driver_version', '--format=csv,noheader'],
        check=True, capture_output=True, text=True, timeout=20).stdout.strip()
    require(inventory.startswith(GPU + ','), 'GPU hardware inventory does not match the registered UUID')
    return dict(python_executable=sys.executable, python_version=sys.version, numpy_version=np.__version__,
        torch_version=str(torch.__version__), torch_cuda_version=torch.version.cuda,
        device_name=properties.name, total_memory_bytes=int(properties.total_memory),
        gpu_uuid=GPU, nvidia_smi_inventory=inventory,
        matmul_allow_tf32=bool(torch.backends.cuda.matmul.allow_tf32),
        cudnn_allow_tf32=bool(torch.backends.cudnn.allow_tf32),
        deterministic_algorithms=bool(torch.are_deterministic_algorithms_enabled()),
        cudnn_deterministic=bool(torch.backends.cudnn.deterministic),
        cudnn_benchmark=bool(torch.backends.cudnn.benchmark),
        float32_matmul_precision=torch.get_float32_matmul_precision(),
        CUBLAS_WORKSPACE_CONFIG=os.environ.get('CUBLAS_WORKSPACE_CONFIG'),
        TORCH_CUDA_ARCH_LIST=os.environ.get('TORCH_CUDA_ARCH_LIST'),
        settings_changed_by_environment_recorder=False,
        environment_recorder_adds_no_timing=True,
        build_timing_source='Use inherited summary.build_or_cache_seconds; build/cache is also included in runtime wall time')


def verify_environment(recorded, component):
    require(type(recorded) is type(component) is dict and
            recorded.get('gpu_uuid') == GPU and recorded.get('settings_changed_by_environment_recorder') is False,
            'Missing or invalid actual environment receipt')
    require(component.get('cuda_matmul_allow_tf32') is False and
            component.get('float32_matmul_precision') == 'highest' and
            component.get('precision_configuration_changed_by_component') is False,
            'Frozen AM precision configuration changed; never correct it automatically')
    pairs = {'matmul_allow_tf32': 'cuda_matmul_allow_tf32', 'cudnn_allow_tf32': 'cudnn_allow_tf32',
             'deterministic_algorithms': 'deterministic_algorithms', 'cudnn_deterministic': 'cudnn_deterministic',
             'cudnn_benchmark': 'cudnn_benchmark', 'float32_matmul_precision': 'float32_matmul_precision',
             'CUBLAS_WORKSPACE_CONFIG': 'CUBLAS_WORKSPACE_CONFIG', 'device_name': 'device_name',
             'torch_version': 'torch', 'torch_cuda_version': 'cuda', 'numpy_version': 'numpy'}
    for left, right in pairs.items():
        require(left in recorded and right in component and same_json(recorded[left], component[right]),
                'Actual runtime/APL environment disagreement: ' + left)
    require(type(recorded['device_name']) is str and '4090' in recorded['device_name'], 'Wrong physical device class')


def check_runtime_record(reg, reg_sha, row, record):
    """Recheck actual bytes, runtime completion and parent chain, not intent."""
    paths = locations(row['row_id'])
    require(record.get('complete') is True and record.get('registration_sha256') == reg_sha and
            same_json(record.get('row'), row) and record.get('training') is False,
            'Runtime record is not bound to this registered row')
    case = Path(reg['cases'][row['row_id']]['case_dir'])
    checkpoint = C.load_checkpoint(paths['output'], record['checkpoint_manifest_sha256'])
    summary = checkpoint['summary']
    verify_environment(record['environment'], summary['apl_component']['environment'])
    require(checkpoint['steps'] == row['steps'] and checkpoint['global_tick_offset'] == row['start'] and
            checkpoint['global_next_tick'] == row['stop'] and summary['nodes'] == 166700 and summary['edges'] == 25582938 and
            summary['diagnostic_mode'] == row['mode'], 'Runtime clock/population/mode differs from plan')
    require(record['case_sha256'] == R.digest(case / 'spec.json') and
            record['requests_sha256'] == R.digest(case / 'requests.npz') and
            record['derivation_sha256'] == R.digest(paths['derivation']), 'Derived input/receipt bytes changed')
    template = R._read_json(case / 'template-spec.json')
    parent_pin = None
    if row['parent_row_id'] is not None:
        parent = verify_execution(reg, reg_sha, row_from(reg, row['parent_row_id']))
        parent_pin = parent['checkpoint_manifest_sha256']
    expected, derivation = derive_case(template, row, parent_pin)
    require(same_json(R._read_json(case / 'spec.json'), expected) and
            same_json(R._read_json(paths['derivation']), derivation), 'Case was not generated by the registered exact rule')
    require(record['parent_checkpoint_manifest_sha256'] == parent_pin and
            summary.get('parent_checkpoint_manifest_sha256') == parent_pin, 'Execution parent SHA mismatch')
    require(record['checkpoint_inventory_sha256'] == R.digest(paths['output'] / C.MANIFEST), 'Checkpoint manifest changed')
    for name in ('state.npz', 'terminal-queues.npz'):
        require(same_json(record['storage'][name], npz_storage(paths['output'] / name)), 'Storage measurement changed')
    return record


def verify_execution(reg, reg_sha, row):
    paths = locations(row['row_id'])
    receipt = R._read_json(paths['execution'])
    require(receipt.get('complete') is True and type(receipt.get('process_exit_code')) is int and receipt['process_exit_code'] == 0 and
            receipt.get('registration_sha256') == reg_sha and receipt.get('row_id') == row['row_id'],
            'Missing completed process receipt')
    for field, key in (('runtime_record_sha256', 'runtime'), ('stdout_sha256', 'stdout'), ('stderr_sha256', 'stderr')):
        require(receipt[field] == R.digest(paths[key]), 'Execution log/runtime receipt changed')
    require(not (EXP / 'audit/stopped.json').exists(), 'A failed/deferred run requires explicit review, not automatic retry')
    record = check_runtime_record(reg, reg_sha, row, R._read_json(paths['runtime']))
    require(receipt['checkpoint_manifest_sha256'] == record['checkpoint_manifest_sha256'], 'Execution/checkpoint linkage mismatch')
    return record


def worker(reg_sha, row_id):
    require(sys.dont_write_bytecode, 'Use the scientific venv with -B')
    reg = preparation().read_registration(reg_sha)
    row = row_from(reg, row_id)
    paths = locations(row_id)
    require(not paths['output'].exists() and not paths['runtime'].exists() and not paths['derivation'].exists(),
            'Never overwrite/retry a partial row')
    budget_before = disk_budget(reg)
    preflight()
    parent_sha = None
    parent_dir = None
    if row['parent_row_id'] is not None:
        parent_record = verify_execution(reg, reg_sha, row_from(reg, row['parent_row_id']))
        parent_sha = parent_record['checkpoint_manifest_sha256']
        parent_dir = locations(row['parent_row_id'])['output']
    case = Path(reg['cases'][row_id]['case_dir'])
    template = R._read_json(case / 'template-spec.json')
    spec, derivation = derive_case(template, row, parent_sha)
    require(not (case / 'spec.json').exists(), 'Derived spec already exists; preserve partial instead of rerun')
    R._save(case / 'spec.json', spec)
    R._save(paths['derivation'], derivation)
    overlay = reg['overlays'][row['overlay']]
    start_utc = datetime.now(timezone.utc).isoformat()
    start = time.perf_counter()
    if parent_dir is None:
        function = R.run_reference if row['mode'] == 'reference' else R.run_light
        function(case, paths['output'], EXP / 'extension-build', overlay['overlay_file'], overlay['overlay_manifest'])
        pin = C.seal_checkpoint(paths['output'])['sha256']
    else:
        result = C.run_continued(case, paths['output'], EXP / 'extension-build',
            overlay['overlay_file'], overlay['overlay_manifest'], parent_dir, mode=row['mode'])
        pin = result['checkpoint_manifest_sha256']
    elapsed = time.perf_counter() - start
    budget_after = disk_budget(reg)
    record = dict(schema='phase4an-engineering-runtime-record-v1', complete=True, training=False,
        registration_sha256=reg_sha, row=row, case_sha256=R.digest(case / 'spec.json'),
        requests_sha256=R.digest(case / 'requests.npz'), derivation_sha256=R.digest(paths['derivation']),
        checkpoint_manifest_sha256=pin, checkpoint_inventory_sha256=R.digest(paths['output'] / C.MANIFEST),
        parent_checkpoint_manifest_sha256=parent_sha, runtime_and_immediate_seal_wall_seconds=elapsed,
        started_utc=start_utc, completed_utc=datetime.now(timezone.utc).isoformat(),
        environment=runtime_environment(),
        budget_before=budget_before, budget_after=budget_after,
        storage={name: npz_storage(paths['output'] / name) for name in ('state.npz', 'terminal-queues.npz')},
        no_result_selected_retry=True, learning_pass_claimed=False)
    # Do not write this or any later report INSIDE the sealed checkpoint.
    check_runtime_record(reg, reg_sha, row, record)
    preparation().read_registration(reg_sha)
    R._save(paths['runtime'], record)
    print('COMPLETED ENGINEERING ROW', row_id, flush=True)


def run_registered(reg_sha):
    require(sys.dont_write_bytecode, 'Use python -B')
    reg = preparation().read_registration(reg_sha)
    require(not (EXP / 'audit/stopped.json').exists() and not (EXP / 'audit/execution-complete.json').exists(),
            'Stopped/completed gate must not be restarted')
    started = time.perf_counter()
    completed = []
    try:
        for row in reg['plan']['rows']:
            paths = locations(row['row_id'])
            if paths['execution'].exists():
                verify_execution(reg, reg_sha, row)
                completed.append(dict(row_id=row['row_id'], execution_sha256=R.digest(paths['execution']),
                                      reverified_existing=True))
                continue
            require(not any(paths[k].exists() for k in ('output', 'runtime', 'stdout', 'stderr', 'derivation')),
                    'Partial row present; no automatic retry')
            disk_budget(reg)
            preflight()
            began = time.perf_counter()
            print('START', row['row_id'], flush=True)
            with paths['stdout'].open('xb') as out, paths['stderr'].open('xb') as err:
                result = subprocess.run([sys.executable, '-B', str(Path(__file__)), 'worker',
                    '--registration-sha256', reg_sha, '--row-id', row['row_id']], cwd=MODEL, stdout=out, stderr=err)
            require(result.returncode == 0, 'Worker failed; preserve its stdout/stderr and partial output')
            record = check_runtime_record(reg, reg_sha, row, R._read_json(paths['runtime']))
            receipt = dict(schema='phase4an-engineering-execution-receipt-v1', complete=True,
                process_exit_code=result.returncode, row_id=row['row_id'], registration_sha256=reg_sha,
                runtime_record_sha256=R.digest(paths['runtime']), stdout_sha256=R.digest(paths['stdout']),
                stderr_sha256=R.digest(paths['stderr']), process_wall_seconds=time.perf_counter() - began,
                checkpoint_manifest_sha256=record['checkpoint_manifest_sha256'], scientific_pass_claimed=False)
            R._save(paths['execution'], receipt)
            verify_execution(reg, reg_sha, row)
            completed.append(dict(row_id=row['row_id'], execution_sha256=R.digest(paths['execution']),
                                  reverified_existing=False))
            print('COMPLETE', row['row_id'], flush=True)
        require(len(completed) == 32, 'Incomplete registered inventory')
        preparation().read_registration(reg_sha)
        R._save(EXP / 'audit/execution-complete.json', dict(complete=True, registration_sha256=reg_sha,
            completed=completed, expected_calls=32, executed_ticks=11074, wall_seconds=time.perf_counter() - started,
            numerical_gate_evaluated=False, training=False, learning_pass_claimed=False))
    except Exception as error:
        R._save(EXP / 'audit/stopped.json', dict(complete=False, type=type(error).__name__, message=str(error),
            traceback=traceback.format_exc(), completed_before_error=completed, partial_outputs_preserved=True,
            no_automatic_retry=True, scientific_failure_not_inferred_from_engineering_failure=True))
        raise


def verified_inventory(reg, reg_sha):
    """A completed child and every planned descriptor must remain auditable."""
    plan = PLAN.validate_plan(reg['plan'])
    receipt = R._read_json(EXP / 'audit/execution-complete.json')
    require(receipt.get('complete') is True and receipt.get('registration_sha256') == reg_sha and
            type(receipt.get('expected_calls')) is int and receipt['expected_calls'] == plan['total_calls'] and
            type(receipt.get('executed_ticks')) is int and receipt['executed_ticks'] == plan['total_ticks'] and
            receipt.get('training') is False and receipt.get('learning_pass_claimed') is False,
            'Execution inventory is not a complete no-training registered gate')
    require(type(receipt.get('completed')) is list and len(receipt['completed']) == plan['total_calls'],
            'Execution inventory truncated')
    records = {}
    for row, entry in zip(plan['rows'], receipt['completed']):
        require(entry.get('row_id') == row['row_id'] and
                entry.get('execution_sha256') == R.digest(locations(row['row_id'])['execution']),
                'Execution completion inventory reordered or receipt changed')
        records[row['row_id']] = verify_execution(reg, reg_sha, row)
    return records


def evaluate(reg_sha):
    reg = preparation().read_registration(reg_sha)
    records = verified_inventory(reg, reg_sha)
    report = importlib.import_module('engineering_analysis').analyze(reg, reg_sha, records)
    require(type(report) is dict and report.get('training') is False and report.get('learning_pass_claimed') is False,
            'Engineering analysis must not claim training or learning')
    preparation().read_registration(reg_sha)
    return reg, report


def analyze(reg_sha):
    require(sys.dont_write_bytecode, 'Use python -B')
    target = EXP / 'audit/engineering-analysis.json'
    require(not target.exists() and not (EXP / 'manifest.json').exists(), 'Analysis/seal already exists; do not overwrite')
    began = time.perf_counter()
    _, report = evaluate(reg_sha)
    R._save(target, report)
    R._save(EXP / 'audit/analysis-execution.json', dict(registration_sha256=reg_sha,
        analysis_sha256=R.digest(target), wall_seconds=time.perf_counter() - began, GPU_used=False,
        training=False, learning_pass_claimed=False))
    print(report['status'], flush=True)


def seal_experiment():
    """Seal only the new engineering experiment, not the unfinished lab model."""
    controls = ('manifest.json', 'checksums.sha256')
    require(not any((EXP / name).exists() for name in controls), 'Do not replace an existing or partial seal')
    paths = sorted(EXP.rglob('*'))
    require(not any(path.is_symlink() for path in paths), 'Cannot seal linked artifacts')
    require(not any('__pycache__' in path.relative_to(EXP).parts for path in paths), 'Unexpected bytecode cache in new experiment')
    rows = [dict(path=path.relative_to(EXP).as_posix(), bytes=path.stat().st_size, sha256=R.digest(path))
            for path in paths if path.is_file()]
    require(rows, 'Cannot seal an empty experiment')
    # Only ROOT controls are excluded from their own hash. Nested checkpoint
    # manifests and source snapshots are ordinary payload and remain included.
    R._save(EXP / controls[0], dict(root=str(EXP), files=rows))
    with (EXP / controls[1]).open('x', encoding='utf-8', newline='\n') as stream:
        stream.writelines(row['sha256'] + '  ' + row['path'] + '\n' for row in rows)
        stream.write(R.digest(EXP / controls[0]) + '  manifest.json\n')
    return dict(manifest_sha256=R.digest(EXP / controls[0]), checksums_sha256=R.digest(EXP / controls[1]))


def finalize(reg_sha):
    require(sys.dont_write_bytecode, 'Use python -B')
    require(not (EXP / 'manifest.json').exists() and not (EXP / 'audit/final-status.json').exists(),
            'Final/partial seal already exists; preserve it')
    reg, actual = evaluate(reg_sha)
    target = EXP / 'audit/engineering-analysis.json'
    require(same_json(actual, R._read_json(target)), 'Recomputed analysis differs from saved evidence')
    execution = R._read_json(EXP / 'audit/analysis-execution.json')
    require(execution.get('registration_sha256') == reg_sha and execution.get('analysis_sha256') == R.digest(target),
            'Saved analysis is not bound to its execution receipt')
    authority = preparation().verify_authority(reg['anchors'])
    old = preparation().verify_old_integrity(reg)
    R._save(EXP / 'regression/AM-seals-after.json', authority)
    R._save(EXP / 'regression/old-integrity-after.json', old)
    preparation().read_registration(reg_sha)
    final = dict(schema='phase4an-engineering-final-v1', status=actual['status'],
        registration_sha256=reg_sha, analysis_sha256=R.digest(target),
        analysis_recomputed_exactly=True, completed_calls=32, executed_ticks=11074,
        old_integrity_pass=True, old_integrity_roots=len(old), AM_seals_preserved=True,
        old_five_reference_evidence='Sealed AM five-trial evidence retained; not rerun by this engineering gate',
        training=False, learning_pass_claimed=False, learned_candidate=False,
        model_directory_sealed=False, used_engineering_source_snapshots_sealed=True,
        budget=disk_budget(reg),
        scope='Finite registered 100 ms input and fixed continuation cuts only; no universal determinism, learning or physiology claim')
    R._save(EXP / 'audit/final-status.json', final)
    pins = seal_experiment()
    helper = preparation().load_module('phase4an_pinned_AI_publication',
        ROOT / preparation().AI_HELPER_REL, preparation().AI_HELPER_SHA)
    verified = helper.verify_seal(EXP)
    require(verified['manifest_sha256'] == pins['manifest_sha256'] and
            verified['checksums_sha256'] == pins['checksums_sha256'], 'Final experiment seal verification failed')
    print(canonical(dict(status=final['status'], **pins)), flush=True)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('action', choices=('run', 'worker', 'analyze', 'finalize'))
    parser.add_argument('--registration-sha256', required=True)
    parser.add_argument('--row-id')
    args = parser.parse_args()
    if args.action == 'worker':
        require(args.row_id is not None, 'Worker requires a row ID')
        worker(args.registration_sha256, args.row_id)
    else:
        require(args.row_id is None, 'Full inventory run has no selected row')
        {'run': run_registered, 'analyze': analyze, 'finalize': finalize}[args.action](args.registration_sha256)


if __name__ == '__main__':
    main()