#!/usr/bin/env python3
"""kev - free, fast (~30 ms) yes/no and choice decisions from KEV (kev.osmike.com).

  kev yesno  "<state>" "<question>" [--exit] [--min P]
  kev ask    "<state>" -q id:"question" [-q id2:"question2" ...] [--json]
  kev choose "<state>" "<question>" opt:"description" opt2:"description" [...] [--json]

State: a text describing the situation (keep it < ~55k chars; longer is cut to the last 55k).
Use "-" as the state to read it from stdin, or --state-file FILE.

Output:
  yesno  -> "yes 0.82" or "no 0.31"  (the number is always p(yes); yes when p >= --min, default 0.5)
           --exit: exit code 0 = yes, 1 = no, 2 = error (for shell `if kev yesno ... --exit; then`)
  ask    -> one line per question: "<id> yes 0.82"
  choose -> first line "<choice> <p>" (+ " near-tie" when the runner-up is within 0.15),
            then every option with its probability, best first.
Tips: KEV is sharper on several yes/no questions (`kev ask`) than on one 3-way choice.
With near-ties, look at the probabilities and apply your own margin.

Token: own revocable GUEST_GATEWAY_KEY; host injects the KEV credential.
URL:   $OSMIKE_KEV_URL (default https://kev.osmike.com). Timeout 20 s. Exit 2 on any error.
"""
import argparse, json, os, re, sys, urllib.error, urllib.request

DEFAULT_URL = "http://172.30.240.1:8770/kev"

MAX_STATE = 55000
TIMEOUT = 20
ID_RE = re.compile(r"^[A-Za-z0-9_.-]{1,40}$")
NEAR_TIE = 0.15


class KevError(Exception):
    pass


def _cfg(name):
    return 'http://172.30.240.1:8770/kev' if name == 'OSMIKE_KEV_URL' else ''


def get_token():
    token=os.environ.get('GUEST_GATEWAY_KEY','')
    if not token: raise KevError('missing guest gateway key')
    return token


def call_kev(state, questions):
    url = (_cfg("OSMIKE_KEV_URL") or DEFAULT_URL).rstrip("/") + "/v1/systemone"
    if len(state) > MAX_STATE:
        print(f"kev: state is {len(state)} chars, keeping the last {MAX_STATE}", file=sys.stderr)
        state = state[-MAX_STATE:]
    if not state.strip():
        raise KevError("empty state")
    body = json.dumps({"model": "kev-latest", "state": state, "questions": questions}).encode()
    req = urllib.request.Request(url, data=body, method="POST", headers={
        "Authorization": "Bearer " + get_token(), "Content-Type": "application/json"})
    try:
        with urllib.request.urlopen(req, timeout=TIMEOUT) as r:
            data = json.loads(r.read())
    except urllib.error.HTTPError as e:
        try:
            detail = e.read().decode("utf-8", "replace")[:300]
        except Exception:
            detail = ""
        hint = {401: "token rejected", 422: "bad request (only noul/choice questions are valid)",
                502: "KEV backend down (GPU box/tunnel)", 503: "KEV unavailable", 504: "KEV timeout"}
        raise KevError(f"KEV HTTP {e.code} ({hint.get(e.code, 'error')}) {detail}".strip())
    except Exception as e:
        raise KevError(f"KEV unreachable ({url}): {e.__class__.__name__}: {e}")
    answers = data.get("answers")
    if not isinstance(answers, dict):
        raise KevError("KEV reply has no answers")
    return answers


# ------------------------------------------------------------ parsing ---
def split_spec(spec, n):
    """'id:text' -> (id, text). The id must be followed directly by the text (no space after ':'),
    so a question like "Note: is it broken?" keeps its colon and gets an automatic id q<n>."""
    if ":" in spec:
        head, tail = spec.split(":", 1)
        if ID_RE.match(head) and tail.strip() and not tail[:1].isspace():
            return head, tail.strip()
    return f"q{n}", spec.strip()


def parse_questions(specs):
    qs, seen = {}, set()
    for i, s in enumerate(specs, 1):
        qid, text = split_spec(s, i)
        if not text:
            raise KevError(f"empty question in {s!r}")
        if qid in seen:
            raise KevError(f"duplicate question id {qid!r}")
        seen.add(qid)
        qs[qid] = {"type": "noul", "instructions": text}
    if not qs:
        raise KevError("no questions (use -q id:\"question\")")
    return qs


def parse_options(specs):
    crit = {}
    for i, s in enumerate(specs, 1):
        if ":" in s:
            name, desc = s.split(":", 1)
            name, desc = name.strip(), desc.strip()
        else:
            name, desc = s.strip(), s.strip()
        if not name:
            raise KevError(f"empty option name in {s!r}")
        if name in crit:
            raise KevError(f"duplicate option {name!r}")
        crit[name] = desc or name
    if len(crit) < 2:
        raise KevError("choose needs at least 2 options (opt:\"description\")")
    return crit


