import hashlib
import json
import os
import re

for _key in ("OMP_NUM_THREADS", "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS"):
    try:
        _threads = int(os.environ.get(_key, "4"))
    except ValueError:
        _threads = 4
    os.environ[_key] = str(max(1, min(4, _threads)))
import numpy as np  # noqa: E402 -- thread limits must precede NumPy initialization

ROLES = ["none", "tool", "arg", "value", "field", "limit"]
MODES = ["single", "sequential", "parallel"]
OUTPUTS = ["all", "last"]


class Diagnostic(ValueError):
    pass


NUMBER = re.compile(r"-?(?:0|[1-9]\d*)(?:\.\d+)?(?:[eE][+-]?\d+)?\Z")
TOKEN = re.compile(r'"(?:\\.|[^"\\])*"|[A-Za-z_$][\w$]*|[+-]?\d[\w.+-]*|[^\s]')


def number_ok(v):
    import math

    if type(v) is int:
        try:
            fv = float(v)
        except OverflowError as error:
            raise Diagnostic("number: JS overflow") from error
        if not math.isfinite(fv) or int(fv) != v:
            raise Diagnostic("number: integer not exactly JS-representable")
    elif type(v) is float:
        if not math.isfinite(v):
            raise Diagnostic("number: nonfinite")
    else:
        raise Diagnostic("number: primitive required")


def catalog_check(tools):
    if not isinstance(tools, list) or not 1 <= len(tools) <= 8:
        raise Diagnostic("catalog: expected 1-8 tools")
    names = set()
    for t in tools:
        if not isinstance(t, dict) or set(t) != {
            "name",
            "arguments",
            "required",
            "returns",
        }:
            raise Diagnostic("catalog: tool shape")
        name = t["name"]
        a = t["arguments"]
        req = t["required"]
        ret = t["returns"]
        if (
            not isinstance(name, str)
            or not re.fullmatch(r"[A-Za-z_$][A-Za-z0-9_$]*", name)
            or name in names
        ):
            raise Diagnostic("catalog: unsafe/duplicate tool")
        names.add(name)
        if (
            not isinstance(a, dict)
            or len(a) > 4
            or any(
                not isinstance(k, str)
                or not k
                or v not in ("string", "number", "boolean")
                for k, v in a.items()
            )
        ):
            raise Diagnostic("catalog: primitive argument schema")
        if (
            not isinstance(req, list)
            or any(not isinstance(k, str) for k in req)
            or len(set(req)) != len(req)
            or set(req) - set(a)
        ):
            raise Diagnostic("catalog: required keys")
        if (
            not isinstance(ret, dict)
            or set(ret) - {"kind", "properties"}
            or ret.get("kind") not in ("text", "object", "array")
        ):
            raise Diagnostic("catalog: return schema")
        props = ret.get("properties", [])
        if (
            not isinstance(props, list)
            or any(not isinstance(k, str) or not k for k in props)
            or len(set(props)) != len(props)
        ):
            raise Diagnostic("catalog: properties")
    return tools


