"""
Solution: the prompt as a budget (Lab 07, 02_context_budget.py and 04_tool_results.py).
"""


def tokens(text: str) -> int:
    """A stand-in for a real tokenizer: one token per word."""
    return len(text.split())


def profile(parts: dict[str, str]) -> list[tuple[str, int, float]]:
    """(part, tokens, percent of the whole) for each part, biggest first."""
    total = sum(tokens(t) for t in parts.values())
    rows = [(name, tokens(t), round(100 * tokens(t) / total, 1)) for name, t in parts.items()]
    return sorted(rows, key=lambda r: -r[1])


def keep_the_end(parts: dict[str, str], budget: int) -> list[str]:
    """Join the parts in order, keep the last `budget` tokens; which parts survive, even partly?"""
    words = [(name, w) for name, t in parts.items() for w in t.split()]
    return list(dict.fromkeys(name for name, _ in words[-budget:])) if budget > 0 else []


def fit_by_priority(rules: str, question: str, memory: str, documents: list[str], history: list[str],
                    tools: str, budget: int) -> dict:
    kept = ["rules", "question"]
    used = tokens(rules) + tokens(question)

    def add(label: str, text: str) -> bool:
        nonlocal used
        if used + tokens(text) > budget:
            return False
        kept.append(label)
        used += tokens(text)
        return True

    add("memory", memory)
    for i, d in enumerate(documents, 1):
        if not add(f"doc {i}", d):
            break
    kept_msgs = 0
    for i in range(len(history), 0, -1):
        if not add(f"msg {i}", history[i - 1]):
            break
        kept_msgs += 1
    add("tools", tools)

    dropped = []
    if kept_msgs < len(history):
        dropped.append(f"{len(history) - kept_msgs} oldest messages")
    docs_kept = sum(1 for k in kept if k.startswith("doc "))
    if docs_kept < len(documents):
        dropped.append(f"{len(documents) - docs_kept} documents")
    if "tools" not in kept:
        dropped.append("tools")
    return {"kept": kept, "dropped": dropped, "tokens": used}


def offload(rows: list[dict], path: str) -> dict:
    return {"saved_to": path, "rows": len(rows), "columns": list(rows[0]) if rows else [], "preview": rows[:2]}


def query(rows: list[dict], where_field: str, where_value, field: str) -> dict:
    hit = [r for r in rows if r.get(where_field) == where_value]
    return {"count": len(hit), "sum": round(sum(r[field] for r in hit), 2)}