#!/usr/bin/env python3
"""Dump a finished Open WebUI chat - full chain of thought, every tool call and
result, and the final answer - into a prompt folder's round directory.

The transcript is read straight from Open WebUI's SQLite database inside the
container, not from the browser. Open WebUI 0.11 stores each assistant turn as
an ordered `output` list of items (reasoning / function_call /
function_call_output / message), so nothing has to be scraped or reconstructed.

Each prompt folder holds one directory per testing round - `v1/`, `v2/`, ... -
carrying that round's `result.txt` and `grade.md`. The transcript is written to
`<prompt_dir>/<round>/result.txt`; with no `--round`, the highest-numbered round
directory that exists is used, so a normal run needs no extra flag.

Usage:
    ./dump_result.py easy/e01-letter-counting            # match by prompt text
    ./dump_result.py easy/e01-letter-counting --chat-id <uuid>
    ./dump_result.py easy/e01-letter-counting --round v1 # grade an older round
    ./dump_result.py easy/e01-letter-counting --force    # overwrite a filled result.txt
    ./dump_result.py --list 20                           # show recent chats to pick from

Matching by prompt text finds the most recent chat whose FIRST user message is
exactly the contents of prompt.txt (whitespace-trimmed), so paste the prompt
verbatim. Multi-turn chats (the model asked a question, you answered) are
dumped in full along the branch that is currently selected in the UI.
"""

import argparse
import json
import os
import re
import subprocess
import sys

CONTAINER = os.environ.get("OWUI_CONTAINER", "open-webui")
DB_PATH = "/app/backend/data/webui.db"

