#!/usr/bin/env python3
"""tool_call_shim — serve-time parser shim (the seatbelt, not the diploma).

Normalizes a model's tool-call output to the strict production format and reports
BOTH scores so you never confuse "model emitted it" with "runtime rescued it":
  - strict_parse_pass : model emitted a real <tool_call>{json}</tool_call>
  - lenient_recovered : model emitted BARE {"name","arguments"} -> shim wrapped it

Every lenient recovery is logged to runs/wrapper_recoveries.jsonl so bare-vs-wrapped
becomes chosen/rejected preference data for a later DPO/ORPO pass.

serve() is the SAFE production entry. Hard rules:
  * NO SCHEMA = NO EXECUTION. tools empty/missing -> intercept all, never execute.
  * RUNTIME MISSING-ARG VALIDATOR. required args validated vs schema before execution;
    missing-arg / unknown-tool calls are intercepted -> clarification, never executed.
  * ATOMIC PARALLEL. if ANY call is intercepted, execute ZERO calls (all-valid-or-none).
  * The shim fixes FORMAT. It never INVENTS missing args, executes unknown tools, or
    partially executes a multi-call response.

  parse only:   calls, info, normalized        = normalize(model_output)
  serve safely: ex, info, normalized, clarify   = serve(model_output, tools)
"""
import ast, re, json, os, time

STRICT = re.compile(r"<tool_call>\s*(\{.*?\})\s*</tool_call>", re.S)
# a bare top-level tool-call object: {"name": "...", "arguments": {...}}
BARE = re.compile(r"\{[^{}]*\"name\"\s*:\s*\"[^\"]+\"[^{}]*\"arguments\"\s*:\s*\{.*?\}\s*\}", re.S)
# LFM often emits pseudo-calls like:
#   [get_weather(city="Tokyo", unit="c"), get_weather(city="London")]
# Treat these as lenient recoveries, never as strict native protocol.
FN_CALL = re.compile(r"\b([A-Za-z_]\w*)\s*\(([^()]*)\)")
_RUNS = os.environ.get("VT_RUNS_DIR", os.path.expanduser("~/vibethinker-tune/runs"))
REC_LOG       = os.path.join(_RUNS, "wrapper_recoveries.jsonl")
MISS_LOG      = os.path.join(_RUNS, "missing_arg_intercepted.jsonl")
COUNTERS_PATH = os.path.join(_RUNS, "shim_counters.json")
LOG_MAX_BYTES = int(os.environ.get("VT_SHIM_LOG_MAX_BYTES", str(10 * 1024 * 1024)))
LOG_BACKUPS = int(os.environ.get("VT_SHIM_LOG_BACKUPS", "3"))

def _rotate_file(path, max_bytes=LOG_MAX_BYTES, backups=LOG_BACKUPS):
    if max_bytes <= 0 or backups <= 0:
        return
    try:
        if not os.path.exists(path) or os.path.getsize(path) < max_bytes:
            return
        oldest = f"{path}.{backups}"
        if os.path.exists(oldest):
            os.remove(oldest)
        for i in range(backups - 1, 0, -1):
            src, dst = f"{path}.{i}", f"{path}.{i + 1}"
            if os.path.exists(src):
                os.replace(src, dst)
        os.replace(path, f"{path}.1")
    except Exception:
        pass

# ---- telemetry counters --------------------------------------------------
COUNTERS = {k:0 for k in ("strict_native_call","shim_recovered_call","missing_arg_intercepted",
    "unknown_tool_intercepted","tool_schema_missing","partial_parallel_blocked","no_tool_output")}
def _bump(key, n=1, persist=True):
    COUNTERS[key]=COUNTERS.get(key,0)+n
    if persist:
        try: json.dump(COUNTERS, open(COUNTERS_PATH,"w"))
        except Exception: pass
def get_counters(): return dict(COUNTERS)

def _loads(s):
    try: return json.loads(s)
    except Exception: return None

def _split_args(s):
    parts=[]; cur=[]; quote=None; esc=False
    for ch in s:
        if esc:
            cur.append(ch); esc=False; continue
        if ch == "\\":
            cur.append(ch); esc=True; continue
        if quote:
            cur.append(ch)
            if ch == quote: quote=None
            continue
        if ch in ("'", '"'):
            quote=ch; cur.append(ch); continue
        if ch == ",":
            part="".join(cur).strip()
            if part: parts.append(part)
            cur=[]; continue
        cur.append(ch)
    part="".join(cur).strip()
    if part: parts.append(part)
    return parts

