#!/usr/bin/env python3
"""Sort a cannabis wholesale inbox into act / confirm / human with TypeSafe's Jev.

    python sort.py sample-messages.csv --dry-run     # prints what would be sent; no key needed
    python sort.py sample-messages.csv               # live; needs TYPESAFE_API_KEY and `pip install typesafe-sdk`
    python sort.py --self-test                       # exercises route() on canned answers; no key needed

CSV columns: from,text. One request per row: the row is the state, questions.json
is the question pack. Every question is evaluated in the same call. Code below
decides the lane; the model only answers. "act" never writes to your ERP or Metrc.
CC0.
"""
import argparse
import csv
import json
import os
import sys
from pathlib import Path

# ---- thresholds: edit these, then rerun. Start here, tune on your own traffic. ----
NEEDS_HUMAN = 0.50        # needs_human noul at/above this -> human
HEALTH_BLOCK = 0.70       # health_claim noul at/above this -> human (do not publish)
HEALTH_REVIEW = 0.35      # health_claim noul at/above this -> confirm (a person reads it)
INTENT_UNSURE = 0.50      # intent confidence below this -> human (we do not know what it is)
INTENT_ACT = 0.75         # intent confidence at/above this -> act, else confirm
IS_ORDER_ACT = 0.85       # is_order noul at/above this (plus INTENT_ACT) -> draft the order
URGENT = 1.5              # urgency score (0..2) at/above this flags the row
MODEL = "jev-latest"      # pin a versioned id (see GET /v1/models) once thresholds are tuned
PRICE_PER_MTOK = 0.042    # USD per million input tokens, TypeSafe models page, Sept 2026
# -----------------------------------------------------------------------------------

HERE = Path(__file__).resolve().parent
QUESTIONS = {k: v for k, v in json.loads((HERE / "questions.json").read_text()).items() if not k.startswith("_")}


def read_messages(path):
    with open(path, newline="", encoding="utf-8") as f:
        rows = list(csv.DictReader(f))
    for r in rows:
        if "from" not in r or "text" not in r:
            sys.exit("CSV needs columns: from,text")
    return rows


def make_state(row):
    """The state: what the model judges. All five questions see this same object."""
    return {"from": row["from"], "message": row["text"]}


def route(a):
    """a: dict of question id -> answer with .noul / .choice / .confidence / .score.
    Returns (lane, reason). The answer picks the direction, confidence picks the lane."""
    if a["needs_human"].noul >= NEEDS_HUMAN:
        return "human", "needs a person"
    if a["health_claim"].noul >= HEALTH_BLOCK:
        return "human", "health claim; do not publish"
    if a["health_claim"].noul >= HEALTH_REVIEW:
        return "confirm", "possible health claim; a person reads it"
    intent = a["intent"]
    if intent.confidence < INTENT_UNSURE:
        return "human", "unsure what this is"
    if intent.choice in ("order", "reorder"):
        if a["is_order"].noul >= IS_ORDER_ACT and intent.confidence >= INTENT_ACT:
            return "act", "draft %s for rep approval" % intent.choice
        return "confirm", "looks like a%s order; confirm quantities" % (" repeat" if intent.choice == "reorder" else "n")
    if intent.choice == "complaint":
        return "human", "complaint"
    action = {"question": "route to rep", "vendor": "file to purchasing", "other": "archive"}[intent.choice]
    if intent.confidence >= INTENT_ACT:
        return "act", action
    return "confirm", "probably %s; confirm" % intent.choice


def print_table(rows, answers):
    print("%-3s %-8s %-3s %-9s %-5s %-5s %-6s %-34s %s" % ("#", "lane", "urg", "intent", "conf", "order", "health", "reason", "from / text"))
    for i, (row, a) in enumerate(zip(rows, answers), 1):
        lane, reason = route(a)
        urgent = "!" if a["urgency"].score >= URGENT else ""
        print("%-3d %-8s %-3s %-9s %-5.2f %-5.2f %-6.2f %-34s %s | %s" % (
            i, lane, urgent, a["intent"].choice, a["intent"].confidence, a["is_order"].noul,
            a["health_claim"].noul, reason[:34], row["from"], row["text"][:48]))


