from preflight import R, g, api, HEAD
import json
import os
import signal
import subprocess
import tempfile
import time
import traceback
from pathlib import Path

plan = json.loads((R / 'plan.json').read_text())
built = json.loads((R / 'build.json').read_text())
active = Path(plan['active'])
overlay = Path(plan['overlay'])
normal = ['docker', 'compose', '--env-file', str(g.ENV), '-f', str(g.REPO / 'docker-compose.yml'), '-f', str(active)]
state = dict(stage='prepared', ai_start=0, limits=dict(total_seconds=240, readiness_seconds=120,
    host_ram_floor_gib=4, gpu_delta_mib=8192, log_bytes=2097152),
    apply_order=['CAS active image', 'CAS env image', 'compose no-start recreate Analytics only', 'container start', 'readiness/stopped'],
    rollback='Before recreate: CAS restore original config. After recreate failure: retain failed image, stop Analytics, restart=no actual/persistent; backups retained.')
path = R / 'reflection.json'

def save():
    path.write_text(json.dumps(state, indent=2))

def atomic(target, data, expected):
    assert g.sha(target) == expected, 'CONFIG_CAS'
    fd, tmp = tempfile.mkstemp(prefix='.rtsp-ready-', dir=target.parent)
    try:
        os.fchmod(fd, target.stat().st_mode & 0o777)
        with os.fdopen(fd, 'wb') as out:
            out.write(data)
            out.flush()
            os.fsync(out.fileno())
        assert g.sha(target) == expected, 'CONFIG_CAS'
        os.replace(tmp, target)
    finally:
        if os.path.exists(tmp):
            os.unlink(tmp)
    return g.sha(target)

def cmd(args, timeout=120):
    result = subprocess.run(args, text=True, capture_output=True, timeout=timeout)
    with (R / 'reflection-commands.log').open('a') as log:
        log.write(result.stdout[-20000:] + result.stderr[-20000:])
    assert result.returncode == 0, 'COMMAND_FAILED'
    return result.stdout

def alarm(*_):
    raise TimeoutError('REFLECTION_HELPER_CAP')

