"""Run Beam as an operator of the local experimental monitor via Modal.

The HF credential remains in a Modal secret; a small remote function proxies
chat requests. CPU/monitor execution happens locally. Each invocation is billed
by Modal and the selected HF inference provider.

Run: OMP_NUM_THREADS=1 .venv/bin/modal run scripts/modal_beam_operator.py
"""
import json
import os
import urllib.request
import urllib.error
import modal

app = modal.App('beamboot-operator')


@app.function(image=modal.Image.debian_slim(python_version='3.12'),
              secrets=[modal.Secret.from_name('huggingface-secret')], timeout=90)
def beam_chat(messages, tools):
    token = (os.environ.get('HF_TOKEN') or os.environ.get('HUGGING_FACE_HUB_TOKEN')
             or os.environ.get('HUGGINGFACE_TOKEN') or os.environ.get('HUGGINGFACEHUB_API_TOKEN'))
    if not token:
        # Do not print the names of unrelated environment variables or secrets.
        raise RuntimeError('Hugging Face token was not found under a supported environment variable name in the Modal secret')
    payload = {'model': 'reflection/Beam-1.0-501B-A23B:baseten', 'messages': messages,
               'tools': tools, 'tool_choice': 'auto', 'max_tokens': 2048}
    req = urllib.request.Request('https://router.huggingface.co/v1/chat/completions',
        data=json.dumps(payload).encode(),
        headers={'Authorization': 'Bearer ' + token, 'Content-Type': 'application/json'})
    try:
        with urllib.request.urlopen(req, timeout=75) as response:
            return json.load(response)
    except urllib.error.HTTPError as error:
        raise RuntimeError(f'Hugging Face router returned HTTP {error.code}: {error.read(1000).decode(errors="replace")}') from None


@app.local_entrypoint()
def main():
    from beamboot.os import OperatingMonitor
    monitor = OperatingMonitor()
    boot = monitor.boot()
    print('Machine boot:', repr(boot))
    tool = {'type': 'function', 'function': {'name': 'monitor_command',
        'description': 'Operate the booted four-bit monitor: help, version, or add. Addition returns the sum modulo 16.',
        'parameters': {'type': 'object', 'properties': {
            'command': {'type': 'string', 'enum': ['help', 'version', 'add']},
            'left': {'type': 'integer', 'minimum': 0, 'maximum': 15},
            'right': {'type': 'integer', 'minimum': 0, 'maximum': 15}},
            'required': ['command'], 'additionalProperties': False}}}
    messages = [{'role': 'system', 'content': 'You are operating an experimental four-bit computer. Use the monitor_command tool. Do not invent its responses.'},
                {'role': 'user', 'content': f'The machine booted with {boot!r}. Use the tool to get the version and calculate 7 + 8, then report the actual results.'}]
    for turn in range(5):
        response = beam_chat.remote(messages, [tool])
        choice = response['choices'][0]
        message = choice['message']
        # Preserve reasoning_content if supplied by the provider. Provider-specific
        # fields that are not part of the message contract are intentionally omitted.
        assistant = {k: message[k] for k in ('role', 'content', 'reasoning_content', 'tool_calls') if k in message}
        messages.append(assistant)
        calls = message.get('tool_calls') or []
        if not calls:
            print('Beam:', message.get('content'))
            return
        for call in calls:
            try:
                if call['function']['name'] != 'monitor_command': raise ValueError('unknown tool')
                args = json.loads(call['function']['arguments'])
                name = args['command']
                if name == 'help': result = monitor.help()
                elif name == 'version': result = monitor.version()
                elif name == 'add': result = str(monitor.add(args['left'], args['right']))
                else: raise ValueError('unsupported command')
            except (KeyError, ValueError, RuntimeError, TypeError) as error:
                result = json.dumps({'error': str(error)})
            print('Tool:', call['function']['arguments'], '->', repr(result))
            messages.append({'role': 'tool', 'tool_call_id': call['id'], 'content': result})
    raise RuntimeError('Beam did not finish within the five-call budget')