#!/usr/bin/env python3
"""
Distribute MCP Server
=====================

Exposes the Distribute Public API as Model Context Protocol tools so
Claude Desktop, Cursor, and any other MCP-compatible client can manage
apps, deployments, environment variables, and billing programmatically.

Configuration
-------------
Set two environment variables before launching:

  DISTRIBUTE_API_TOKEN  · a token created in your dashboard (Settings → API tokens)
  DISTRIBUTE_BASE_URL   · default https://distribute.app

Usage with Claude Desktop
-------------------------
Add the following to ~/Library/Application Support/Claude/claude_desktop_config.json:

  {
    "mcpServers": {
      "distribute": {
        "command": "python3",
        "args": ["/path/to/mcp_server.py"],
        "env": {
          "DISTRIBUTE_API_TOKEN": "your-token-here",
          "DISTRIBUTE_BASE_URL":  "https://distribute.app"
        }
      }
    }
  }

After restart, Claude can call: list_apps, get_app, restart_app, get_logs,
get_env, set_env, get_balance, list_domains, list_workspaces.

Wire format
-----------
This server speaks the JSON-RPC 2.0 stdio variant of MCP. Most production
MCP servers use the `mcp` Python SDK, but to keep this file zero-dependency
(important when bundled with the Distribute platform · operators don't
need to pip install anything), we implement the protocol directly.
"""
import json
import os
import sys
import urllib.error
import urllib.request


# ── Config ────────────────────────────────────────────────────────────
TOKEN    = os.getenv('DISTRIBUTE_API_TOKEN', '').strip()
BASE_URL = os.getenv('DISTRIBUTE_BASE_URL', 'https://distribute.app').rstrip('/')
TIMEOUT  = int(os.getenv('DISTRIBUTE_TIMEOUT', '30'))


def _err(msg):
    """Stderr-log without disturbing the JSON-RPC stdout stream."""
    print(f'[mcp-distribute] {msg}', file=sys.stderr, flush=True)


if not TOKEN:
    _err('FATAL · DISTRIBUTE_API_TOKEN env var is required')
    sys.exit(1)


# ── HTTP helper ────────────────────────────────────────────────────────
def api_call(method, path, body=None, follow_redirects=True):
    """Make an authenticated request to the Distribute API.

    Returns the parsed JSON response. Raises RuntimeError on HTTP error.
    Follows 307s automatically (the v1 namespace uses 307 redirects to
    the canonical implementation endpoints, preserving method + body).
    """
    url = BASE_URL + path
    data = json.dumps(body).encode() if body is not None else None
    headers = {
        'Authorization': f'Bearer {TOKEN}',
        'User-Agent':    'Distribute-MCP-Server/1.0',
        'Accept':        'application/json',
    }
    if data is not None:
        headers['Content-Type'] = 'application/json'
    req = urllib.request.Request(url, data=data, method=method, headers=headers)
    try:
        with urllib.request.urlopen(req, timeout=TIMEOUT) as resp:
            raw = resp.read().decode('utf-8', errors='replace')
            try:
                return json.loads(raw)
            except json.JSONDecodeError:
                return {'raw': raw[:500]}
    except urllib.error.HTTPError as e:
        body = e.read().decode('utf-8', errors='replace')[:500]
        try:
            err = json.loads(body)
        except Exception:
            err = {'error': body}
        raise RuntimeError(f'HTTP {e.code}: {err.get("error", err)}') from e
    except urllib.error.URLError as e:
        raise RuntimeError(f'Network error: {e.reason}') from e


# ── Tool implementations · thin wrappers around api_call ──────────────
def tool_list_apps(args):
    return api_call('GET', '/api/v1/apps')


def tool_get_app(args):
    name = args.get('name')
    if not name:
        raise ValueError('name is required')
    return api_call('GET', f'/api/v1/apps/{name}')


def tool_get_logs(args):
    name = args.get('name')
    if not name:
        raise ValueError('name is required')
    tail = int(args.get('tail', 100))
    return api_call('GET', f'/api/v1/apps/{name}/logs?tail={tail}')


def tool_restart_app(args):
    name = args.get('name')
    if not name:
        raise ValueError('name is required')
    return api_call('POST', f'/api/v1/apps/{name}/restart')


def tool_stop_app(args):
    name = args.get('name')
    if not name:
        raise ValueError('name is required')
    return api_call('POST', f'/api/v1/apps/{name}/stop')


def tool_start_app(args):
    name = args.get('name')
    if not name:
        raise ValueError('name is required')
    return api_call('POST', f'/api/v1/apps/{name}/start')


def tool_delete_app(args):
    name = args.get('name')
    if not name:
        raise ValueError('name is required')
    confirm = args.get('confirm')
    if confirm != name:
        raise ValueError(
            f'Pass confirm="{name}" to delete this app · '
            f'this is destructive and cannot be undone')
    return api_call('DELETE', f'/api/v1/apps/{name}')


def tool_get_env(args):
    name = args.get('name')
    if not name:
        raise ValueError('name is required')
    return api_call('GET', f'/api/v1/apps/{name}/env')


