#!/usr/bin/env python3
"""CSCS MarketSphere — local MCP server (stdio). REAL, runnable, stdlib-only.

Exposes CSCS tools to any MCP-capable AI agent (Claude Desktop, copilots). Each tool call goes to the
CSCS MarketSphere API with your API key, so it is authenticated and METERED (credits debited) by the
same billing spine as REST — no separate billing. Sandbox-only until the CSCS production licence lands.

Claude Desktop config (claude_desktop_config.json):
  {
    "mcpServers": {
      "cscs-marketsphere": {
        "command": "python",
        "args": ["/path/to/mcp_server.py"],
        "env": {
          "CSCS_MCP_API_KEY": "cscs_sk_sandbox_...",
          "CSCS_MCP_BASE": "https://<your MarketSphere deployment>"
        }
      }
    }
  }

The deployment is REQUIRED configuration: set CSCS_MCP_BASE (or pass ``--base <url>``) to the
MarketSphere deployment that issued your key. There is deliberately no default -- a default would
send your key and every question to whichever deployment happened to be written here.

A failed call (bad key, no plan scope, no credits, network error) is returned as a tool error
(``isError: true``) so an agent never mistakes it for data. Large results are cut at a whole
record, never mid-JSON: the text is always valid JSON, with a ``_truncated`` marker saying which
list was cut and how many records were left out.

Implements MCP JSON-RPC 2.0 over stdio: initialize, tools/list, tools/call. No third-party deps.
"""
import os
import re
import sys
import json
import urllib.parse
import urllib.request

def _base_from(argv: list, env) -> str:
    """The deployment to call: ``--base <url>`` / ``--base=<url>`` on the command line, else
    CSCS_MCP_BASE. Empty when neither is given -- there is no default deployment."""
    args = list(argv[1:])
    for i, a in enumerate(args):
        if a == "--base" and i + 1 < len(args):
            return args[i + 1].strip().rstrip("/")
        if a.startswith("--base="):
            return a.split("=", 1)[1].strip().rstrip("/")
    return (env.get("CSCS_MCP_BASE", "") or "").strip().rstrip("/")


BASE = _base_from(sys.argv, os.environ)
KEY = os.environ.get("CSCS_MCP_API_KEY", "")
MAX_RESULT_CHARS = 8000      # what one tool result may carry back to the agent

# tool name -> (api path template, needs_symbol, one-line description)
TOOLS = {
    "cscs_securities":  ("/api/v1/securities", False, "Priced securities master."),
    "cscs_prices":      ("/api/v1/prices", False, "Latest prices for all securities."),
    "cscs_symbols":     ("/api/v1/symbols", False, "Symbol directory (code + name)."),
    "cscs_ask_simba":   ("/api/v1/ask?q={q}", False, "Grounded natural-language Q&A over CSCS data. Arg: q (question)."),
    "cscs_news_intel":  ("/api/v1/news-intel/{symbol}", True, "AI news-intelligence for a security. Arg: symbol."),
    "cscs_stock360":    ("/api/v1/stock/360/{symbol}", True, "360° single-stock analysis (PREMIUM): CSCS identity + "
                         "quant + ownership + grounded research + plain-English verdict. Full report on Standard+ keys; "
                         "a glimpse + upgrade prompt on free/developer keys (RBAC + credits enforced by your key). "
                         "Args: symbol; optional focus (comma list of market,quant,ownership,sentiment,research) and "
                         "depth (standard|deep) to steer the report."),
}


# Investor identity must never reach a third-party agent. The server already masks every
# /api/v1 response (a key carries no role); this is a second, client-side layer so a
# misconfigured or impersonated server still cannot make this script forward a CHN.
_IDENTITY_KEYS = {"chn", "investorchn", "top_chn", "top_owner_chn", "chno", "account", "acountnumber",
                  "accountnumber", "cscsaccountno", "account_name", "accountname", "full_name", "fullname",
                  "investor", "owner", "bankaccname"}
_CHN_RE = re.compile(r"\bC\d{7,12}[A-Z]{2}\b")


def _mask(v):
    s = str(v)
    return ("\u2022" * max(0, len(s) - 4) + s[-4:]) if len(s) > 4 else "\u2022" * len(s)