def read_state(arg, state_file):
    if state_file:
        try:
            return open(state_file, encoding="utf-8", errors="replace").read()
        except OSError as e:
            raise KevError(f"cannot read state file: {e}")
    if arg == "-":
        return sys.stdin.read()
    if arg is None:
        raise KevError("missing state")
    return arg


def noul_p(ans):
    try:
        return float(ans["noul"])
    except Exception:
        raise KevError(f"unexpected KEV answer: {json.dumps(ans)[:200]}")


def fmt_yes(p, thr):
    return f"{'yes' if p >= thr else 'no'} {p:.2f}"


def rank_choice(ans):
    probs = ans.get("probabilities") or {}
    ranked = sorted(((k, float(v)) for k, v in probs.items()), key=lambda kv: -kv[1])
    if not ranked and ans.get("choice"):
        ranked = [(ans["choice"], float(ans.get("confidence") or 0))]
    if not ranked:
        raise KevError(f"unexpected KEV answer: {json.dumps(ans)[:200]}")
    return ranked


# --------------------------------------------------------------- main ---
def build_parser():
    p = argparse.ArgumentParser(prog="kev", description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = p.add_subparsers(dest="cmd", required=True)

    y = sub.add_parser("yesno", help="one yes/no question")
    y.add_argument("state", nargs="?")
    y.add_argument("question", nargs="?")
    y.add_argument("--exit", action="store_true", help="exit 0 = yes, 1 = no, 2 = error")
    y.add_argument("--min", type=float, default=0.5, help="p(yes) threshold for 'yes' (default 0.5)")
    y.add_argument("--json", action="store_true")
    y.add_argument("--state-file")

    a = sub.add_parser("ask", help="several yes/no questions at once")
    a.add_argument("state", nargs="?")
    a.add_argument("-q", "--question", action="append", default=[], metavar='ID:"QUESTION"')
    a.add_argument("--min", type=float, default=0.5)
    a.add_argument("--json", action="store_true")
    a.add_argument("--state-file")

    c = sub.add_parser("choose", help="pick one of several options")
    c.add_argument("state", nargs="?")
    c.add_argument("question", nargs="?")
    c.add_argument("options", nargs="*", metavar='OPT:"DESCRIPTION"')
    c.add_argument("--json", action="store_true")
    c.add_argument("--state-file")
    return p


def main(argv=None):
    p = build_parser()
    args = p.parse_args(argv)
    # with --state-file the positional "state" slot holds the next argument
    if args.state_file:
        pos = [args.state] + ([args.question] if hasattr(args, "question") else []) + list(getattr(args, "options", []))
        pos = [x for x in pos if x is not None]
        args.state = None
        if args.cmd == "yesno":
            args.question = pos[0] if pos else None
        elif args.cmd == "choose":
            args.question = pos[0] if pos else None
            args.options = pos[1:]
        elif pos:
            raise KevError(f"unexpected argument {pos[0]!r} (use -q)")
    try:
        state = read_state(args.state, args.state_file)
        if args.cmd == "yesno":
            if not args.question:
                raise KevError('usage: kev yesno "<state>" "<question>"')
            ans = call_kev(state, {"q": {"type": "noul", "instructions": args.question}})
            pv = noul_p(ans.get("q") or {})
            yes = pv >= args.min
            print(json.dumps({"answer": "yes" if yes else "no", "p_yes": round(pv, 4)}) if args.json else fmt_yes(pv, args.min))
            return (0 if yes else 1) if args.exit else 0
        if args.cmd == "ask":
            qs = parse_questions(args.question)
            ans = call_kev(state, qs)
            out = {}
            for qid in qs:
                pv = noul_p(ans.get(qid) or {})
                out[qid] = {"answer": "yes" if pv >= args.min else "no", "p_yes": round(pv, 4)}
            if args.json:
                print(json.dumps(out))
            else:
                for qid, v in out.items():
                    print(f"{qid} {v['answer']} {v['p_yes']:.2f}")
            return 0
        if args.cmd == "choose":
            if not args.question:
                raise KevError('usage: kev choose "<state>" "<question>" opt:"desc" opt2:"desc"')
            crit = parse_options(args.options)
            ans = call_kev(state, {"c": {"type": "choice", "instructions": args.question, "criteria": crit}})
            ranked = rank_choice(ans.get("c") or {})
            tie = len(ranked) > 1 and ranked[0][1] - ranked[1][1] < NEAR_TIE
            if args.json:
                print(json.dumps({"choice": ranked[0][0], "p": round(ranked[0][1], 4), "near_tie": tie,
                                  "probabilities": {k: round(v, 4) for k, v in ranked}}))
            else:
                print(f"{ranked[0][0]} {ranked[0][1]:.2f}" + (" near-tie" if tie else ""))
                for k, v in ranked:
                    print(f"  {k} {v:.2f}")
            return 0
    except KevError as e:
        print(f"kev: {e}", file=sys.stderr)
        return 2
    except KeyboardInterrupt:
        return 2
    return 2


if __name__ == "__main__":
    try:
        sys.exit(main())
    except KevError as e:
        print(f"kev: {e}", file=sys.stderr)
        sys.exit(2)