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

Fill in the TODOs, then run the tests:
    pytest test_context_practice.py        (locally)
or press "Run tests" on the chapter page.
"""


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


def profile(parts: dict[str, str]) -> list[tuple[str, int, float]]:
    """Measure before you cut: (part, tokens, percent of the whole prompt) for each part,
    biggest first. Round the percent to 1 place.

    TODO 1
    """
    return []


def keep_the_end(parts: dict[str, str], budget: int) -> list[str]:
    """The naive cut. Join the parts in order (dicts keep their order), keep only the LAST
    `budget` tokens (words), and return the names of the parts that survive, even partly,
    in prompt order.

    TODO 2
    """
    return list(parts)


def fit_by_priority(rules: str, question: str, memory: str, documents: list[str], history: list[str],
                    tools: str, budget: int) -> dict:
    """Cut by priority instead.

    Always keep the rules and the question. Then add, while each next item still fits in `budget`:
      1. the memory
      2. documents one at a time, best first (stop at the first that doesn't fit)
      3. messages one at a time, NEWEST first (stop at the first that doesn't fit)
      4. the tools, last
    Labels: "rules", "question", "memory", "doc 1".., "msg 1".. (msg N is the newest), "tools".

    Return {"kept": [labels in the order added], "tokens": tokens used,
            "dropped": a list of what didn't make it, in this order and only if something was dropped:
                       "<n> oldest messages", "<n> documents", "tools"}

    TODO 3
    """
    return {"kept": ["rules", "question"], "dropped": [], "tokens": tokens(rules) + tokens(question)}


def offload(rows: list[dict], path: str) -> dict:
    """What the tool returns instead of the rows: a short summary of the file it saved.

    {"saved_to": path, "rows": number of rows, "columns": the first row's keys, "preview": the first 2 rows}

    TODO 4
    """
    return {"saved_to": path, "rows": rows}


def query(rows: list[dict], where_field: str, where_value, field: str) -> dict:
    """The query tool: count the rows where row[where_field] == where_value, and add up `field`.

    Return {"count": n, "sum": total rounded to 2 places}. Code does the arithmetic, not the model.

    TODO 5
    """
    return {"count": 0, "sum": 0.0}