#!/usr/bin/env python3
"""EMT Compute MCP stdio server. No dependencies, network listener or credential storage."""
import json
import sys
from urllib.parse import quote
from emt_compute import EMTCompute

VERSION = '2025-11-25'
STRING = {'type': 'string', 'minLength': 1, 'maxLength': 500}
UUID = {**STRING, 'maxLength': 36}
HOURS = {'type': 'integer', 'minimum': 1, 'maximum': 720}
KEY = {'type': 'string', 'minLength': 8, 'maxLength': 128}

def tool(name, description, properties=None, required=None, read_only=True, destructive=False):
    return {'name': name, 'description': description,
            'inputSchema': {'type': 'object', 'properties': properties or {}, 'required': required or [], 'additionalProperties': False},
            'annotations': {'readOnlyHint': read_only, 'destructiveHint': destructive, 'openWorldHint': True}}

TOOLS = [
    tool('emt_catalog', 'Browse currently offered GPU nodes and EMT prices.'),
    tool('emt_clusters', 'Browse available whole-node clusters and interconnects.'),
    tool('emt_templates', 'List public container templates.'),
    tool('emt_orders', 'List the authenticated customer’s GPU orders.'),
    tool('emt_monitor', 'Read live node metrics and lifecycle events. Missing metrics are unknown, not zero.', {'order': UUID, 'member': UUID}, ['order']),
    tool('emt_logs', 'Read bounded container logs for one owned node.', {'order': UUID, 'member': UUID}, ['order']),
    tool('emt_volumes', 'List owned persistent volumes.'),
    tool('emt_storage_costs', 'Read metered storage and transfer charges.'),
    tool('emt_ledger', 'Read the customer’s billing ledger.'),
    tool('emt_backups', 'List backups for one node, including retained backups.', {'order': UUID, 'member': UUID, 'page': {'type': 'integer', 'minimum': 1, 'maximum': 10000}}, ['order']),
    tool('emt_quote', 'Create an expiring quote; does not start a GPU or charge credit.', {'node': UUID, 'count': {'type': 'integer', 'minimum': 1, 'maximum': 256}, 'hours': HOURS}, ['node', 'count', 'hours'], False),
    tool('emt_cluster_quote', 'Quote selected whole nodes from one available cluster.', {'cluster': STRING, 'nodes': {'type': 'array', 'items': UUID, 'minItems': 2, 'maxItems': 16}, 'hours': HOURS}, ['cluster', 'nodes', 'hours'], False),
    tool('emt_rent', 'Start a paid GPU rental from a fresh quote. Persist the idempotency key and reuse it after an uncertain result. Use a public SSH key only.', {'quote': UUID, 'name': STRING, 'template': UUID, 'ssh_public_key': {**STRING, 'maxLength': 10000}, 'idempotency_key': KEY, 'jupyter': {'type': 'boolean'}, 'volume': UUID, 'backups': {'type': 'boolean'}, 'ports': {'type': 'integer', 'minimum': 0, 'maximum': 200}}, ['quote', 'name', 'template', 'ssh_public_key', 'idempotency_key'], False),
    tool('emt_rent_cluster', 'Start a paid multi-node cluster from a fresh quote. Reuse the same persisted idempotency key on retry.', {'quote': UUID, 'name': STRING, 'ssh_public_key': {**STRING, 'maxLength': 10000}, 'idempotency_key': KEY}, ['quote', 'name', 'ssh_public_key', 'idempotency_key'], False),
    tool('emt_create_volume', 'Create a persistent volume with metered storage and transfer charges.', {'name': STRING, 'description': STRING, 'idempotency_key': KEY}, ['name', 'idempotency_key'], False),
    tool('emt_backup', 'Request a metered backup after configuring its path. Reuse the same persisted idempotency key on retry.', {'order': UUID, 'member': UUID, 'name': STRING, 'idempotency_key': KEY}, ['order', 'name', 'idempotency_key'], False),
    tool('emt_safe_stop', 'Request termination only after all node backups satisfy the qualified verification policy. Close and flush checkpoints first. Persistent volumes and backups remain billed.', {'order': UUID, 'path': STRING, 'frequency': {'type': 'integer', 'minimum': 1, 'maximum': 168}, 'retention': {'type': 'integer', 'minimum': 1, 'maximum': 365}, 'checkpoint_closed': {'type': 'boolean', 'const': True}}, ['order', 'path', 'frequency', 'retention', 'checkpoint_closed'], False, True),
    tool('emt_terminate', 'Terminate all nodes in an order. Local disks are deleted. Persistent volumes and backups remain retained and billed.', {'order': UUID, 'confirm_local_data_loss': {'type': 'boolean', 'const': True}}, ['order', 'confirm_local_data_loss'], False, True),
]

