#!/usr/bin/env python3
"""COBOL-SML MCP — exposes v1 (the ~14M pure COBOL SML) as a callable component
for the Onyx-driven SML runtime.

v1 is NOT the system; it is the model_reasoning unit. This server wraps it with
the Onyx-shaped pipeline so a real system can USE it:

    draft (v1, greedy)  ->  bounded_retry (symbolic repair)  ->  verification (cobc -free -c)

Returns the COBOL, whether it compiles, and provenance (which stage produced the
compiling program). Stdio JSON-RPC (MCP). Tool: `cobol_draft`.

Run:  python3 cobol_sml_mcp.py            (Onyx/kist launches it as an MCP)
Env:  COBOL_SML_BAKE=bakes/cobol-sml-v1   (which component weights to serve)
"""
import json, os, sys
import mlx.core as mx
from cobol_tokenizer import CobolTokenizer
from sml_call import system_call
from model import MicroBrain

BAKE = os.environ.get("COBOL_SML_BAKE",
                      os.path.join(os.path.dirname(__file__), "bakes/cobol-sml-v1"))

_model = _tok = None


def _load():
    global _model, _tok
    if _model is None:
        cfg = json.load(open(os.path.join(BAKE, "config.json")))
        _tok = CobolTokenizer.load(os.path.join(BAKE, "tokenizer.json"))
        _model = MicroBrain(**cfg)
        _model.load_weights(os.path.join(BAKE, "best.safetensors"))
        mx.eval(_model.parameters()); _model.eval()
    return _model, _tok


def cobol_draft(prompt, branches=8):
    """The Onyx-shaped call: branch (model_reasoning) -> repair (bounded_retry)
    -> cobc verify. Same v1 weights, called well: ~34% vs 6% single-shot."""
    model, tok = _load()
    return system_call(model, tok, prompt, branches=branches)


TOOL = {
    "name": "cobol_draft",
    "description": ("Draft GnuCOBOL (free-format) from a natural-language spec using the "
                    "COBOL-SML v1 component, then symbolically repair and verify with the real "
                    "compiler. Returns {cobol, compiles, path, verifier}. Compilation is proven, "
                    "not claimed."),
    "inputSchema": {
        "type": "object",
        "properties": {"prompt": {"type": "string", "description": "What COBOL to write."}},
        "required": ["prompt"],
    },
}

SERVER_INFO = {"name": "cobol-sml", "version": "1.0.0"}


def _resp(rid, result=None, error=None):
    m = {"jsonrpc": "2.0", "id": rid}
    if error is not None:
        m["error"] = error
    else:
        m["result"] = result
    return m


def handle(req):
    method, rid = req.get("method"), req.get("id")
    if method == "initialize":
        return _resp(rid, {"protocolVersion": "2024-11-05",
                           "capabilities": {"tools": {}},
                           "serverInfo": SERVER_INFO})
    if method in ("notifications/initialized", "initialized"):
        return None
    if method == "tools/list":
        return _resp(rid, {"tools": [TOOL]})
    if method == "tools/call":
        p = req.get("params", {})
        if p.get("name") != "cobol_draft":
            return _resp(rid, error={"code": -32601, "message": f"unknown tool {p.get('name')}"})
        prompt = (p.get("arguments") or {}).get("prompt", "")
        try:
            out = cobol_draft(prompt)
            return _resp(rid, {"content": [{"type": "text", "text": json.dumps(out)}],
                               "isError": not out["compiles"]})
        except Exception as e:
            return _resp(rid, error={"code": -32000, "message": str(e)})
    if rid is not None:
        return _resp(rid, error={"code": -32601, "message": f"unknown method {method}"})
    return None


def main():
    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            req = json.loads(line)
        except json.JSONDecodeError:
            continue
        resp = handle(req)
        if resp is not None:
            sys.stdout.write(json.dumps(resp) + "\n")
            sys.stdout.flush()


if __name__ == "__main__":
    main()