#!/usr/bin/env python3
"""Remove account-bound state from one Codex rollout, preserving original history."""
import argparse
import collections
import hashlib
import json
import os
from pathlib import Path
import shutil
import tempfile
from datetime import datetime, timezone


def repair(raw: bytes) -> tuple[bytes, collections.Counter]:
    kept = []
    removed = collections.Counter()
    for line in raw.splitlines(keepends=True):
        item = json.loads(line)
        payload = item.get('payload', {})
        if item.get('type') == 'response_item' and payload.get('type') in ('reasoning', 'compaction') and payload.get('encrypted_content'):
            removed[payload['type']] += 1
            continue
        if item.get('type') == 'compacted' and any(x.get('encrypted_content') for x in payload.get('replacement_history') or []):
            # Removing the checkpoint restores the earlier original history on replay.
            # Removing just its encrypted item would silently discard that context.
            removed['encrypted compaction checkpoints'] += 1
            continue
        kept.append(line)
    result = b''.join(kept)
    def encrypted(value):
        if isinstance(value, dict):
            return any((k == 'encrypted_content' and bool(v)) or encrypted(v) for k, v in value.items())
        return isinstance(value, list) and any(encrypted(v) for v in value)
    if any(encrypted(json.loads(line)) for line in result.splitlines()):
        raise ValueError('Unexpected encrypted data remains; refusing to change session')
    return result, removed


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('session', type=Path)
    parser.add_argument('--apply', action='store_true', help='Back up and atomically replace the rollout; close this chat first')
    args = parser.parse_args()
    path = args.session.resolve(strict=True)
    raw = path.read_bytes()
    result, removed = repair(raw)
    print(json.dumps({'session': str(path), 'removed': dict(removed), 'before_bytes': len(raw), 'after_bytes': len(result)}))
    if not args.apply or result == raw:
        return
    backup = path.with_name(path.name + '.before-account-repair-' + datetime.now(timezone.utc).strftime('%Y%m%dT%H%M%S%fZ') + '.bak')
    with backup.open('xb') as f:
        os.chmod(backup, 0o600)
        f.write(raw)
        f.flush()
        os.fsync(f.fileno())
    fd, name = tempfile.mkstemp(prefix='.' + path.name, dir=path.parent)
    try:
        with os.fdopen(fd, 'wb') as f:
            f.write(result)
            f.flush()
            os.fsync(f.fileno())
        shutil.copystat(path, name)
        if hashlib.sha256(path.read_bytes()).digest() != hashlib.sha256(raw).digest():
            raise RuntimeError('Session changed during repair; close the chat and try again')
        os.replace(name, path)
    finally:
        if os.path.exists(name):
            os.unlink(name)
    print('Backup:', backup)
    print('Repaired:', path)


if __name__ == '__main__':
    main()