# Runs INSIDE the container. Reads the DB read-only and prints the rendered
# transcript to stdout; the host side decides where it goes.
_INNER = r'''
import json, sqlite3, sys, datetime

args = json.loads(sys.stdin.read())
con = sqlite3.connect("file:%s?mode=ro" % args["db"], uri=True)

def utc(ts):
    if not ts: return "?"
    if ts > 1e12: ts = ts / 1000.0
    return datetime.datetime.fromtimestamp(ts, datetime.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")

def first_user_text(chat):
    msgs = (chat.get("history") or {}).get("messages") or {}
    roots = [m for m in msgs.values() if m.get("role") == "user" and not m.get("parentId")]
    roots.sort(key=lambda m: m.get("timestamp") or 0)
    return (roots[0].get("content") or "").strip() if roots else ""

if args.get("list"):
    rows = con.execute("select c.id, c.title, c.updated_at, u.name, c.chat from chat c left join user u on u.id=c.user_id order by c.updated_at desc limit ?", (args["list"],)).fetchall()
    for cid, title, upd, uname, chat in rows:
        chat = json.loads(chat)
        model = ",".join(chat.get("models") or [])
        head = first_user_text(chat)[:70].replace("\n", " ")
        print("%s  %s  %-8s %-20s %-40s | %s" % (cid, utc(upd), (uname or "?")[:8], model[:20], (title or "")[:40], head))
    sys.exit(0)

if args.get("chat_id"):
    row = con.execute("select c.id, c.title, c.updated_at, u.name, c.chat from chat c left join user u on u.id=c.user_id where c.id=?", (args["chat_id"],)).fetchone()
    if not row:
        print("ERROR: no chat with id %s" % args["chat_id"], file=sys.stderr); sys.exit(2)
else:
    want = args["prompt"].strip()
    row = None
    for r in con.execute("select c.id, c.title, c.updated_at, u.name, c.chat from chat c left join user u on u.id=c.user_id order by c.updated_at desc").fetchall():
        if first_user_text(json.loads(r[4])) == want:
            row = r; break
    if not row:
        print("ERROR: no chat whose first user message equals prompt.txt (paste the prompt verbatim, or pass --chat-id)", file=sys.stderr); sys.exit(2)

cid, title, upd, uname, chat = row
chat = json.loads(chat)
hist = chat.get("history") or {}
msgs = hist.get("messages") or {}

# Walk the selected branch from currentId back to the root, then reverse.
chain = []
mid = hist.get("currentId")
seen = set()
while mid and mid in msgs and mid not in seen:
    seen.add(mid); chain.append(msgs[mid]); mid = msgs[mid].get("parentId")
chain.reverse()

out = []
def sec(title): out.append("\n=== %s ===" % title)

def parts_text(parts):
    if isinstance(parts, str): return parts
    bits = []
    for p in parts or []:
        if not isinstance(p, dict): bits.append(str(p)); continue
        t = p.get("type") or ""
        if "image" in t:
            u = p.get("image_url") or p.get("url") or ""
            bits.append("[image: %s, %d chars]" % (u[:30], len(u)))
        elif p.get("text") is not None:
            bits.append(p["text"])
        else:
            bits.append(json.dumps(p)[:2000])
    return "\n".join(bits)

tool_counts = {}
first_model = None
for m in chain:
    if m.get("role") == "assistant" and not first_model:
        first_model = m.get("model")
    for o in m.get("output") or []:
        if o.get("type") == "function_call":
            tool_counts[o.get("name")] = tool_counts.get(o.get("name"), 0) + 1

hdr = []
hdr.append("#" * 72)
hdr.append("# owuitools prompt test result")
hdr.append("# prompt dir : %s" % args.get("label", "?"))
hdr.append("# chat id    : %s" % cid)
hdr.append("# chat title : %s" % (title or ""))
hdr.append("# user       : %s" % (uname or "?"))
hdr.append("# model      : %s" % (first_model or ",".join(chat.get("models") or [])))
hdr.append("# turns      : %d messages on selected branch" % len(chain))
hdr.append("# tool calls : %d  %s" % (sum(tool_counts.values()), ", ".join("%s x%d" % kv for kv in sorted(tool_counts.items(), key=lambda kv: -kv[1])) or "(none)"))
hdr.append("# dumped at  : %s" % utc(datetime.datetime.now(datetime.timezone.utc).timestamp()))
hdr.append("#" * 72)
out.append("\n".join(hdr))

for m in chain:
    role = m.get("role")
    ts = utc(m.get("timestamp"))
    if role == "user":
        sec("USER  (%s)" % ts)
        out.append((m.get("content") or "").rstrip())
        for f in m.get("files") or []:
            out.append("[attached file: %s]" % (f.get("name") or f.get("url") or "?"))
        continue

    sec("ASSISTANT  (%s, model %s)" % (ts, m.get("model") or "?"))
    u = m.get("usage") or {}
    if u:
        bits = []
        for k in ("prompt_eval_count", "eval_count", "prompt_tokens", "completion_tokens"):
            if k in u: bits.append("%s=%s" % (k, u[k]))
        if "total_duration" in u: bits.append("total=%.1fs" % (u["total_duration"] / 1e9))
        if "prompt_eval_duration" in u: bits.append("prefill=%.1fs" % (u["prompt_eval_duration"] / 1e9))
        if "eval_duration" in u: bits.append("decode=%.1fs" % (u["eval_duration"] / 1e9))
        if "response_token/s" in u: bits.append("%.1f tok/s" % u["response_token/s"])
        out.append("usage: " + " ".join(bits))

    items = m.get("output") or []
    if not items:
        out.append("(no structured output items - raw content follows)")
        out.append((m.get("content") or "").rstrip())
    else:
        n = 0
        calls = {}
        for o in items:
            n += 1
            t = o.get("type")
            if t == "reasoning":
                dur = o.get("duration")
                sec("[%d] THINKING%s" % (n, " (%ss)" % dur if dur is not None else ""))
                out.append(parts_text(o.get("content")).rstrip())
            elif t == "function_call":
                calls[o.get("call_id")] = o.get("name")
                sec("[%d] TOOL CALL: %s  (%s)" % (n, o.get("name"), o.get("call_id")))
                a = o.get("arguments")
                try:
                    a = json.dumps(json.loads(a), indent=2, ensure_ascii=False) if isinstance(a, str) else json.dumps(a, indent=2, ensure_ascii=False)
                except Exception:
                    pass
                out.append(str(a))
            elif t == "function_call_output":
                sec("[%d] TOOL RESULT: %s  (%s)" % (n, calls.get(o.get("call_id"), "?"), o.get("call_id")))
                out.append(parts_text(o.get("output")).rstrip())
            elif t == "message":
                sec("[%d] ASSISTANT TEXT" % n)
                out.append(parts_text(o.get("content")).rstrip())
            else:
                sec("[%d] %s" % (n, t))
                out.append(json.dumps(o, ensure_ascii=False)[:4000])

    sh = m.get("statusHistory") or []
    if sh:
        sec("STATUS HISTORY (live progress indicator)")
        for s in sh:
            out.append("- %s" % (s.get("description") or ""))

print("\n".join(out))
'''


