with open('/home/andymccall/development/oxygen/docs/xenon_spec.md', 'r') as f:
    text = f.read()

# Remove all existing ```json, ```cbor, ``` blocks to start clean
import re
text = re.sub(r'```[a-z]*\n', '', text)
text = re.sub(r'```\n', '', text)

# Now, apply the blocks properly
text = text.replace('Receipt {\n', '```cbor\nReceipt {\n')
text = text.replace('  duration_ms    [8]: u64\n}\n', '  duration_ms    [8]: u64\n}\n```\n')

text = text.replace('Attest {\n', '```json\nAttest {\n')
text = text.replace('      "reproducible": bool\n    }\n  }\n}\n', '      "reproducible": bool\n    }\n  }\n}\n```\n')

text = text.replace('Signature {\n', '```cbor\nSignature {\n')
text = text.replace('  ml  [3]: bytes(3309)                 // Post-quantum signature\n}\n', '  ml  [3]: bytes(3309)                 // Post-quantum signature\n}\n```\n')

text = text.replace('PolicyRule {\n', '```json\nPolicyRule {\n')
text = text.replace('      signer_roles: [ "official-worker" ]\n    }\n  ]\n}\n', '      signer_roles: [ "official-worker" ]\n    }\n  ]\n}\n```\n')

# KDL block is already fixed mostly, let's ensure it's clean
if '```kdl' not in text:
    text = text.replace('schema-version 1\n', '```kdl\nschema-version 1\n')
    text = text.replace('    }\n}\n\n### 3.2', '    }\n}\n```\n\n### 3.2')

with open('/home/andymccall/development/oxygen/docs/xenon_spec.md', 'w') as f:
    f.write(text)