def _scrub(o):
    if isinstance(o, dict):
        return {k: (_mask(v) if k.lower() in _IDENTITY_KEYS and isinstance(v, (str, int)) and v != ""
                    else _scrub(v)) for k, v in o.items()}
    if isinstance(o, list):
        return [_scrub(x) for x in o]
    if isinstance(o, str):
        return _CHN_RE.sub(lambda m: _mask(m.group(0)), o)
    return o


def _base_ok(base: str) -> bool:
    """https only -- the API key and the data travel over this connection. Plain http is
    accepted for a local development server only."""
    u = urllib.parse.urlsplit(base)
    return u.scheme == "https" or (u.scheme == "http" and u.hostname in ("localhost", "127.0.0.1", "::1"))


# _call_api marks a failed call with this key; _tools_call pops it and sets isError.
_FAILED = "_mcp_call_failed"


def _call_api(path: str) -> dict:
    req = urllib.request.Request(BASE + path, headers={"X-API-Key": KEY, "Accept": "application/json"})
    try:
        with urllib.request.urlopen(req, timeout=30) as r:
            return json.loads(r.read().decode())
    except urllib.error.HTTPError as e:
        try:
            detail = (e.read() or b"").decode(errors="replace")[:300]
        except Exception:
            detail = ""
        return {_FAILED: True, "error": f"HTTP {e.code}", "detail": detail}
    except Exception as e:
        return {_FAILED: True, "error": str(e) or e.__class__.__name__}


def _dumps(o) -> str:
    return json.dumps(o, indent=2, ensure_ascii=False)


def _largest_list(o, path=()):
    """(path, list) of the longest list in the payload, depth-first; (None, None) if there is none."""
    best = (None, None)
    if isinstance(o, list):
        best = (path, o)
        items = enumerate(o)
    elif isinstance(o, dict):
        items = o.items()
    else:
        return best
    for k, v in items:
        p, lst = _largest_list(v, path + (k,))
        if lst is not None and (best[1] is None or len(lst) > len(best[1])):
            best = (p, lst)
    return best


def _replace_at(o, path, value):
    if not path:
        return value
    if isinstance(o, dict):
        return {k: (_replace_at(v, path[1:], value) if k == path[0] else v) for k, v in o.items()}
    return [(_replace_at(v, path[1:], value) if i == path[0] else v) for i, v in enumerate(o)]


def _fit(data, limit: int) -> str:
    """Serialise ``data`` into at most ``limit`` characters of VALID JSON. Over the limit, the
    largest list is cut at a whole record (binary search on how many fit) and a ``_truncated``
    marker says which list, how many records are shown and how many more exist. With no list to
    cut, the longest top-level string is shortened instead, and the marker says so."""
    text = _dumps(data)
    if len(text) <= limit:
        return text
    wrap = data if isinstance(data, dict) else {"result": data}
    path, lst = _largest_list(wrap)
    if lst:
        field = ".".join(str(p) for p in path)

        def attempt(n: int) -> str:
            marked = _replace_at(wrap, path, lst[:n])
            marked["_truncated"] = {"field": field, "shown": n, "more": len(lst) - n,
                                    "note": "result cut to fit one tool reply; narrow the query for the rest"}
            return _dumps(marked)
        lo, hi = 0, len(lst)
        while lo < hi:
            mid = (lo + hi + 1) // 2
            if len(attempt(mid)) <= limit:
                lo = mid
            else:
                hi = mid - 1
        out = attempt(lo)
        if len(out) <= limit:
            return out
    strs = [(k, v) for k, v in wrap.items() if isinstance(v, str)]
    if strs:
        k, v = max(strs, key=lambda kv: len(kv[1]))
        room = len(v)
        while room >= 0:
            cut = {**wrap, k: v[:room], "_truncated": {"field": k, "shown_chars": room, "more_chars": len(v) - room}}
            out = _dumps(cut)
            if len(out) <= limit:
                return out
            room -= max(1, len(out) - limit)
    return _dumps({"_truncated": {"field": "", "note": "result too large for one tool reply; narrow the query"}})


