#!/usr/bin/env python3
"""
TENGWAR MCP Server

Model Context Protocol server that exposes Tengwar as a tool
to any MCP-compatible AI agent (Claude, etc).

Tools exposed:
  - tengwar_eval: Execute Tengwar source code
  - tengwar_eval_binary: Execute base64-encoded binary AST  
  - tengwar_sandbox: Execute in sandboxed mode (safe for untrusted code)

Resources exposed:
  - tengwar://builtins: List of all available builtins
  - tengwar://syntax: Quick syntax reference

Usage:
  python -m tengwar.mcp_server           # stdio mode
  python -m tengwar.mcp_server --port 8080  # HTTP mode (if supported)

MCP config (claude_desktop_config.json):
  {
    "mcpServers": {
      "tengwar": {
        "command": "python",
        "args": ["-m", "tengwar.mcp_server"]
      }
    }
  }
"""

import sys
import json
import traceback
from typing import Any

# Import Tengwar
from .interpreter import Interpreter, TengwarValue
from .binary_ast import run_b64, ASTBuilder


# === MCP Protocol Implementation (stdio JSON-RPC) ===

class TengwarMCPServer:
    """MCP server exposing Tengwar as tools for AI agents."""
    
    def __init__(self):
        self.interpreter = Interpreter()
        self.sandbox_interpreter = Interpreter(sandbox=True)
        self.sandbox_interpreter.max_steps = 1_000_000
        
    def handle_request(self, request: dict) -> dict:
        """Handle a JSON-RPC request."""
        method = request.get("method", "")
        req_id = request.get("id")
        params = request.get("params", {})
        
        try:
            if method == "initialize":
                return self._initialize(req_id, params)
            elif method == "tools/list":
                return self._list_tools(req_id)
            elif method == "tools/call":
                return self._call_tool(req_id, params)
            elif method == "resources/list":
                return self._list_resources(req_id)
            elif method == "resources/read":
                return self._read_resource(req_id, params)
            elif method == "ping":
                return self._result(req_id, {})
            elif method == "notifications/initialized":
                return None  # No response needed
            else:
                return self._error(req_id, -32601, f"Method not found: {method}")
        except Exception as e:
            return self._error(req_id, -32603, str(e))
    
    def _result(self, req_id, result):
        return {"jsonrpc": "2.0", "id": req_id, "result": result}
    
    def _error(self, req_id, code, message):
        return {"jsonrpc": "2.0", "id": req_id, "error": {"code": code, "message": message}}
    
    def _initialize(self, req_id, params):
        return self._result(req_id, {
            "protocolVersion": "2024-11-05",
            "capabilities": {
                "tools": {},
                "resources": {}
            },
            "serverInfo": {
                "name": "tengwar",
                "version": "0.3.0"
            }
        })
    
    def _list_tools(self, req_id):
        return self._result(req_id, {
            "tools": [
                {
                    "name": "tengwar_eval",
                    "description": (
                        "Execute Tengwar code. Tengwar is a functional language optimized for AI — "
                        "23% fewer tokens than Python for data operations, with 80+ builtins including "
                        "map, filter, reduce, scan, chunks, partition, unique, flat-map, pattern matching, "
                        "dictionaries, JSON, and Python interop. "
                        "Example: (sum (for x (range 1 11) when (even? x) (* x 2))) → 60"
                    ),
                    "inputSchema": {
                        "type": "object",
                        "properties": {
                            "code": {
                                "type": "string",
                                "description": "Tengwar source code to execute"
                            }
                        },
                        "required": ["code"]
                    }
                },
                {
                    "name": "tengwar_sandbox",
                    "description": (
                        "Execute Tengwar code in a secure sandbox. No file I/O, no network, "
                        "no Python imports, with a 1M step limit. Safe for untrusted code. "
                        "Pure computation only."
                    ),
                    "inputSchema": {
                        "type": "object",
                        "properties": {
                            "code": {
                                "type": "string",
                                "description": "Tengwar source code to execute (sandboxed)"
                            },
                            "max_steps": {
                                "type": "integer",
                                "description": "Maximum execution steps (default: 1000000)",
                                "default": 1000000
                            }
                        },
                        "required": ["code"]
                    }
                },
                {
                    "name": "tengwar_eval_binary",
                    "description": (
                        "Execute a base64-encoded Tengwar binary AST. "
                        "Zero syntax errors possible — the binary format maps directly to AST nodes. "
                        "Use tengwar_eval to generate base64 with: (encode-b64 \"(+ 1 2)\")"
                    ),
                    "inputSchema": {
                        "type": "object",
                        "properties": {
                            "b64": {
                                "type": "string",
                                "description": "Base64-encoded binary AST"
                            }
                        },
                        "required": ["b64"]
                    }
                },
                {
                    "name": "tengwar_multi",
                    "description": (
                        "Execute multiple Tengwar expressions sharing the same environment. "
                        "Define functions in earlier expressions, use them in later ones. "
                        "Returns the result of the last expression."
                    ),
                    "inputSchema": {
                        "type": "object",
                        "properties": {
                            "expressions": {
                                "type": "array",
                                "items": {"type": "string"},
                                "description": "Array of Tengwar expressions to execute in sequence"
                            }
                        },
                        "required": ["expressions"]
                    }
                }
            ]
        })
    
    def _call_tool(self, req_id, params):
        tool_name = params.get("name", "")
        args = params.get("arguments", {})
        
        if tool_name == "tengwar_eval":
            return self._eval(req_id, args.get("code", ""), sandbox=False)
        elif tool_name == "tengwar_sandbox":
            max_steps = args.get("max_steps", 1_000_000)
            self.sandbox_interpreter.max_steps = max_steps
            self.sandbox_interpreter.step_count = 0
            return self._eval(req_id, args.get("code", ""), sandbox=True)
        elif tool_name == "tengwar_eval_binary":
            return self._eval_binary(req_id, args.get("b64", ""))
        elif tool_name == "tengwar_multi":
            return self._eval_multi(req_id, args.get("expressions", []))
        else:
            return self._error(req_id, -32602, f"Unknown tool: {tool_name}")
    
    def _eval(self, req_id, code: str, sandbox: bool = False):
        try:
            interp = self.sandbox_interpreter if sandbox else self.interpreter
            result = interp.run_source(code)
            return self._result(req_id, {
                "content": [{
                    "type": "text",
                    "text": repr(result)
                }]
            })
        except Exception as e:
            return self._result(req_id, {
                "content": [{
                    "type": "text",
                    "text": f"Error: {e}"
                }],
                "isError": True
            })
    
    def _eval_binary(self, req_id, b64: str):
        try:
            result = run_b64(b64, self.interpreter)
            return self._result(req_id, {
                "content": [{
                    "type": "text",
                    "text": repr(result)
                }]
            })
        except Exception as e:
            return self._result(req_id, {
                "content": [{
                    "type": "text",
                    "text": f"Error: {e}"
                }],
                "isError": True
            })
    
    def _eval_multi(self, req_id, expressions: list):
        try:
            result = None
            for expr in expressions:
                result = self.interpreter.run_source(expr)
            return self._result(req_id, {
                "content": [{
                    "type": "text",
                    "text": repr(result)
                }]
            })
        except Exception as e:
            return self._result(req_id, {
                "content": [{
                    "type": "text",
                    "text": f"Error: {e}"
                }],
                "isError": True
            })
    
    def _list_resources(self, req_id):
        return self._result(req_id, {
            "resources": [
                {
                    "uri": "tengwar://builtins",
                    "name": "Tengwar Builtins",
                    "description": "Complete list of all 80+ built-in functions",
                    "mimeType": "text/plain"
                },
                {
                    "uri": "tengwar://syntax",
                    "name": "Tengwar Syntax Reference",
                    "description": "Quick syntax guide for writing Tengwar code",
                    "mimeType": "text/plain"
                }
            ]
        })
    
    def _read_resource(self, req_id, params):
        uri = params.get("uri", "")
        
        if uri == "tengwar://builtins":
            return self._result(req_id, {
                "contents": [{
                    "uri": uri,
                    "mimeType": "text/plain",
                    "text": BUILTINS_REFERENCE
                }]
            })
        elif uri == "tengwar://syntax":
            return self._result(req_id, {
                "contents": [{
                    "uri": uri,
                    "mimeType": "text/plain",
                    "text": SYNTAX_REFERENCE
                }]
            })
        else:
            return self._error(req_id, -32602, f"Unknown resource: {uri}")
    
    def run_stdio(self):
        """Run as stdio MCP server."""
        for line in sys.stdin:
            line = line.strip()
            if not line:
                continue
            try:
                request = json.loads(line)
                response = self.handle_request(request)
                if response is not None:
                    sys.stdout.write(json.dumps(response) + "\n")
                    sys.stdout.flush()
            except json.JSONDecodeError:
                error = {
                    "jsonrpc": "2.0",
                    "id": None,
                    "error": {"code": -32700, "message": "Parse error"}
                }
                sys.stdout.write(json.dumps(error) + "\n")
                sys.stdout.flush()


