"""Score a line of conversation. Needs torch, transformers, safetensors.

  python predict.py "I'll send the deck by Friday."
  python predict.py --context "A: can someone own the budget deck?" --speaker B "sure, I'll do it."
  echo '{"text": "...", "context": [{"speaker": "A", "text": "..."}], "speaker": "B"}' | python predict.py --jsonl

--context is repeatable, oldest turn first; only the last 3 are used. --model is the model folder
(default: the folder this file is in).
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

from cue import CueModel, parse_turn


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default=str(Path(__file__).resolve().parent))
    ap.add_argument("--context", action="append", default=[], help='"SPEAKER: text", oldest first; repeatable')
    ap.add_argument("--speaker", default=None, help="speaker of the target line (needed to tell same/other speaker)")
    ap.add_argument("--jsonl", action="store_true", help="read {text, context?, speaker?} rows from stdin")
    ap.add_argument("--device", default="cpu")
    ap.add_argument("text", nargs="*")
    a = ap.parse_args()

    m = CueModel.from_pretrained(a.model, device=a.device)
    if a.jsonl:
        rows = [json.loads(line) for line in sys.stdin if line.strip()]
    else:
        ctx = [parse_turn(c) for c in a.context]
        rows = [{"text": t, "context": ctx, "speaker": a.speaker} for t in a.text]
    if not rows:
        ap.error("give a text argument or --jsonl")
    heads = m.product_heads
    for r, s in zip(rows, m.score_many(rows)):
        fired = m.decide(s)
        print(json.dumps({"text": r["text"], **{h: round(s[h], 4) for h in heads},
                          "fired": [h for h in heads if fired.get(h)]}))


if __name__ == "__main__":
    main()