#!/usr/bin/env python3
"""Voice2Policy Text Router v0.1 — dependency-light inference (numpy + safetensors only, no pickle, no scikit-learn).

A character n-gram (char_wb, 1-4) TF-IDF + logistic-regression classifier for Korean business voice commands
(after ASR). It predicts the command frame (intent / action / kind) and whether a booking slot value (time or phone)
was mentioned. It does NOT resolve entity names, check tool arguments, or execute anything.

    python predict.py "내일 오후 세 시에 스케일링 예약해 주세요"
    python predict.py --file utterances.txt --json
"""
from __future__ import annotations

import argparse
import json
import re
import sys
from pathlib import Path

import numpy as np
from safetensors.numpy import load_file

HERE = Path(__file__).resolve().parent
_WS = re.compile(r"\s\s+")

# Default tool family for each action (the final tool also depends on the workspace state, e.g. which pending item to
# confirm, and whether a registered service/menu name was found). See README "From frame to tool".
TOOL_HINT = {"book": "update_slots", "order": "update_basket", "confirm": "confirm_{kind}", "decline": "decline_{kind}"}


class TextRouter:
    def __init__(self, model_dir: str | Path = HERE):
        d = Path(model_dir)
        self.cfg = json.loads((d / "config.json").read_text(encoding="utf-8"))
        self.vocab = json.loads((d / "vocab.json").read_text(encoding="utf-8"))
        w = load_file(str(d / "model.safetensors"))
        self.idf, self.W, self.b = w["tfidf.idf"], w["frame.coef"], w["frame.intercept"]
        self.Ws, self.bs = w["slot.coef"][0], float(w["slot.intercept"][0])
        self.labels = [tuple(x) for x in self.cfg["frame_labels"]]
        lo, hi = self.cfg["tfidf"]["ngram_range"]
        self.n_range = range(lo, hi + 1)

    def _ngrams(self, text: str):
        # same as scikit-learn TfidfVectorizer(analyzer="char_wb", lowercase=True)
        t = _WS.sub(" ", text.lower())
        out = []
        for w in t.split():
            w = " " + w + " "
            L = len(w)
            for n in self.n_range:
                off = 0
                out.append(w[off:off + n])
                while off + n < L:
                    off += 1
                    out.append(w[off:off + n])
                if off == 0:
                    break
        return out

    def features(self, text: str) -> np.ndarray:
        x = np.zeros(len(self.idf))
        for g in self._ngrams(text or ""):
            j = self.vocab.get(g)
            if j is not None:
                x[j] += 1.0
        nz = x > 0
        x[nz] = 1.0 + np.log(x[nz])          # sublinear tf
        x *= self.idf
        n = np.linalg.norm(x)
        return x / n if n > 0 else x

    def predict(self, text: str) -> dict:
        x = self.features(text)
        z = self.W @ x + self.b
        p = np.exp(z - z.max()); p /= p.sum()
        k = int(p.argmax())
        intent, action, kind = self.labels[k]
        ps = 1.0 / (1.0 + np.exp(-(self.Ws @ x + self.bs)))
        hint = TOOL_HINT.get(action)
        if hint and "{kind}" in hint:
            hint = hint.format(kind=kind) if kind else hint.replace("{kind}", "<pending>")
        return {"intent": intent, "action": action, "kind": kind, "frame_score": round(float(p[k]), 4),
                "frame_scores": {json.dumps(list(l), ensure_ascii=False): round(float(v), 4) for l, v in zip(self.labels, p)},
                "slot_mentioned": bool(ps >= 0.5), "slot_score": round(float(ps), 4),
                "tool_hint": "handoff" if intent == "delegate" else hint}


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("text", nargs="?")
    ap.add_argument("--file")
    ap.add_argument("--json", action="store_true")
    a = ap.parse_args()
    r = TextRouter()
    texts = [a.text] if a.text else [l.strip() for l in open(a.file, encoding="utf-8") if l.strip()]
    for t in texts:
        o = r.predict(t)
        if a.json:
            print(json.dumps({"text": t, **o}, ensure_ascii=False))
        else:
            print(f"{t}\n  -> intent={o['intent']} action={o['action']} kind={o['kind']} ({o['frame_score']:.3f}) "
                  f"slot_mentioned={o['slot_mentioned']} ({o['slot_score']:.3f}) tool_hint={o['tool_hint']}")
    return 0


if __name__ == "__main__":
    sys.exit(main())