# === Reference Content ===

BUILTINS_REFERENCE = """TENGWAR BUILTINS (v0.3)

COLLECTIONS:
  (map f vec)              Apply f to each element
  (filter f vec)           Keep elements where f returns true
  (reduce f init vec)      Fold left with accumulator
  (flat-map f vec)         Map then flatten one level
  (find f vec)             First element where f is true
  (index-of f vec)         Index of first match
  (take n vec)             First n elements
  (drop n vec)             Remove first n elements
  (take-while f vec)       Take while f is true
  (drop-while f vec)       Drop while f is true
  (zip-with f v1 v2)       Combine two vectors element-wise
  (group-by f vec)         Group by key function → dict
  (unique vec)             Remove duplicates
  (frequencies vec)        Count occurrences → dict
  (partition f vec)        Split into (matches, non-matches)
  (scan f init vec)        Running accumulation
  (chunks n vec)           Split into groups of n
  (interleave v1 v2)       Alternate elements
  (repeat n val)           Vector of n copies
  (iterate f init n)       Generate by repeated application
  (sort vec)               Sort ascending
  (sort-by f vec)          Sort by key function
  (reverse vec)            Reverse order
  (concat v1 v2)           Concatenate vectors
  (len vec)                Length
  (head vec)               First element
  (tail vec)               All but first
  (last vec)               Last element
  (nth n vec)              Element at index n
  (range a b)              Numbers from a to b-1
  (sum vec)                Sum all elements
  (product vec)            Multiply all elements
  (any? f vec)             True if any match
  (all? f vec)             True if all match
  (count f vec)            Count matches
  (for-each f vec)         Execute f for side effects

DICTIONARIES:
  (dict k1 v1 k2 v2 ...)  Create dictionary
  (dict-get d key)         Get value (optional default)
  (dict-set d key val)     Set key → new dict
  (dict-del d key)         Remove key → new dict
  (dict-keys d)            All keys as vector
  (dict-vals d)            All values as vector
  (dict-pairs d)           Key-value pairs as vector of tuples
  (dict-has? d key)        Key exists?
  (dict-merge d1 d2)       Merge two dicts
  (dict-size d)            Number of entries

STRINGS:
  (fmt template args...)   String interpolation: (fmt "Hi {}!" name)
  (split str sep)          Split string
  (join sep vec)           Join with separator
  (upper str)              Uppercase
  (lower str)              Lowercase
  (trim str)               Strip whitespace
  (replace str old new)    Replace substring
  (starts-with? str pre)   Prefix check
  (ends-with? str suf)     Suffix check
  (chars str)              String → vector of chars
  (char-at str n)          Character at index
  (pad-left str n ch)      Left pad to width
  (pad-right str n ch)     Right pad to width

MATH:
  (abs x)                  Absolute value
  (min a b ...) or (min v) Minimum
  (max a b ...) or (max v) Maximum
  (clamp x lo hi)          Clamp to range
  (floor x)                Floor
  (ceil x)                 Ceiling
  (round x)                Round
  pi e inf nan             Constants
  (rand)                   Random float [0, 1)
  (rand-int lo hi)         Random int [lo, hi)

TYPE PREDICATES:
  (int? x) (float? x) (num? x) (str? x) (bool? x)
  (vec? x) (tuple? x) (dict? x) (fn? x) (nil? x) (py? x)
  (type x)                 Type name as string
  (zero? x) (pos? x) (neg? x) (even? x) (odd? x)
  (divides? n d)           n divisible by d?
  (empty? x)               Empty collection or string?

JSON:
  (json-parse str)         Parse JSON → Tengwar value
  (json-encode val)        Encode to JSON string

FILE I/O:
  (read-file path)         Read file contents
  (write-file path str)    Write string to file
  (append-file path str)   Append to file
  (file-exists? path)      Check if file exists

HTTP:
  (http-get url)           GET request → dict with status, headers, body
  (http-post url body)     POST request

PYTHON INTEROP:
  (py-import "module")     Import Python module
  (py-call obj "method" args...)  Call method
  (py-attr obj "attr")     Get attribute
  (py-eval "expression")   Evaluate Python expression
  obj.attr                 Dot access on Python objects

SYSTEM:
  (time-ms)                Current time in milliseconds
  (sleep ms)               Sleep for milliseconds
  (uuid)                   Generate UUID
  (env-get name)           Environment variable

COMPREHENSION:
  (for x coll body)        → (map (fn x body) coll)
  (for x coll when pred body) → filter + map
"""