def frontend(item):
    if not isinstance(item, dict) or not isinstance(item.get("instruction"), str):
        raise Diagnostic("input: instruction string required")
    tools = catalog_check(item.get("tools"))
    s = item["instruction"]
    toks = TOKEN.findall(s)
    if not toks:
        raise Diagnostic("input: empty instruction")
    if len(toks) > 128:
        raise Diagnostic("token bound")
    names = {t["name"] for t in tools}
    args = {a for t in tools for a in t["arguments"]}
    fields = {a for t in tools for a in t["returns"].get("properties", [])}
    ids = []
    features = []
    values = []
    for x in toks:
        literal = False
        v = x
        if x.startswith('"'):
            try:
                v = json.loads(x)
                literal = True
            except Exception as error:
                raise Diagnostic("malformed literal") from error
        elif x in ("true", "false"):
            v = x == "true"
            literal = True
        elif x[0].isdigit() or (x[0] in "+-" and len(x) > 1 and x[1].isdigit()):
            if not NUMBER.fullmatch(x):
                raise Diagnostic("number: malformed JSON syntax")
            try:
                v = json.loads(x)
                number_ok(v)
            except (ValueError, OverflowError) as e:
                raise Diagnostic(str(e)) from e
            from decimal import Decimal

            decimal = Decimal(x)
            if v == 0 and decimal != 0:
                raise Diagnostic("number: JS underflow")
            if decimal == decimal.to_integral_value() and int(v) != int(decimal):
                raise Diagnostic("number: integer not exactly JS-representable")
            if (
                type(v) is float
                and v.is_integer()
                and decimal != decimal.to_integral_value()
            ):
                raise Diagnostic("number: fractional part lost in JS")
            literal = True
        f = [
            float(x in names and not literal),
            float(x in args and not literal),
            float(x in fields and not literal),
            float(literal and isinstance(v, str)),
            float(literal and type(v) in (int, float)),
            float(literal and type(v) is bool),
        ]
        f += [
            float(not literal and any(t["arguments"].get(x) == typ for t in tools))
            for typ in ("string", "number", "boolean")
        ]
        f += [
            float(
                not literal
                and any(t["name"] == x and t["returns"]["kind"] == kind for t in tools)
            )
            for kind in ("text", "object", "array")
        ]
        norm = "<schema>" if any(f[:3]) else "<literal>" if literal else x.lower()
        ids.append(
            int.from_bytes(hashlib.sha256(norm.encode()).digest()[:4], "little") % 1024
        )
        features.append(f)
        values.append(v)
    for i, x in enumerate(toks):
        if x == "=" and (i + 1 >= len(toks) or not any(features[i + 1][3:6])):
            raise Diagnostic("literal: assignment requires JSON string/number/boolean")
    return (
        toks,
        np.array(ids, dtype=np.int64),
        np.array(features, dtype=np.float32),
        values,
    )


def validate(plan, tools):
    catalog_check(tools)
    if not isinstance(plan, dict) or set(plan) != {"mode", "calls", "output"}:
        raise Diagnostic("plan: shape")
    catalog = {t["name"]: t for t in tools}
    calls = plan["calls"]
    out = plan["output"]
    if not isinstance(calls, list) or not 1 <= len(calls) <= 3:
        raise Diagnostic("call bound")
    if (
        not isinstance(plan["mode"], str)
        or plan["mode"] not in MODES
        or (len(calls) == 1) != (plan["mode"] == "single")
    ):
        raise Diagnostic("mode mismatch")
    if (
        not isinstance(out, dict)
        or set(out) - {"kind", "fields", "limit"}
        or out.get("kind") not in OUTPUTS
    ):
        raise Diagnostic("output kind/shape")
    for c in calls:
        if (
            not isinstance(c, dict)
            or set(c) != {"tool", "args"}
            or not isinstance(c["tool"], str)
            or c["tool"] not in catalog
        ):
            raise Diagnostic("call: unknown tool/shape")
        t = catalog[c["tool"]]
        a = c["args"]
        if (
            not isinstance(a, dict)
            or len(a) > 4
            or any(not isinstance(k, str) for k in a)
            or set(a) - set(t["arguments"])
            or set(t["required"]) - set(a)
        ):
            raise Diagnostic("unknown/missing arguments")
        for k, v in a.items():
            typ = (
                "boolean"
                if type(v) is bool
                else "number"
                if type(v) in (int, float)
                else "string"
                if type(v) is str
                else "invalid"
            )
            if typ != t["arguments"][k]:
                raise Diagnostic("argument type")
            if typ == "number":
                number_ok(v)
    if "fields" in out or "limit" in out:
        t = catalog[calls[0]["tool"]]
        if len(calls) != 1 or t["returns"]["kind"] != "array":
            raise Diagnostic("shaping scope")
        if "fields" in out:
            fs = out["fields"]
            if (
                not isinstance(fs, list)
                or not 1 <= len(fs) <= 3
                or any(not isinstance(k, str) for k in fs)
                or len(set(fs)) != len(fs)
                or set(fs) - set(t["returns"].get("properties", []))
            ):
                raise Diagnostic("projection fields")
        if "limit" in out and (
            type(out["limit"]) is not int or not 0 <= out["limit"] <= 100
        ):
            raise Diagnostic("limit range")
    return plan