def _parse_fn_args(s):
    args={}
    for part in _split_args(s):
        if "=" not in part:
            return None
        k,v=part.split("=",1); k=k.strip(); v=v.strip()
        if not re.match(r"^[A-Za-z_]\w*$", k):
            return None
        try:
            val=ast.literal_eval(v)
        except Exception:
            try:
                val=json.loads(v)
            except Exception:
                val=v.strip("'\"")
        args[k]=val
    return args

def normalize(text, source="serve", log=True, tool_names=None):
    """Return (calls, info, normalized_text).
    info = {strict_parse_pass, lenient_recovered, n_calls, format}.
    tool_names: if provided, the FN_CALL dialect recovery only fires for names in this
    set — prevents fabricating a tool call from prose like 'compute f(x=1)'. When None
    (callers without a schema) it falls back to a hardcoded prose/keyword skip-list."""
    # 1) strict: already production format
    strict=[c for c in (_loads(m) for m in STRICT.findall(text)) if isinstance(c,dict) and "name" in c]
    if strict:
        return strict, {"strict_parse_pass":True,"lenient_recovered":False,"n_calls":len(strict),"format":"strict"}, text
    # 2) lenient: recover bare JSON tool-call(s) and wrap them
    bare=[]; spans=[]
    for m in BARE.finditer(text):
        d=_loads(m.group(0))
        if isinstance(d,dict) and "name" in d and "arguments" in d:
            bare.append(d); spans.append(m.span())
    if bare:
        norm=text; off=0
        for (s,e),d in zip(spans,bare):
            w="<tool_call>\n"+json.dumps(d,ensure_ascii=False)+"\n</tool_call>"
            norm=norm[:s+off]+w+norm[e+off:]; off+=len(w)-(e-s)
        if log:
            try:
                os.makedirs(os.path.dirname(REC_LOG),exist_ok=True)
                _rotate_file(REC_LOG)
                with open(REC_LOG,"a",encoding="utf-8") as f:
                    # chosen = wrapped (norm), rejected = raw bare (text) -> DPO pair
                    f.write(json.dumps({"ts":int(time.time()),"source":source,"rejected":text,"chosen":norm,"calls":bare},ensure_ascii=False)+"\n")
            except Exception: pass
        return bare, {"strict_parse_pass":False,"lenient_recovered":True,"n_calls":len(bare),"format":"lenient"}, norm
    # 3) lenient: recover function-call syntax emitted by some models.
    fn_calls=[]; spans=[]
    for m in FN_CALL.finditer(text):
        name=m.group(1)
        # Is this match inside an explicit [ ... ] tool-call list (the LFM dialect
        # `[get_weather(city='x')]`)? Brackets are a strong "this IS a tool call"
        # signal, even when the NAME is unknown/hallucinated.
        pre=text[:m.start()]
        bracketed = pre.rfind('[') > pre.rfind(']') and ']' in text[m.end():]
        # Gate the loose dialect to REAL tools when we know them; otherwise this regex
        # fabricates a "call" from any prose `word(k=v)` (e.g. 'compute f(x=1)').
        # EXCEPTION: a bracketed [name(args)] is the explicit dialect — recover it even
        # for an unknown name so serve() intercepts it as an unknown-tool clarification
        # instead of the harness leaking the raw `[get_weather(...)]` text as content.
        if tool_names is not None:
            if name not in tool_names and not bracketed:
                continue
        # No schema passed: fall back to skipping common prose/code pseudo-functions.
        elif name in {"if", "for", "while", "print", "json", "dict", "list"} and not bracketed:
            continue
        args=_parse_fn_args(m.group(2))
        if args is not None:
            fn_calls.append({"name":name,"arguments":args}); spans.append(m.span())
    if fn_calls:
        blocks=["<tool_call>\n"+json.dumps(d,ensure_ascii=False)+"\n</tool_call>" for d in fn_calls]
        # For this dialect, the clean normalized form should be just strict tool_call blocks.
        # Keeping surrounding "[" / "]" from pseudo-call lists would be poor preference data.
        norm="\n".join(blocks)
        if log:
            try:
                os.makedirs(os.path.dirname(REC_LOG),exist_ok=True)
                _rotate_file(REC_LOG)
                with open(REC_LOG,"a",encoding="utf-8") as f:
                    f.write(json.dumps({"ts":int(time.time()),"source":source,"rejected":text,"chosen":norm,"calls":fn_calls},ensure_ascii=False)+"\n")
            except Exception: pass
        return fn_calls, {"strict_parse_pass":False,"lenient_recovered":True,"n_calls":len(fn_calls),"format":"lenient_fn_call"}, norm
    # 3) no tool call at all
    return [], {"strict_parse_pass":False,"lenient_recovered":False,"n_calls":0,"format":"none"}, text