if __name__ == '__main__':
    assert not path.exists()
    tests = json.loads((R / 'image-verify.json').read_text())
    assert tests['tests'] == 2 and tests['errors'] == tests['failures'] == 0
    assert g.sh(['git', '-C', str(g.REPO), 'rev-parse', 'HEAD']).strip() == HEAD
    assert g.inspect()['Id'] == plan['analytics']['Id']
    assert g.others() == plan['other_containers']
    assert all(s['state'] == 'stopped' for s in api('/api/v1/streams')['streams'])
    assert g.sha(overlay) == plan['overlay_hash']
    assert all(g.sha(Path(p)) == h for p, h in plan['model_hashes'].items())
    assert g.sha(active) == plan['active_hash'] and g.sha(g.ENV) == plan['env_hash']
    cfg = g.yaml.safe_load(active.read_text())
    assert cfg['services']['analytics']['restart'] == 'on-failure:3'
    assert all(not s.get('autostart', False) for s in g.yaml.safe_load(overlay.read_text()).get('added', []))
    save()
    signal.signal(signal.SIGALRM, alarm)
    signal.alarm(240)
    recreated = False
    changed_active = changed_env = None
    try:
        cfg['services']['analytics']['image'] = built['image_id']
        changed_active = atomic(active, g.yaml.safe_dump(cfg).encode(), plan['active_hash'])
        state['active_after'] = changed_active
        save()
        lines = g.ENV.read_text().splitlines(keepends=True)
        indices = [i for i, line in enumerate(lines) if line.split('=', 1)[0] == 'CONTEXTRA_IMAGE']
        assert len(indices) == 1
        lines[indices[0]] = 'CONTEXTRA_IMAGE=' + built['tag'] + '\n'
        changed_env = atomic(g.ENV, ''.join(lines).encode(), plan['env_hash'])
        state['env_after'] = changed_env
        save()
        cmd(normal + ['config', '--quiet'])
        assert all(s['state'] == 'stopped' for s in api('/api/v1/streams')['streams'])
        state['stage'] = 'recreating'
        recreated = True
        save()
        ledger = json.loads((R / 'ledger.json').read_text())
        ledger['reflected'] += 1
        (R / 'ledger.json').write_text(json.dumps(ledger, indent=2))
        cmd(normal + ['up', '--no-start', '--no-deps', '--force-recreate', '--pull', 'never', 'analytics'])
        cmd(['docker', 'start', g.NAME])
        end = time.monotonic() + 120
        while time.monotonic() < end:
            current = g.inspect()
            assert current['Image'] == built['image_id'] and current['State']['Running'] and current['RestartCount'] == 0
            assert g.gpu() <= plan['gpu_mib'] + 8192
            available = next(int(line.split()[1]) for line in Path('/proc/meminfo').read_text().splitlines() if line.startswith('MemAvailable:'))
            assert available > 4 * 1024**2
            try:
                ready = api('/ready')
                if current['State'].get('Health', {}).get('Status') == 'healthy':
                    break
            except Exception:
                pass
            time.sleep(1)
        else:
            raise TimeoutError('READINESS_TIMEOUT')
        streams = api('/api/v1/streams')['streams']
        assert len(streams) == len(plan['streams']) and all(s['state'] == 'stopped' for s in streams)
        schema = api('/openapi.json')
        assert '/api/v1/streams' in schema['paths'] and '/api/v1/cameras/start' not in schema['paths']
        assert {'server_id', 'source_id', 'frame_time_anchor'} <= set(schema['components']['schemas']['StreamSpec']['properties'])
        (R / 'openapi.json').write_text(json.dumps(schema))
        probe = 'import json,urllib.request; print(json.dumps({m:json.load(urllib.request.urlopen("http://localhost:8000/v2/models/"+m,timeout=3))["versions"] for m in ("RFDETRFireSmoke","RFDETRMedium","RFDETRKeypointPreview")}))'
        versions = json.loads(cmd(['docker', 'exec', g.NAME, 'python3', '-c', probe], 20))
        assert all(v == ['2'] for v in versions.values())
        current = g.inspect()
        assert current['HostConfig']['RestartPolicy'] == plan['normal_restart']
        env_before = dict(x.split('=', 1) for x in plan['analytics']['Config']['Env'])
        env_after = dict(x.split('=', 1) for x in current['Config']['Env'])
        assert env_before == env_after, 'CONTAINER_ENV_DRIFT'
        mounts = lambda c: sorted((x['Source'], x['Destination'], x['RW']) for x in c['Mounts'])
        assert mounts(current) == mounts(plan['analytics']), 'MOUNT_DRIFT'
        assert g.sha(overlay) == plan['overlay_hash']
        assert g.others() == plan['other_containers']
        state.update(stage='PASS', image=built['image_id'], tag=built['tag'], readiness=ready,
                     streams=streams, model_versions=versions, config_preserved=True,
                     restart=current['HostConfig']['RestartPolicy'], container_id=current['Id'])
    except Exception:
        state['stage'] = 'FAIL'
        state['error'] = traceback.format_exc()
        signal.alarm(0)
        if recreated:
            try:
                current = g.inspect()
                logs = subprocess.run(['docker', 'logs', '--tail', '3000', g.NAME], capture_output=True, timeout=15)
                (R / 'failure-runtime.log').write_bytes((logs.stdout + logs.stderr)[-2097152:])
                cmd(['docker', 'update', '--restart=no', g.NAME])
                cmd(['docker', 'stop', '-t', '15', g.NAME], 30)
                data = g.yaml.safe_load(active.read_text())
                data['services']['analytics']['restart'] = 'no'
                state['safe_active_hash'] = atomic(active, g.yaml.safe_dump(data).encode(), changed_active)
                state['safe_stopped'] = True
            except Exception:
                state['cleanup_error'] = traceback.format_exc()
        else:
            if changed_env:
                atomic(g.ENV, (R / 'env.before').read_bytes(), changed_env)
            if changed_active:
                atomic(active, (R / 'active.before').read_bytes(), changed_active)
    finally:
        signal.alarm(0)
        save()
    print(json.dumps({k: state.get(k) for k in ('stage', 'image', 'model_versions', 'config_preserved', 'error', 'safe_stopped', 'cleanup_error')}))
    assert state['stage'] == 'PASS'