"""Text renderings for prompts (mint packets and the zero-shot judge). Pure functions over items / sims / start DB."""
from __future__ import annotations

import json
import re

ORDER_RE = re.compile(r"#?W\d{7}")
RES_RE = re.compile(r"\b[A-Z0-9]{6}\b")
UID_RE = re.compile(r"\b[a-z]+_[a-z]+_\d{4}\b")
EMAIL_RE = re.compile(r"[\w.+-]+@[\w-]+\.[\w.]+")


def j(x, n=None):
    s = json.dumps(x, sort_keys=True, default=str)
    return s if n is None or len(s) <= n else s[:n] + f"...[+{len(s) - n} chars]"


def transcript(messages, tool_out=700, max_chars=24000, agent=True):
    lines = []
    for m in messages:
        r = m.get("role")
        if r == "user" and isinstance(m.get("content"), str) and m["content"].strip():
            lines.append(f"[user] {m['content'].strip()}")
        elif r == "assistant" and agent:
            if isinstance(m.get("content"), str) and m["content"].strip():
                lines.append(f"[agent] {m['content'].strip()}")
            for t in m.get("tool_calls") or []:
                lines.append(f"[agent tool call] {t['name']}({j(t.get('arguments'))})")
        elif r == "tool" and agent:
            c = m.get("content") or ""
            c = c if len(c) <= tool_out else c[:tool_out] + f"...[+{len(c) - tool_out} chars]"
            lines.append(f"[tool result{' ERROR' if m.get('error') else ''}] {c}")
    s = "\n".join(lines)
    if len(s) > max_chars:
        s = s[: max_chars // 2] + f"\n...[{len(s) - max_chars} chars of the middle omitted]...\n" + s[-max_chars // 2:]
    return s


def leaf_diff(before, after, path=""):
    out = []
    if isinstance(before, dict) and isinstance(after, dict):
        for k in sorted(set(before) | set(after), key=str):
            if before.get(k) != after.get(k):
                out += leaf_diff(before.get(k), after.get(k), f"{path}.{k}" if path else str(k))
    elif isinstance(before, list) and isinstance(after, list) and len(before) == len(after):
        for i, (a, b) in enumerate(zip(before, after)):
            if a != b:
                out += leaf_diff(a, b, f"{path}[{i}]")
    else:
        out.append((path, before, after))
    return out


def changes_text(changes, max_chars=8000):
    if not changes:
        return "(no database changes: the agent made no successful write)"
    lines = []
    for c in changes:
        head = f"{c['table']}[{c['key']!r}] {c['op'].upper()}"
        if c["op"] == "update":
            lines.append(head + ":")
            for p, a, b in leaf_diff(c["before"], c["after"]):
                lines.append(f"    {p}: {j(a, 400)} -> {j(b, 400)}")
        elif c["op"] == "insert":
            lines.append(head + ": " + j(c["after"], 1500))
        else:
            lines.append(head + " (record removed): " + j(c["before"], 600))
    s = "\n".join(lines)
    return s if len(s) <= max_chars else s[:max_chars] + f"\n...[+{len(s) - max_chars} chars]"


def _order_brief(o):
    return {"order_id": o["order_id"], "status": o["status"], "user_id": o["user_id"],
            "items": [{"name": i["name"], "item_id": i["item_id"], "product_id": i["product_id"], "price": i["price"],
                       "options": i["options"]} for i in o["items"]],
            "payment_history": o["payment_history"]}


def relevant_start(domain, start, user_text, changes, max_chars=14000):
    """Start-DB records relevant to this dialogue: touched records, their owners, ids mentioned by the user, the
    identified user's orders/reservations (brief), and products / flights of touched records."""
    parts = []
    users = set(); orders = set(); res = set(); products = set(); flights = set()
    for c in changes:
        rec = c["before"] or c["after"] or {}
        if c["table"] == "users":
            users.add(c["key"])
        if c["table"] == "orders":
            orders.add(c["key"]); users.add(rec.get("user_id"))
            for side in (c["before"] or {}, c["after"] or {}):
                for i in side.get("items") or []:
                    products.add(i.get("product_id"))
        if c["table"] == "reservations":
            res.add(c["key"]); users.add(rec.get("user_id"))
            for side in (c["before"] or {}, c["after"] or {}):
                for f in side.get("flights") or []:
                    flights.add((f.get("flight_number"), f.get("date")))
    for uid in UID_RE.findall(user_text):
        if uid in start.get("users", {}):
            users.add(uid)
    for em in EMAIL_RE.findall(user_text):
        for uid, u in start.get("users", {}).items():
            if u.get("email") == em:
                users.add(uid)
    if domain == "retail":
        for oid in ORDER_RE.findall(user_text):
            oid = oid if oid.startswith("#") else "#" + oid
            if oid in start["orders"]:
                orders.add(oid); users.add(start["orders"][oid]["user_id"])
    else:
        for rid in RES_RE.findall(user_text):
            if rid in start["reservations"]:
                res.add(rid); users.add(start["reservations"][rid]["user_id"])
    users.discard(None)
    for uid in sorted(users):
        u = start["users"].get(uid)
        if not u:
            continue
        parts.append(f"users[{uid!r}] = {j(u, 2500)}")
        if domain == "retail":
            for oid in u.get("orders", []):
                o = start["orders"].get(oid)
                if o:
                    parts.append(f"  orders[{oid!r}] (brief) = {j(_order_brief(o), 1200)}")
                    if oid in orders:
                        for i in o["items"]:
                            products.add(i["product_id"])
        else:
            for rid in u.get("reservations", []):
                r = start["reservations"].get(rid)
                if r:
                    parts.append(f"  reservations[{rid!r}] = {j(r, 1500)}")
                    if rid in res:
                        for f in r.get("flights") or []:
                            flights.add((f.get("flight_number"), f.get("date")))
    if domain == "retail":
        for pid in sorted(p for p in products if p):
            p = start["products"].get(pid)
            if p:
                parts.append(f"products[{pid!r}] = {j(p, 3000)}")
    else:
        for fn, d in sorted(f for f in flights if f[0]):
            fl = start["flights"].get(fn)
            if fl:
                brief = {k: v for k, v in fl.items() if k != "dates"}
                brief["dates"] = {d: fl.get("dates", {}).get(d)}
                parts.append(f"flights[{fn!r}] (date {d}) = {j(brief, 1200)}")
    s = "\n".join(parts) if parts else "(no records identified)"
    return s if len(s) <= max_chars else s[:max_chars] + f"\n...[+{len(s) - max_chars} chars]"