# ---- runtime missing-arg validator ---------------------------------------
def _tool_required(tools, name):
    """Required arg names for tool `name` from an OpenAI-style tools list."""
    for t in tools or []:
        if not isinstance(t, dict): continue
        fn = t.get("function", t)
        if isinstance(fn, dict) and fn.get("name") == name:
            params = fn.get("parameters") or {}
            return list(params.get("required") or [])
    return []

def validate_call(call, tools):
    """(ok, missing). Required args must be present AND non-empty/non-null.
    Unknown tool name -> ok=False, missing=['<unknown_tool>'] (never execute blind).
    NOTE: serve() enforces the no-schema-no-execution rule BEFORE this is reached."""
    name = call.get("name"); args = call.get("arguments") or {}
    known = any(isinstance(t,dict) and (t.get("function",t) or {}).get("name")==name for t in (tools or []))
    if not known:
        return (False, ["<unknown_tool>"])
    req = _tool_required(tools, name)
    missing = [a for a in req if a not in args or args.get(a) in (None, "", [], {})]
    return (len(missing)==0, missing)

def _log_intercept(source, intercepted, text):
    try:
        os.makedirs(os.path.dirname(MISS_LOG),exist_ok=True)
        _rotate_file(MISS_LOG)
        with open(MISS_LOG,"a",encoding="utf-8") as f:
            f.write(json.dumps({"ts":int(time.time()),"source":source,"intercepted":intercepted,"raw":text[:1200]},ensure_ascii=False)+"\n")
    except Exception: pass

def _build_clar(intercepted):
    if any("<tool_schema_missing>" in it["missing"] for it in intercepted):
        return "I can't safely execute a tool call because the available tool schema is missing."
    slots=sorted({m for it in intercepted for m in it["missing"] if m not in ("<unknown_tool>","<tool_schema_missing>")})
    names=sorted({it["name"] for it in intercepted})
    if slots:
        return (f"Before I can use {', '.join(names)}, I need: {', '.join(slots)}. "
                f"Could you provide {'it' if len(slots)==1 else 'them'}?")
    return f"I can't call {', '.join(names)} — it isn't an available tool here."

def serve(text, tools, source="serve", log=True):
    """Production entry. Returns (executable_calls, info, normalized_text, clarification).
    NO SCHEMA -> no execution. Missing-arg/unknown-tool -> intercept. ANY intercept ->
    execute ZERO calls (atomic). The shim corrects FORMAT; it never fabricates args,
    executes unknown tools, or partially executes a multi-call response."""
    def _tool_name(t):
        if not isinstance(t, dict): return None
        fn = t.get("function")
        fn = fn if isinstance(fn, dict) else t
        n = fn.get("name") if isinstance(fn, dict) else None
        return n if isinstance(n, str) and n else None
    tool_names={_tool_name(t) for t in (tools or [])}
    tool_names.discard(None)
    calls, info, norm = normalize(text, source=source, log=log, tool_names=(tool_names or None))
    info=dict(info)
    # parse-outcome telemetry
    if not calls: _bump("no_tool_output")
    elif info["strict_parse_pass"]: _bump("strict_native_call", len(calls))
    elif info["lenient_recovered"]: _bump("shim_recovered_call", len(calls))

    # RULE 1: no (valid) schema = no execution. Use tool_names so a malformed-but-truthy
    # `tools` (e.g. ["junk"]) can't bypass this and reach validation/execution.
    if calls and not tool_names:
        intercepted=[{"name":c.get("name"),"missing":["<tool_schema_missing>"],"call":c} for c in calls]
        _bump("tool_schema_missing", len(intercepted))
        info["intercepted"]=len(intercepted); info["executable"]=0
        clar="I can't safely execute a tool call because the available tool schema is missing."
        if log: _log_intercept(source, intercepted, text)
        return [], info, norm, clar

    # RULE 2: validate required args vs schema
    executable=[]; intercepted=[]
    for c in calls:
        try:
            ok, missing = validate_call(c, tools)
        except Exception:
            ok, missing = False, ["<unknown_tool>"]   # contain validator crashes -> intercept
        if ok: executable.append(c)
        else:
            intercepted.append({"name":c.get("name"),"missing":missing,"call":c})
            _bump("unknown_tool_intercepted" if missing==["<unknown_tool>"] else "missing_arg_intercepted")
    info["intercepted"]=len(intercepted)

    # RULE 3: atomic parallel — any bad call blocks ALL execution
    if intercepted:
        if executable: _bump("partial_parallel_blocked", len(executable))  # valid calls we are NOT executing
        info["executable"]=0
        clar=_build_clar(intercepted)
        if log: _log_intercept(source, intercepted, text)
        return [], info, norm, clar

    info["executable"]=len(executable)
    return executable, info, norm, None