SYNTAX_REFERENCE = """TENGWAR SYNTAX REFERENCE (v0.3)

LITERALS:
  42  3.14  "hello"  true  false  nil

COLLECTIONS:
  ⟦1 2 3⟧  or  [1 2 3]     Vector
  ⟨1 2 3⟩                   Tuple
  (dict "a" 1 "b" 2)        Dictionary

FUNCTIONS:
  (λ x y (+ x y))           Lambda (or: fn x y (+ x y))
  {+ _ 1}                   Short lambda (_ is parameter)
  (defn name args... body)  Define named function

CONTROL FLOW:
  (? cond then else)        Conditional (or: if)
  (cond (t1 v1) (t2 v2) (_ default))  Multi-branch
  (match expr (pat body) ...)  Pattern matching

BINDINGS:
  (≡ name value)            Define (or: def)
  (let x 1 y 2 (+ x y))    Local bindings
  (>> e1 → name e2)         Sequence + bind

ITERATION:
  (for x coll body)         Comprehension (map)
  (for x coll when p body)  Comprehension (filter + map)
  (pipe val f1 f2 f3)       Pipeline

RECURSION:
  (↺ name (λ args body))    Recursive definition (or: rec)
  (defn name args body)     Named function (auto-recursive)

ERROR HANDLING:
  (throw expr)              Throw error
  (catch expr handler)      Catch error
  (try thunk handler)       Try/catch returning (ok?, result)

PARALLEL:
  (‖ e1 e2 e3)              Parallel execution (or: par)

PYTHON:
  (py-import "module")      Import
  (py-call obj "method" args...)
  obj.attr                  Dot access
"""


def main():
    server = TengwarMCPServer()
    server.run_stdio()


if __name__ == "__main__":
    main()
