#!/usr/bin/env python3
"""stem-mcp SymPy worker. JSON-lines over stdio: {"id","op","args"} -> {"id","ok","result"|"error"}.

Parsing policy (spec §4): latex2sympy2_extended preferred; else sympy parse_latex(strict=True);
NEVER non-strict ANTLR — it silently mis-parses (e.g. 'x -' -> x) instead of raising.

Framing: exactly one JSON object per line, in and out. json.dumps defaults to ensure_ascii=True
so a response can never contain a raw newline or a non-ASCII byte.
"""
import json
import sys

import sympy
from sympy import simplify, expand, factor, solve, diff, integrate, latex as to_latex

# Node writes UTF-8; don't let a C/POSIX locale turn that into a decode error.
for _stream in (sys.stdin, sys.stdout):
    try:
        _stream.reconfigure(encoding="utf-8")
    except Exception:
        pass

PARSER = "sympy-strict"
try:
    from latex2sympy2_extended import latex2sympy  # type: ignore

    PARSER = "latex2sympy2_extended"
except Exception:
    latex2sympy = None


def parse(tex):
    if latex2sympy is not None:
        return latex2sympy(tex)
    from sympy.parsing.latex import parse_latex

    return parse_latex(tex, strict=True)


def handle(op, a):
    if op == "probe":
        return {"sympy": sympy.__version__, "parser": PARSER}
    if op == "parse":
        e = parse(a["latex"])
        return {"sympy": str(e), "latex": to_latex(e)}
    if op == "transform":
        e = parse(a["latex"])
        act = a["action"]
        sym = sympy.Symbol(a.get("symbol") or "x")
        actions = {
            "simplify": lambda: simplify(e),
            "expand": lambda: expand(e),
            "factor": lambda: factor(e),
            "solve": lambda: solve(e, sym),
            "diff": lambda: diff(e, sym),
            "integrate": lambda: integrate(e, sym),
        }
        if act not in actions:
            raise ValueError(f"unknown action: {act}")
        r = actions[act]()
        return {"latex": to_latex(r), "plain": str(r)}
    if op == "verify":
        d = simplify(parse(a["lhs"]) - parse(a["rhs"]))
        return {"equivalent": d == 0}
    raise ValueError(f"unknown op: {op}")


def main():
    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            req = json.loads(line)
        except Exception:
            continue  # unframed junk: ignore rather than die
        try:
            out = {"id": req["id"], "ok": True, "result": handle(req["op"], req.get("args", {}))}
        except Exception as e:
            out = {"id": req.get("id"), "ok": False, "error": {"message": f"{type(e).__name__}: {e}"}}
        sys.stdout.write(json.dumps(out) + "\n")
        sys.stdout.flush()


if __name__ == "__main__":
    main()