#!/usr/bin/env python3
"""Exercise the real Chio MCP edge without depending on a model's tool choice."""
import json
import os
import selectors
import shutil
import subprocess
import sys
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parent
os.chdir(ROOT)
if len(sys.argv) == 1:
    policy = ROOT / "policy.yaml"
elif len(sys.argv) == 3 and sys.argv[1] == "--policy":
    policy = Path(sys.argv[2]).resolve()
else:
    raise SystemExit("Usage: python3 run.py [--policy PATH]")
if not policy.is_file():
    raise SystemExit(f"Policy file does not exist: {policy}")
STATE = ROOT / '.chio'
STATE.mkdir(mode=0o700, exist_ok=True)
binary = shutil.which('chio')
if not binary:
    raise SystemExit('Install chio and put it on PATH before running this example.')
if os.getuid() == 0:
    raise SystemExit('Run this example as a non-root user; the cage policy refuses a root execution identity.')

SERVER_ID = 'workspace'
SERVER_NAME = 'Larkfield pricing workspace'
SERVER_VERSION = '1'
PYTHON = os.path.realpath(sys.executable)
SERVER_SCRIPT = str(ROOT / 'server.py')

# chio mcp serve requires a signed native-launch policy and a signed tool
# manifest, both bound to this exact interpreter, script path and working
# directory. chio security provision-native-mcp-demo mints that material; run
# it fresh here for this copy's own path instead of shipping a static artifact
# that would only match wherever it was first generated.
LAUNCH = STATE / 'launch'
if LAUNCH.exists():
    shutil.rmtree(LAUNCH)
provision = subprocess.run(
    [binary, 'security', 'provision-native-mcp-demo',
     '--output-dir', str(LAUNCH), '--discover-tools',
     '--target', PYTHON, '--target-arg', SERVER_SCRIPT,
     '--working-directory', str(ROOT),
     '--execution-uid', str(os.getuid()), '--execution-gid', str(os.getgid()),
     '--server-id', SERVER_ID, '--server-name', SERVER_NAME, '--server-version', SERVER_VERSION],
    capture_output=True, text=True)
if provision.returncode != 0:
    raise SystemExit(f'Provisioning the cage policy and signed manifest failed:\n{provision.stderr}')
signed_manifest = LAUNCH / 'signed-manifest.json'
manifest_public_key = (LAUNCH / 'manifest-public-key').read_text().strip()
cage_policy = LAUNCH / 'cage-launch-policy.json'
cage_policy_signer = (LAUNCH / 'cage-policy-signer').read_text().strip()

TARGET = ROOT / 'workspace' / 'src' / 'pricing.py'
original = TARGET.read_bytes()
command = [binary, '--receipt-db', str(STATE / 'receipts.db'), '--session-db', str(STATE / 'session.db'),
           'mcp', 'serve', '--policy', str(policy), '--server-id', SERVER_ID,
           '--server-name', SERVER_NAME, '--server-version', SERVER_VERSION,
           '--signed-manifest', str(signed_manifest), '--manifest-public-key', manifest_public_key,
           '--cage-policy', str(cage_policy), '--cage-policy-signer', cage_policy_signer,
           '--', PYTHON, SERVER_SCRIPT]
with (STATE / 'server.log').open('w') as log:
    process = subprocess.Popen(command, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
                               stderr=log, bufsize=0)
    selector = selectors.DefaultSelector()
    selector.register(process.stdout, selectors.EVENT_READ)
    pending = bytearray()
    def send(method, params, identity=None):
        message = {'jsonrpc': '2.0', 'method': method, 'params': params}
        if identity is not None:
            message['id'] = identity
        process.stdin.write((json.dumps(message) + '\n').encode())
        process.stdin.flush()
        if identity is None:
            return None
        deadline = time.monotonic() + 30
        while True:
            while b"\n" not in pending:
                remaining = deadline - time.monotonic()
                if remaining <= 0 or not selector.select(remaining):
                    raise RuntimeError('Timed out; inspect .chio/server.log')
                chunk = os.read(process.stdout.fileno(), 65536)
                if not chunk:
                    raise RuntimeError('Chio stopped; inspect .chio/server.log')
                pending.extend(chunk)
                if len(pending) > 4 * 1024 * 1024:
                    raise RuntimeError('Response exceeds example buffer bound')
            line, _, rest = pending.partition(b"\n")
            pending[:] = rest
            response = json.loads(line)
            if response.get('id') == identity:
                return response
    try:
        initialized = send('initialize', {'protocolVersion': '2025-11-25', 'capabilities': {},
                           'clientInfo': {'name': 'security-docs', 'version': '1.0.0'}}, 1)
        assert 'result' in initialized, initialized
        send('notifications/initialized', {})
        tools = send('tools/list', {}, 2)
        assert 'read_file' in {t['name'] for t in tools['result']['tools']}, tools
        allowed = send('tools/call', {'name': 'read_file', 'arguments': {'path': 'src/pricing.py'}}, 3)
        assert 'error' not in allowed and not allowed['result'].get('isError'), allowed
        assert original.decode() in json.dumps(allowed, ensure_ascii=False).replace('\\n', '\n'), allowed
        print('ALLOW read_file: src/pricing.py returned')
        denied = send('tools/call', {'name': 'write_file', 'arguments': {'path': 'src/pricing.py', 'text': 'TIER_PERCENT = 80'}}, 4)
        assert 'error' in denied or denied.get('result', {}).get('isError'), denied
        assert TARGET.read_bytes() == original, 'Denied write changed src/pricing.py'
        print('DENY write_file: src/pricing.py unchanged')
        (STATE / 'responses.json').write_text(json.dumps({'tools': tools, 'allow': allowed, 'deny': denied}, indent=2) + '\n')
    finally:
        selector.close()
        process.terminate()
        try:
            process.wait(timeout=5)
        except subprocess.TimeoutExpired:
            process.kill()
            process.wait()
print('Evidence: .chio/receipts.db and .chio/responses.json')