def _tools_list() -> list:
    out = []
    for name, (path, needs_sym, desc) in TOOLS.items():
        props, required = {}, []
        if "{q}" in path:
            props["q"] = {"type": "string", "description": "The question to ask"}; required.append("q")
        if needs_sym:
            props["symbol"] = {"type": "string", "description": "Security symbol"}; required.append("symbol")
        if name == "cscs_stock360":
            props["focus"] = {"type": "string", "description": "Optional: comma list of sections to focus "
                              "(market,quant,ownership,sentiment,research). Full-tier keys only."}
            props["depth"] = {"type": "string", "enum": ["standard", "deep"], "description": "Report depth."}
        out.append({"name": name, "description": desc,
                    "inputSchema": {"type": "object", "properties": props, "required": required}})
    return out


def tool_path(name: str, args: dict) -> str:
    """The /api/v1 path (with query) a tool call maps to. Arguments are URL-quoted with NO safe
    characters, so a symbol or question can never add a path segment ("../") or a query parameter.
    Shared by this stdio server and the remote transport (backend/mcp_remote.py). KeyError for an
    unknown tool."""
    path, needs_sym, _ = TOOLS[name]
    args = args or {}

    def q(v) -> str:
        return urllib.parse.quote(str(v or ""), safe="")
    if "{q}" in path:
        path = path.format(q=q(args.get("q", "")))
    if needs_sym:
        path = path.format(symbol=q(args.get("symbol", "")))
    if name == "cscs_stock360":
        qs = []
        if args.get("focus"):
            qs.append("focus=" + q(args["focus"]))
        if args.get("depth"):
            qs.append("depth=" + q(args["depth"]))
        if qs:
            path += ("&" if "?" in path else "?") + "&".join(qs)
    return path


def tool_result(raw, failed: bool, limit: int = 0) -> dict:
    """An MCP tool result: identity scrubbed, fitted to ``limit`` as valid JSON, isError on failure."""
    out = {"content": [{"type": "text", "text": _fit(_scrub(raw), limit or MAX_RESULT_CHARS)}]}
    if failed:
        out["isError"] = True
    return out


def error_result(message: str) -> dict:
    return {"content": [{"type": "text", "text": message}], "isError": True}


def _tools_call(name: str, args: dict) -> dict:
    if name not in TOOLS:
        return error_result(f"Unknown tool: {name}")
    if not KEY:
        return error_result("No API key. Set CSCS_MCP_API_KEY (get an API key from the CSCS developer portal).")
    if not BASE:
        return error_result("No MarketSphere deployment configured. Set CSCS_MCP_BASE (or start the server with "
                            "--base <url>) to the https:// address of the deployment that issued your API key.")
    if not _base_ok(BASE):
        return error_result("Refusing to send your API key over plain http: CSCS_MCP_BASE must be an https:// URL.")
    raw = _call_api(tool_path(name, args))
    failed = isinstance(raw, dict) and bool(raw.pop(_FAILED, False))
    return tool_result(raw, failed)


def _send(obj: dict):
    sys.stdout.write(json.dumps(obj) + "\n")
    sys.stdout.flush()


def main():
    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            msg = json.loads(line)
        except Exception:
            continue
        mid = msg.get("id")
        method = msg.get("method")
        if method == "initialize":
            _send({"jsonrpc": "2.0", "id": mid, "result": {
                "protocolVersion": "2024-11-05",
                "serverInfo": {"name": "cscs-marketsphere", "version": "0.1.0"},
                "capabilities": {"tools": {}}}})
        elif method == "tools/list":
            _send({"jsonrpc": "2.0", "id": mid, "result": {"tools": _tools_list()}})
        elif method == "tools/call":
            p = msg.get("params", {})
            _send({"jsonrpc": "2.0", "id": mid, "result": _tools_call(p.get("name"), p.get("arguments", {}))})
        elif method in ("notifications/initialized", "notifications/cancelled"):
            pass  # notifications: no response
        elif mid is not None:
            _send({"jsonrpc": "2.0", "id": mid, "error": {"code": -32601, "message": f"Method not found: {method}"}})


if __name__ == "__main__":
    main()