#!/usr/bin/env python3
"""Ask Mouse-1 one four-option question about a task, and record the answer.

    bin/triage AUD-171            # what does this task need next
    bin/triage AUD-171 --ask redo # a different question set

Mouse-1 answers a typed multiple-choice question about a state you describe. It
does not write code, run tools, or decide whether work is finished. It labels a
situation, and the label is written on the task so a person can see why a machine
suggested something.

The cold start is real: the first call after an idle period answers 503 after
about thirty seconds, and the retry succeeds. This retries exactly once, says so,
and never pretends the first call worked.
"""

import json
import os
import sys
import time
import urllib.error
import urllib.request

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "lib"))
import harness as h  # noqa: E402

ENDPOINT = "https://api.mouse.dev/v1/decisions"

QUESTIONS = {
    "next": {
        "instructions": "This Paperclip task has reached the state described. What does it need next?",
        "criteria": {
            "merge": "the work is complete, its checks pass, and it should be merged now",
            "person": "a person has to decide or look at it; a machine should not continue",
            "redo": "the work is wrong or incomplete, and the same agent should try again",
            "environment": "something outside the code is wrong: a missing tool, a service, a path or a permission",
        },
    },
    "failure": {
        "instructions": "This run failed. What does the failure need?",
        "criteria": {
            "redo": "the work is incomplete or wrong; another attempt with the same agent is likely to succeed",
            "stronger_model": "the task needs more capability than this model has",
            "environment": "something outside the code is wrong: a missing tool, a service, a path",
            "person": "a decision only the person can make",
        },
    },
}


def mouse_key():
    key = os.environ.get("MOUSE_API_KEY", "")
    if key:
        return key
    path = os.path.expanduser("~/.config/environment.d/10-secrets.conf")
    if os.path.isfile(path):
        for line in open(path):
            if line.startswith("MOUSE_API_KEY="):
                return line.split("=", 1)[1].strip().strip('"').strip("'")
    return ""


def ask(state, question):
    key = mouse_key()
    if not key:
        raise h.Refused(2, "no MOUSE_API_KEY in the environment or ~/.config/environment.d/10-secrets.conf")
    body = json.dumps({"model": "mouse-1", "state": state,
                       "questions": {"q": dict(type="choice", **question)}}).encode()
    for attempt in (1, 2):
        started = time.time()
        req = urllib.request.Request(ENDPOINT, data=body, method="POST",
                                     headers={"Authorization": "Bearer " + key,
                                              "Content-Type": "application/json"})
        try:
            with urllib.request.urlopen(req, timeout=180) as r:
                return json.loads(r.read().decode()), time.time() - started, attempt
        except urllib.error.HTTPError as e:
            detail = e.read().decode()[:300]
            if e.code == 503 and attempt == 1:
                h.say("  the model was idle; it takes minutes to start. retrying once…")
                continue
            raise h.Refused(3, "Mouse-1 answered %d: %s" % (e.code, detail))


def main():
    args = [a for a in sys.argv[1:] if not a.startswith("--")]
    if not args:
        h.say("usage: bin/triage <task> [--ask next|failure]")
        return 2
    which = "next"
    if "--ask" in sys.argv:
        which = sys.argv[sys.argv.index("--ask") + 1]
    if which not in QUESTIONS:
        raise h.Refused(2, "no question set called %r" % which)

    issue = h.find_issue(args[0])
    worktree, branch, base = (None, "", "")
    try:
        worktree, branch, base = h.worktree_of(issue)
    except h.Refused:
        pass

    runs = h.api("GET", "/api/issues/%s/runs" % issue["id"])
    last = runs[0] if runs else {}
    state = {
        "task": issue["identifier"] + ": " + issue["title"],
        "description": (issue.get("description") or "")[:1500],
        "status": issue.get("status"),
        "last_run": {k: last.get(k) for k in ("status", "invocationSource", "errorCode", "finishedAt")},
        "worktree": worktree or "none",
        "branch": branch or "none",
    }
    if not last.get("errorCode") and last.get("status") in ("failed", "cancelled"):
        last_status = last.get("status")
        state["last_run"]["status"] = last_status

    h.say("%s %s" % (issue["identifier"], issue["title"]))
    h.say("  asking Mouse-1: %s" % which)
    answer, seconds, attempt = ask(state, QUESTIONS[which])
    a = (answer.get("answers") or {}).get("q") or {}
    choice = a.get("choice", "?")
    confidence = a.get("confidence", 0)
    probs = a.get("probabilities") or {}

    h.say("  answer: %s (confidence %.2f, %.1fs, attempt %d)" % (choice, confidence, seconds, attempt))
    for option, p in sorted(probs.items(), key=lambda kv: -kv[1]):
        h.say("    %-14s %.2f  %s" % (option, p, QUESTIONS[which]["criteria"].get(option, "")))

    lines = ["## Triage (Mouse-1)", "",
             "Question: %s" % QUESTIONS[which]["instructions"], "",
             "- answer: **%s**" % choice,
             "- confidence: %.2f" % confidence,
             "- took %.1fs on attempt %d" % (seconds, attempt), "",
             "| option | probability | what it means |",
             "|---|---|---|"]
    for option, p in sorted(probs.items(), key=lambda kv: -kv[1]):
        lines.append("| `%s` | %.2f | %s |" % (option, p, QUESTIONS[which]["criteria"].get(option, "")))
    lines += ["", "The model labelled the situation. It did not change anything, and it does not decide a merge."]
    h.comment(issue["id"], "\n".join(lines))
    h.say("  written on the task as a comment")
    return 0


if __name__ == "__main__":
    h.main_guard(main)