def dry_run(rows):
    print("DRY RUN: nothing sent. Each request below would go to POST https://api.typesafe.ai/v1/systemone\n")
    for i, row in enumerate(rows, 1):
        body = {"model": MODEL, "state": make_state(row), "questions": QUESTIONS}
        print("--- request %d of %d ---" % (i, len(rows)))
        print(json.dumps(body, indent=2)[:900] + ("\n  ... (%d question keys: %s)" % (len(QUESTIONS), ", ".join(QUESTIONS))))
    print("\n%d requests, %d questions each. Thresholds: needs_human>=%.2f human; health>=%.2f human, >=%.2f confirm; "
          "intent conf<%.2f human, >=%.2f act; is_order>=%.2f draft; urgency>=%.1f flag." % (
              len(rows), len(QUESTIONS), NEEDS_HUMAN, HEALTH_BLOCK, HEALTH_REVIEW, INTENT_UNSURE, INTENT_ACT, IS_ORDER_ACT, URGENT))
    print("Set TYPESAFE_API_KEY and drop --dry-run to run it live.")


def live(rows):
    if not os.environ.get("TYPESAFE_API_KEY"):
        sys.exit("TYPESAFE_API_KEY is not set. Use --dry-run to see what would be sent.")
    try:
        from typesafe_sdk import TypeSafeClient  # pip install typesafe-sdk  (Python >= 3.10)
    except ImportError:
        sys.exit("pip install typesafe-sdk   (or run with --dry-run)")
    answers, tokens = [], 0
    with TypeSafeClient(model=MODEL) as client:   # reads TYPESAFE_API_KEY; retries 429/529 with backoff by default
        for row in rows:
            r = client.system_one(make_state(row), QUESTIONS)   # raw question dicts are accepted as-is
            answers.append(r.answers)
            tokens += r.usage.input_tokens or 0
            served = r.model
    print_table(rows, answers)
    print("\n%d messages, %d input tokens, about $%.4f at $%.3f/Mtok (model served: %s)" % (
        len(rows), tokens, tokens / 1e6 * PRICE_PER_MTOK, PRICE_PER_MTOK, served))


class _A:  # stand-in answer object for --self-test
    def __init__(self, **kw):
        self.__dict__.update(kw)


def self_test():
    cases = [
        ("clean reorder", dict(needs_human=_A(noul=0.10), health_claim=_A(noul=0.02), intent=_A(choice="reorder", confidence=0.92),
                               is_order=_A(noul=0.97), urgency=_A(score=1.3)), "act"),
        ("vague order", dict(needs_human=_A(noul=0.20), health_claim=_A(noul=0.01), intent=_A(choice="order", confidence=0.80),
                             is_order=_A(noul=0.60), urgency=_A(score=0.4)), "confirm"),
        ("complaint", dict(needs_human=_A(noul=0.88), health_claim=_A(noul=0.01), intent=_A(choice="complaint", confidence=0.95),
                           is_order=_A(noul=0.05), urgency=_A(score=1.1)), "human"),
        ("menu health claim", dict(needs_human=_A(noul=0.30), health_claim=_A(noul=0.91), intent=_A(choice="other", confidence=0.70),
                                   is_order=_A(noul=0.01), urgency=_A(score=0.2)), "human"),
        ("unsure intent", dict(needs_human=_A(noul=0.10), health_claim=_A(noul=0.02), intent=_A(choice="question", confidence=0.41),
                               is_order=_A(noul=0.30), urgency=_A(score=0.5)), "human"),
        ("vendor quote", dict(needs_human=_A(noul=0.15), health_claim=_A(noul=0.01), intent=_A(choice="vendor", confidence=0.93),
                              is_order=_A(noul=0.03), urgency=_A(score=0.6)), "act"),
    ]
    failed = 0
    for name, answers, expected in cases:
        lane, reason = route(answers)
        ok = lane == expected
        failed += not ok
        print("%s %-18s -> %-8s (%s)%s" % ("ok " if ok else "FAIL", name, lane, reason, "" if ok else "  expected " + expected))
    sys.exit(1 if failed else 0)


def main():
    p = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    p.add_argument("csv", nargs="?", help="CSV with columns from,text")
    p.add_argument("--dry-run", action="store_true", help="print the requests; do not call the API")
    p.add_argument("--self-test", action="store_true", help="run route() on canned answers")
    args = p.parse_args()
    if args.self_test:
        self_test()
    if not args.csv:
        p.error("give a CSV path, or --self-test")
    rows = read_messages(args.csv)
    dry_run(rows) if args.dry_run else live(rows)


if __name__ == "__main__":
    main()