if __name__=="__main__":
    WX={"type":"function","function":{"name":"get_weather","parameters":{"type":"object","properties":{"city":{"type":"string"}},"required":["city"]}}}
    tests=[
      ('<think>x</think>\n<tool_call>\n{"name":"get_weather","arguments":{"city":"Tokyo"}}\n</tool_call>', "strict"),
      ('<think>x</think>\n{"name":"get_weather","arguments":{"city":"Tokyo"}}', "lenient"),
      ('<think>x</think>\n{"name":"a","arguments":{"k":1}}\n{"name":"b","arguments":{"k":2}}', "lenient(2)"),
      ('[get_weather(city="Tokyo", unit="c"), get_weather(city="London")]', "lenient_fn(2)"),
      ('The capital of France is Paris.', "none"),
    ]
    ok=0
    for t,label in tests:
        calls,info,norm=normalize(t,log=False)
        wrapped_ok = ("<tool_call>" in norm) if info["n_calls"] else True
        print(f"[{label:11s}] strict={info['strict_parse_pass']} lenient={info['lenient_recovered']} n={info['n_calls']} wrapped_ok={wrapped_ok}")
        ok += wrapped_ok
    # validator + safety tests
    ex,_,_,clar = serve('{"name":"get_weather","arguments":{"city":"Tokyo"}}', [WX], log=False)
    v1=(len(ex)==1 and clar is None);                         print(f"[validate ok    ] exec={len(ex)} clar={clar!r} -> {'OK' if v1 else 'FAIL'}")
    ex,_,_,clar = serve('{"name":"get_weather","arguments":{}}', [WX], log=False)
    v2=(len(ex)==0 and clar and "city" in clar);             print(f"[missing arg    ] exec={len(ex)} clar={clar!r} -> {'OK' if v2 else 'FAIL'}")
    ex,_,_,clar = serve('{"name":"launch_missiles","arguments":{"x":1}}', [WX], log=False)
    v3=(len(ex)==0 and clar and "available" in clar);        print(f"[unknown tool   ] exec={len(ex)} clar={clar!r} -> {'OK' if v3 else 'FAIL'}")
    ex,_,_,clar = serve('{"name":"get_weather","arguments":{"city":"Tokyo"}}', [], log=False)
    v4=(len(ex)==0 and clar and "schema is missing" in clar);print(f"[no schema      ] exec={len(ex)} clar={clar!r} -> {'OK' if v4 else 'FAIL'}")
    ex,info,_,clar = serve('{"name":"get_weather","arguments":{"city":"Tokyo"}}\n{"name":"get_weather","arguments":{}}', [WX], log=False)
    v5=(len(ex)==0 and clar and "city" in clar and info["intercepted"]==1);print(f"[atomic parallel] exec={len(ex)} intercepted={info['intercepted']} -> {'OK' if v5 else 'FAIL'}")
    ok += v1+v2+v3+v4+v5
    print("COUNTERS:", get_counters())
    print(f"SELFTEST {ok}/{len(tests)+5} ok")