def assemble(item, roles, mode, output):
    toks, _, _, vals = frontend(item)
    if (
        not hasattr(roles, "__len__")
        or len(roles) != len(toks)
        or any(
            not isinstance(r, (int, np.integer))
            or isinstance(r, (bool, np.bool_))
            or r not in range(6)
            for r in roles
        )
    ):
        raise Diagnostic("roles: length/value mismatch")
    if (
        not isinstance(mode, (int, np.integer))
        or not isinstance(output, (int, np.integer))
        or isinstance(mode, (bool, np.bool_))
        or isinstance(output, (bool, np.bool_))
        or mode not in range(3)
        or output not in range(2)
    ):
        raise Diagnostic("heads: invalid class")
    calls = []
    pending = None
    fields = []
    limit = None
    for x, v, r in zip(toks, vals, roles, strict=False):
        if r == 1:
            if pending is not None:
                raise Diagnostic("unbound argument")
            calls.append({"tool": x, "args": {}})
        elif r == 2:
            if not calls or pending is not None:
                raise Diagnostic("argument structure")
            pending = x
        elif r == 3:
            if pending is None or pending in calls[-1]["args"]:
                raise Diagnostic("value structure/duplicate")
            calls[-1]["args"][pending] = v
            pending = None
        elif r == 4:
            fields.append(x)
        elif r == 5:
            if limit is not None:
                raise Diagnostic("multiple limits")
            limit = v
    if pending is not None:
        raise Diagnostic("missing value")
    out = {"kind": OUTPUTS[output]}
    if fields:
        out["fields"] = fields
    if limit is not None:
        out["limit"] = limit
    return validate({"mode": MODES[mode], "calls": calls, "output": out}, item["tools"])


def compile_js(plan, tools):
    validate(plan, tools)

    def j(x):
        return json.dumps(x, ensure_ascii=True, separators=(",", ":"))

    inv = [
        f"tools[{j(c['tool'])}](JSON.parse({j(j(c['args']))}))" for c in plan["calls"]
    ]
    pre = "const err=e=>String(e && e.message || e); const invoke=async f=>{try{return {ok:true,value:await f()}}catch(e){return {ok:false,error:err(e)}}};\n"
    if plan["mode"] == "parallel":
        pre += (
            "const settled=await Promise.allSettled(["
            + ",".join("invoke(()=>" + x + ")" for x in inv)
            + ']); const r=settled.map(s=>s.status==="fulfilled"?s.value:{ok:false,error:err(s.reason)});\n'
        )
    else:
        pre += "const r=[];\n" + "".join(
            "r.push(await invoke(()=>" + x + "));\n" for x in inv
        )
    out = plan["output"]
    if "fields" in out or "limit" in out:
        pre += 'if(r[0].ok){try{let v=r[0].value; if(!Array.isArray(v))throw Error("result-schema: expected array");'
        if "fields" in out:
            pre += (
                "const fields="
                + j(out["fields"])
                + '; v=v.map(row=>{if(row===null||typeof row!=="object"||Array.isArray(row)||fields.some(k=>!Object.hasOwn(row,k)))throw Error("result-schema: missing field"); return Object.fromEntries(fields.map(k=>[k,row[k]]));});'
            )
        if "limit" in out:
            pre += f"v=v.slice(0,{out['limit']});"
        pre += "r[0]={ok:true,value:v};}catch(e){r[0]={ok:false,error:err(e)}}}\n"
    return pre + ("return r;" if out["kind"] == "all" else "return r[r.length-1];")