"""Ask Decider 2B questions from the terminal. Runs on a GPU, or on a CPU (slower).

  # A whole Jev/SystemOne request from a file (or - for stdin)
  decider --request request.json

  # Quick questions about a text or JSON file
  decider --state ticket.txt \
      --choice "team=Which team should handle this?:billing|Payments, charges and refunds,technical|App bugs and outages,account|Login and profile" \
      --yes-no "urgent=Is this urgent?" \
      --score "upset=How upset is the customer?:calm,frustrated,angry"

  # Many states at once (one per line, or a JSON list) -> one JSON answer per line
  decider --states tickets.txt --yes-no "urgent=Is this urgent?"

Describe choice options as key|description: the model is much more accurate when options say what
they mean (6-ticket routing check: 3/6 with bare keys, 5/6 with descriptions).

Options: --model (repo id or folder, default Jeetu456/decider-2b), --device (cuda, mps or cpu; default:
cuda if available), --full (print the whole response, not just the picks).
"""
import argparse
import json
import os
import sys


def read_state(path):
    text = sys.stdin.read() if path == "-" else open(path, encoding="utf-8").read()
    try:
        return json.loads(text)
    except ValueError:
        return text.strip()


def parse_questions(args):
    questions = {}
    for spec in args.choice or []:
        qid, rest = spec.split("=", 1)
        instructions, options = rest.rsplit(":", 1)
        criteria = {}
        for option in (o.strip() for o in options.split(",")):
            if option:
                key, _, description = option.partition("|")
                criteria[key.strip()] = description.strip() or None
        questions[qid] = {"type": "choice", "instructions": instructions, "criteria": criteria}
    for spec in args.yes_no or []:
        qid, instructions = spec.split("=", 1)
        questions[qid] = {"type": "noul", "instructions": instructions}
    for spec in args.score or []:
        qid, rest = spec.split("=", 1)
        instructions, levels = rest.rsplit(":", 1)
        questions[qid] = {"type": "score", "instructions": instructions,
                          "criteria": [level.strip() for level in levels.split(",") if level.strip()]}
    return questions


def picks(response):
    out = {}
    for qid, answer in response["answers"].items():
        if answer["type"] == "choice":
            out[qid] = {"choice": answer["choice"], "confidence": round(answer["confidence"], 3)}
        elif answer["type"] == "score":
            out[qid] = {"score": round(answer["score"], 2), "confidence": round(answer["confidence"], 3)}
        else:
            out[qid] = {"yes": answer["noul"] >= 0.5, "p_yes": round(answer["noul"], 3)}
    return out


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--model", default=os.environ.get("DECIDER_MODEL", "Jeetu456/decider-2b"))
    parser.add_argument("--device", default=None)
    parser.add_argument("--request", help="a Jev/SystemOne request JSON file, or - for stdin")
    parser.add_argument("--state", help="a text or JSON file, or - for stdin")
    parser.add_argument("--states", help="many states: one per line, or a JSON list")
    parser.add_argument("--choice", action="append", help='"id=question:key|description,key|description,..."')
    parser.add_argument("--yes-no", action="append", help='"id=question"')
    parser.add_argument("--score", action="append", help='"id=question:lowest,...,highest"')
    parser.add_argument("--full", action="store_true", help="print the full Jev response")
    args = parser.parse_args()

    import torch
    path = args.model
    if not os.path.isdir(path):
        from huggingface_hub import snapshot_download
        path = snapshot_download(path)
    from decider._core import load_release_model, systemone

    device = args.device or ("cuda" if torch.cuda.is_available() else "cpu")
    model = load_release_model(path, device=device)
    show = (lambda r: r) if args.full else picks

    if args.request:
        print(json.dumps(show(systemone(model, read_state(args.request))), indent=2))
        return
    questions = parse_questions(args)
    if not questions:
        parser.error("give --request, or at least one --choice, --yes-no or --score")
    if args.states:
        text = open(args.states, encoding="utf-8").read()
        try:
            states = json.loads(text)
        except ValueError:
            states = [line for line in text.splitlines() if line.strip()]
        for state in states:
            print(json.dumps(show(systemone(model, {"state": state, "questions": questions}))))
        return
    if not args.state:
        parser.error("give --state (or --states, or --request)")
    print(json.dumps(show(systemone(model, {"state": read_state(args.state), "questions": questions})), indent=2))


if __name__ == "__main__":
    main()