def tool_set_env(args):
    """Set environment variables on an app. Pass {name, env: {KEY: VALUE, ...}}.
    Existing keys not in the payload are preserved · pass null to delete."""
    name = args.get('name')
    env  = args.get('env')
    if not name:
        raise ValueError('name is required')
    if not isinstance(env, dict):
        raise ValueError('env must be an object {KEY: VALUE}')
    return api_call('POST', f'/api/v1/apps/{name}/env', body={'env': env})


def tool_get_balance(args):
    return api_call('GET', '/api/v1/billing/balance')


def tool_list_domains(args):
    return api_call('GET', '/api/v1/domains')


def tool_list_workspaces(args):
    return api_call('GET', '/api/v1/workspaces')


def tool_whoami(args):
    return api_call('GET', '/api/v1/me')


def tool_deploy_app(args):
    """Deploy a new app from a publicly-accessible zip URL.

    The MCP server lives on the user's machine but the deploy happens
    on Distribute's servers · so we can't upload local files directly
    over JSON-RPC stdio (the protocol doesn't carry binary bodies well).
    The pattern: the user's tool produces a zip and uploads it somewhere
    Distribute can fetch it (GitHub release, S3 presigned URL, transfer.sh,
    etc.), then passes that URL here.

    Required args: app_name, source_url
    Optional:      port, custom_domain, env (object), instance_plan,
                   description, tags
    """
    if not args.get('app_name'):
        raise ValueError('app_name is required')
    if not args.get('source_url'):
        raise ValueError(
            'source_url is required · must be a public http(s) URL '
            'to a zip archive of your app source')
    body = {
        'app_name':   args['app_name'],
        'source_url': args['source_url'],
    }
    for k in ('port', 'custom_domain', 'instance_plan', 'description', 'tags'):
        if k in args:
            body[k] = args[k]
    if 'env' in args:
        if not isinstance(args['env'], dict):
            raise ValueError('env must be an object {KEY: VALUE}')
        body['env'] = args['env']
    return api_call('POST', '/api/v1/apps/deploy', body=body)


# ── Tool registry · maps name → (handler, schema) ──────────────────────
TOOLS = {
    'whoami': {
        'description': 'Verify the API token works and identify the user',
        'handler':     tool_whoami,
        'inputSchema': {'type': 'object', 'properties': {}},
    },
    'list_apps': {
        'description': 'List all apps deployed by the current user',
        'handler':     tool_list_apps,
        'inputSchema': {'type': 'object', 'properties': {}},
    },
    'deploy_app': {
        'description': (
            'Deploy a new app to Distribute from a public zip URL. '
            'Returns the live HTTPS URL once the build completes. '
            'The source_url must be a public http(s) URL to a zip archive '
            '(e.g. GitHub release asset, S3 presigned URL, transfer.sh). '
            'Use this when the user asks to deploy code · prompt them for '
            'a URL if their code is not yet hosted somewhere fetchable.'
        ),
        'handler':     tool_deploy_app,
        'inputSchema': {
            'type': 'object',
            'properties': {
                'app_name':      {'type': 'string',
                                  'description': 'Lowercase name, 1-50 chars, [a-z0-9_-]'},
                'source_url':    {'type': 'string',
                                  'description': 'Public URL to a zip of your app source'},
                'port':          {'type': 'integer',
                                  'description': 'Port the app listens on inside the container · auto-detected if omitted'},
                'custom_domain': {'type': 'string',
                                  'description': 'Custom domain like myapp.com · DNS CNAME setup required'},
                'env':           {'type': 'object',
                                  'description': 'Environment variables · {KEY: VALUE} pairs'},
                'instance_plan': {'type': 'string',
                                  'description': 'nano (free), micro, small, medium, large',
                                  'default': 'nano'},
                'description':   {'type': 'string'},
                'tags':          {'type': 'string', 'description': 'Comma-separated tags'},
            },
            'required': ['app_name', 'source_url'],
        },
    },
    'get_app': {
        'description': 'Get full metadata for one app · status, URL, plan, custom domains',
        'handler':     tool_get_app,
        'inputSchema': {
            'type': 'object',
            'properties': {'name': {'type': 'string', 'description': 'App name'}},
            'required': ['name'],
        },
    },
    'get_logs': {
        'description': 'Read recent stdout/stderr from an app (default last 100 lines)',
        'handler':     tool_get_logs,
        'inputSchema': {
            'type': 'object',
            'properties': {
                'name': {'type': 'string', 'description': 'App name'},
                'tail': {'type': 'integer', 'description': 'Lines to return', 'default': 100},
            },
            'required': ['name'],
        },
    },
    'restart_app': {
        'description': 'Restart the app container · short downtime, no data loss',
        'handler':     tool_restart_app,
        'inputSchema': {
            'type': 'object',
            'properties': {'name': {'type': 'string'}},
            'required': ['name'],
        },
    },
    'stop_app': {
        'description': 'Stop the app · billing for compute pauses, data preserved',
        'handler':     tool_stop_app,
        'inputSchema': {
            'type': 'object',
            'properties': {'name': {'type': 'string'}},
            'required': ['name'],
        },
    },
    'start_app': {
        'description': 'Start a stopped app',
        'handler':     tool_start_app,
        'inputSchema': {
            'type': 'object',
            'properties': {'name': {'type': 'string'}},
            'required': ['name'],
        },
    },
    'delete_app': {
        'description': ('Permanently delete an app · destructive, cannot be undone. '
                        'Pass confirm equal to the app name as a safety check.'),
        'handler':     tool_delete_app,
        'inputSchema': {
            'type': 'object',
            'properties': {
                'name':    {'type': 'string'},
                'confirm': {'type': 'string', 'description': 'Must equal the app name'},
            },
            'required': ['name', 'confirm'],
        },
    },
    'get_env': {
        'description': 'Read environment variables of an app',
        'handler':     tool_get_env,
        'inputSchema': {
            'type': 'object',
            'properties': {'name': {'type': 'string'}},
            'required': ['name'],
        },
    },
    'set_env': {
        'description': 'Set environment variables on an app. Pass null as value to delete a key.',
        'handler':     tool_set_env,
        'inputSchema': {
            'type': 'object',
            'properties': {
                'name': {'type': 'string'},
                'env':  {'type': 'object',
                         'description': 'Object of KEY: VALUE pairs to set. Null deletes.'},
            },
            'required': ['name', 'env'],
        },
    },
    'get_balance': {
        'description': 'Read your current wallet balance and gift credit',
        'handler':     tool_get_balance,
        'inputSchema': {'type': 'object', 'properties': {}},
    },
    'list_domains': {
        'description': 'List all custom domains attached to your apps',
        'handler':     tool_list_domains,
        'inputSchema': {'type': 'object', 'properties': {}},
    },
    'list_workspaces': {
        'description': 'List your workspaces (team membership and access)',
        'handler':     tool_list_workspaces,
        'inputSchema': {'type': 'object', 'properties': {}},
    },
}