def latest_round(prompt_dir):
    """Highest-numbered vN directory in a prompt folder ('v2' beats 'v10'? no - numeric)."""
    rounds = []
    for name in os.listdir(prompt_dir):
        m = re.fullmatch(r"v(\d+)", name)
        if m and os.path.isdir(os.path.join(prompt_dir, name)):
            rounds.append((int(m.group(1)), name))
    if not rounds:
        sys.exit("no round directory (v1, v2, ...) in %s - see README.md" % prompt_dir)
    return max(rounds)[1]


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("prompt_dir", nargs="?", help="folder containing prompt.txt and the round directories")
    ap.add_argument("--round", metavar="vN", help="round directory to write into (default: the highest vN present)")
    ap.add_argument("--chat-id", help="dump this chat instead of matching prompt.txt")
    ap.add_argument("--force", action="store_true", help="overwrite a non-empty result.txt")
    ap.add_argument("--stdout", action="store_true", help="print instead of writing result.txt")
    ap.add_argument("--list", type=int, metavar="N", help="list the N most recent chats and exit")
    a = ap.parse_args()

    payload = {"db": DB_PATH}
    if a.list:
        payload["list"] = a.list
    else:
        if not a.prompt_dir:
            ap.error("prompt_dir is required unless --list is given")
        prompt_path = os.path.join(a.prompt_dir, "prompt.txt")
        rnd = a.round or latest_round(a.prompt_dir)
        round_dir = os.path.join(a.prompt_dir, rnd)
        if not os.path.isdir(round_dir):
            sys.exit("no such round directory: %s" % round_dir)
        result_path = os.path.join(round_dir, "result.txt")
        with open(prompt_path, encoding="utf-8") as f:
            payload["prompt"] = f.read()
        payload["label"] = "%s (round %s)" % (
            os.path.relpath(a.prompt_dir, os.path.dirname(os.path.abspath(__file__))), rnd)
        if a.chat_id:
            payload["chat_id"] = a.chat_id
        if not a.stdout and not a.force and os.path.exists(result_path) and os.path.getsize(result_path) > 0:
            sys.exit("refusing to overwrite non-empty %s (use --force)" % result_path)

    proc = subprocess.run(
        ["docker", "exec", "-i", CONTAINER, "python3", "-c", _INNER],
        input=json.dumps(payload), capture_output=True, text=True,
    )
    if proc.returncode != 0:
        sys.stderr.write(proc.stderr)
        sys.exit(proc.returncode)
    if a.list or a.stdout:
        sys.stdout.write(proc.stdout)
        return
    with open(result_path, "w", encoding="utf-8") as f:
        f.write(proc.stdout)
    lines = proc.stdout.count("\n")
    print("wrote %s (%d lines, %d bytes)" % (result_path, lines, len(proc.stdout.encode())))


if __name__ == "__main__":
    main()