import argparse
import base64
import hashlib
import json
import os
import subprocess
from datetime import datetime, timezone
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
CONFIG = '/tmp/codex/ltb1021/2026-11-01_icerti-web-ssh-connect_r01/ssh.conf'
HOST = 'contextra-icerti-web-diagnostic'


def capture(name, timeout=60):
    os.umask(0o077)
    source = ROOT / 'helpers' / (name + '.ps1')
    output = ROOT / 'private' / (name + '.json')
    if output.exists():
        raise RuntimeError('Evidence exists; use a new name')
    prefix = (
        "$ErrorActionPreference = 'Stop'\n"
        "$ProgressPreference = 'SilentlyContinue'\n"
        "[Console]::OutputEncoding = New-Object System.Text.UTF8Encoding($false)\n"
    )
    script = prefix + source.read_text(encoding='utf-8')
    encoded = base64.b64encode(script.encode('utf-16le')).decode('ascii')
    started = datetime.now(timezone.utc).isoformat()
    try:
        result = subprocess.run(
            ['ssh', '-F', CONFIG, HOST, 'powershell.exe', '-NoLogo', '-NoProfile',
             '-NonInteractive', '-EncodedCommand', encoded],
            capture_output=True, timeout=timeout,
        )
    except subprocess.TimeoutExpired:
        (ROOT / 'private' / (name + '.timeout.json')).write_text(
            json.dumps({'name': name, 'timeout_seconds': timeout,
                        'query_sha256': hashlib.sha256(script.encode('utf-8')).hexdigest()}) + '\n',
            encoding='utf-8',
        )
        raise RuntimeError('Remote query timed out: ' + name) from None
    if result.returncode:
        error = ROOT / 'private' / (name + '.error.txt')
        error.write_bytes(result.stderr)
        raise RuntimeError('Remote read failed; restricted error saved: ' + error.name)
    data = json.loads(result.stdout.decode('utf-8-sig'))
    record = {
        'started_utc': started,
        'finished_utc': datetime.now(timezone.utc).isoformat(),
        'query_sha256': hashlib.sha256(script.encode('utf-8')).hexdigest(),
        'data': data,
    }
    output.write_text(json.dumps(record, ensure_ascii=False, indent=2) + '\n', encoding='utf-8')
    print(json.dumps({'saved': str(output), 'top_level_keys': list(data)}, ensure_ascii=False))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('name')
    parser.add_argument('--timeout', type=int, default=60)
    args = parser.parse_args()
    capture(args.name, args.timeout)