# ── JSON-RPC 2.0 / MCP wire protocol ──────────────────────────────────
def jr_response(req_id, result=None, error=None):
    """Build a JSON-RPC 2.0 response object."""
    out = {'jsonrpc': '2.0', 'id': req_id}
    if error is not None:
        out['error'] = error
    else:
        out['result'] = result
    return out


def handle_request(req):
    """Dispatch one JSON-RPC request to the appropriate MCP handler."""
    method = req.get('method', '')
    req_id = req.get('id')
    params = req.get('params') or {}

    if method == 'initialize':
        # Standard MCP handshake · advertise capabilities + server info
        return jr_response(req_id, {
            'protocolVersion': '2024-11-05',
            'capabilities':    {'tools': {'listChanged': False}},
            'serverInfo': {
                'name':    'distribute',
                'version': '1.0',
            },
        })

    if method == 'notifications/initialized':
        # Client telling us setup is complete · no response expected
        # (this is a notification, not a request · req_id will be None)
        return None

    if method == 'tools/list':
        # Return all registered tools with their schemas
        tools = [{
            'name':        name,
            'description': spec['description'],
            'inputSchema': spec['inputSchema'],
        } for name, spec in TOOLS.items()]
        return jr_response(req_id, {'tools': tools})

    if method == 'tools/call':
        name = params.get('name', '')
        args = params.get('arguments') or {}
        spec = TOOLS.get(name)
        if not spec:
            return jr_response(req_id, error={
                'code':    -32602,
                'message': f'Unknown tool: {name}',
            })
        try:
            result = spec['handler'](args)
            # MCP expects the tool result wrapped in a content array
            return jr_response(req_id, {
                'content': [{'type': 'text', 'text': json.dumps(result, ensure_ascii=False, indent=2)}],
                'isError': False,
            })
        except Exception as e:
            _err(f'tool {name} failed: {e}')
            return jr_response(req_id, {
                'content': [{'type': 'text', 'text': f'Error: {e}'}],
                'isError': True,
            })

    if method == 'ping':
        return jr_response(req_id, {})

    return jr_response(req_id, error={
        'code':    -32601,
        'message': f'Method not found: {method}',
    })


# ── Main loop · stdin/stdout JSON-RPC ──────────────────────────────────
def main():
    _err(f'starting · base={BASE_URL}, token=...{TOKEN[-4:] if TOKEN else "?"}')
    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            req = json.loads(line)
        except json.JSONDecodeError as e:
            _err(f'bad JSON: {e}')
            continue
        try:
            resp = handle_request(req)
            if resp is not None:
                sys.stdout.write(json.dumps(resp) + '\n')
                sys.stdout.flush()
        except Exception as e:
            _err(f'dispatch error: {e}')
            sys.stdout.write(json.dumps(jr_response(
                req.get('id'), error={'code': -32603, 'message': str(e)},
            )) + '\n')
            sys.stdout.flush()


if __name__ == '__main__':
    main()