def validate(value, schema):
    kind = schema['type']
    if kind == 'object':
        if not isinstance(value, dict): raise ValueError('Expected object')
        if set(value) - set(schema['properties']): raise ValueError('Unknown argument')
        if set(schema.get('required', [])) - set(value): raise ValueError('Missing required argument')
        for key, item in value.items(): validate(item, schema['properties'][key])
    elif kind == 'string':
        if not isinstance(value, str) or not schema.get('minLength', 0) <= len(value) <= schema.get('maxLength', 10000): raise ValueError('Invalid string')
    elif kind == 'integer':
        if type(value) is not int or not schema.get('minimum', -10**9) <= value <= schema.get('maximum', 10**9): raise ValueError('Invalid integer')
    elif kind == 'boolean':
        if type(value) is not bool or ('const' in schema and value is not schema['const']): raise ValueError('Explicit acknowledgement required')
    elif kind == 'array':
        if not isinstance(value, list) or not schema.get('minItems', 0) <= len(value) <= schema.get('maxItems', 1000): raise ValueError('Invalid list')
        for item in value: validate(item, schema['items'])

class Server:
    def __init__(self, client=None):
        self.client = client or EMTCompute()
        self.initialized = False
        self.negotiated = False
    def dispatch(self, name, a):
        c = self.client
        simple = {'emt_catalog': 'catalog', 'emt_clusters': 'catalog/clusters', 'emt_templates': 'templates', 'emt_orders': 'orders', 'emt_volumes': 'volumes', 'emt_storage_costs': 'storage/costs', 'emt_ledger': 'ledger'}
        if name in simple: return c.request(simple[name])
        if name == 'emt_monitor': return c.monitor(a['order'], a.get('member'))
        if name == 'emt_logs': return c.request('orders/'+quote(a['order'], safe='')+('/members/'+quote(a['member'], safe='') if a.get('member') else '')+'/logs')
        if name == 'emt_backups': return c.backups(a['order'], a.get('member'), a.get('page', 1))
        if name == 'emt_quote': return c.quote(a['node'], a['count'], a['hours'])
        if name == 'emt_cluster_quote': return c.cluster_quote(a['cluster'], a['nodes'], a['hours'])
        if name == 'emt_rent': return c.rent(a['quote'], a['name'], a['template'], a['ssh_public_key'], a['idempotency_key'], a.get('jupyter', False), a.get('volume'), a.get('backups', False), a.get('ports', 10))
        if name == 'emt_rent_cluster': return c.rent_cluster(a['quote'], a['name'], a['ssh_public_key'], a['idempotency_key'])
        if name == 'emt_create_volume': return c.create_volume(a['name'], a.get('description', ''), a['idempotency_key'])
        if name == 'emt_backup': return c.backup(a['order'], a['name'], a['idempotency_key'], a.get('member'))
        if name == 'emt_safe_stop': return c.safe_stop(a['order'], a['path'], a['frequency'], a['retention'], a['checkpoint_closed'])
        if name == 'emt_terminate': return c.remove(a['order'])
        raise ValueError('Unknown tool')
    def handle(self, request):
        if not isinstance(request, dict) or request.get('jsonrpc') != '2.0' or not isinstance(request.get('method'), str):
            return {'jsonrpc': '2.0', 'id': None, 'error': {'code': -32600, 'message': 'Invalid Request'}}
        method = request['method']; rid = request.get('id'); params = request.get('params') or {}
        if 'id' not in request:
            if method == 'notifications/initialized' and self.negotiated: self.initialized = True
            return None
        def answer(result): return {'jsonrpc': '2.0', 'id': rid, 'result': result}
        def error(code, message): return {'jsonrpc': '2.0', 'id': rid, 'error': {'code': code, 'message': message}}
        if not isinstance(params, dict): return error(-32602, 'Invalid params')
        if method == 'initialize':
            self.negotiated = True
            version = params.get('protocolVersion')
            return answer({'protocolVersion': version if version in [VERSION, '2025-06-18', '2024-11-05'] else VERSION, 'capabilities': {'tools': {'listChanged': False}}, 'serverInfo': {'name': 'emt-compute', 'version': '0.3.0'}})
        if method == 'ping': return answer({})
        if not self.initialized: return error(-32002, 'Initialize the session first')
        if method == 'tools/list': return answer({'tools': TOOLS})
        if method != 'tools/call': return error(-32601, 'Method not found')
        found = next((t for t in TOOLS if t['name'] == params.get('name')), None)
        if not found: return error(-32602, 'Unknown tool')
        try: validate(params.get('arguments', {}), found['inputSchema'])
        except ValueError as e: return error(-32602, str(e))
        try:
            result = self.dispatch(found['name'], params.get('arguments', {}))
            return answer({'content': [{'type': 'text', 'text': json.dumps(result, ensure_ascii=False)}], 'isError': False})
        except Exception as e:
            return answer({'content': [{'type': 'text', 'text': str(e)[:2000]}], 'isError': True})

def main():
    server = Server()
    for line in sys.stdin:
        if len(line) > 1_048_576:
            result = {'jsonrpc': '2.0', 'id': None, 'error': {'code': -32600, 'message': 'Request too large'}}
        else:
            try: result = server.handle(json.loads(line))
            except (ValueError, TypeError): result = {'jsonrpc': '2.0', 'id': None, 'error': {'code': -32700, 'message': 'Parse error'}}
        if result is not None:
            print(json.dumps(result, ensure_ascii=False), flush=True)

if __name__ == '__main__': main()