#!/usr/bin/env python3
"""patterns_toolkit.py: agent design patterns with scripted mock models. Standard library only.

Part 1  Five workflow patterns: prompt_chain, route, sectioning, voting, orchestrate, refine
Part 2  Reasoning and acting loops: react, reflexion, codeact (and a JSON-action loop), rewoo
Part 3  Context management: Session (an append-only event log with branches), build_context,
        observation masking, compaction, branch summaries, isolation, and a cost simulation
Part 4  A demo with seven sections. Run: python patterns_toolkit.py

Every "model" here is a mock: a list of canned replies, or a small rule that I wrote. The mocks
show how the patterns work. They say nothing about how well a real model does. All facts, prices
and tasks are made up. "Tokens" are a rough proxy (words and punctuation marks).
Run the tests with: python test_patterns_toolkit.py
"""
from __future__ import annotations

import ast
import builtins
import contextlib
import io
import json
import operator
import re
import sys
import threading
from collections import Counter
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple


def count_tokens(text: str) -> int:
    """A rough token proxy: words plus punctuation marks. Real tokenizers differ; only ratios matter here."""
    return len(re.findall(r"\w+|[^\w\s]", text))


# ----------------------------------------------------------------------------- mock models and a meter
class ScriptExhausted(RuntimeError):
    """The scripted model was called more times than it has replies."""


class ScriptedLM:
    """Returns canned replies in order and records every prompt. Thread-safe."""

    def __init__(self, replies: Sequence[str]):
        self.replies = list(replies)
        self.prompts: List[str] = []
        self._lock = threading.Lock()

    def __call__(self, prompt: str) -> str:
        with self._lock:
            self.prompts.append(prompt)
            i = len(self.prompts) - 1
        if i >= len(self.replies):
            raise ScriptExhausted("script has %d replies, call number %d" % (len(self.replies), i + 1))
        return self.replies[i]


class RuleLM:
    """Wraps a function prompt -> reply, and records every prompt. Thread-safe."""

    def __init__(self, fn: Callable[[str], str]):
        self.fn = fn
        self.prompts: List[str] = []
        self._lock = threading.Lock()

    def __call__(self, prompt: str) -> str:
        with self._lock:
            self.prompts.append(prompt)
        return self.fn(prompt)


class Meter:
    """Counts model calls and tokens. meter.wrap(llm) returns a model that reports to the meter."""

    def __init__(self):
        self.calls = 0
        self.tokens_in = 0
        self.tokens_out = 0
        self._lock = threading.Lock()

    def wrap(self, llm: Callable[[str], str]) -> Callable[[str], str]:
        def metered(prompt: str) -> str:
            out = llm(prompt)
            with self._lock:
                self.calls += 1
                self.tokens_in += count_tokens(prompt)
                self.tokens_out += count_tokens(out)
            return out
        return metered


# ----------------------------------------------------------------------------- tiny tools (all made up)
_OPS = {ast.Add: operator.add, ast.Sub: operator.sub, ast.Mult: operator.mul, ast.Div: operator.truediv,
        ast.FloorDiv: operator.floordiv, ast.Mod: operator.mod, ast.Pow: operator.pow,
        ast.USub: operator.neg, ast.UAdd: operator.pos}


def safe_eval(expr: str) -> float:
    """Evaluate plain arithmetic with ast. No names, calls or attributes."""
    def ev(n):
        if isinstance(n, ast.Expression):
            return ev(n.body)
        if isinstance(n, ast.Constant) and type(n.value) in (int, float):
            return n.value
        if isinstance(n, ast.BinOp) and type(n.op) in _OPS:
            a, b = ev(n.left), ev(n.right)
            if isinstance(n.op, ast.Pow) and abs(b) > 100:
                raise ValueError("exponent too large")
            return _OPS[type(n.op)](a, b)
        if isinstance(n, ast.UnaryOp) and type(n.op) in _OPS:
            return _OPS[type(n.op)](ev(n.operand))
        raise ValueError("unsupported expression")
    return ev(ast.parse(expr.strip(), mode="eval"))


def calc(arg: str) -> str:
    """Tool: arithmetic. Returns text, never raises."""
    try:
        v = safe_eval(arg)
    except (ValueError, SyntaxError, ZeroDivisionError, OverflowError) as e:
        return "error: %s" % e
    return str(int(v)) if float(v).is_integer() else str(round(v, 6))


# A made-up knowledge base: (keywords, answer). search() picks the entry that shares most words with the query.
KB: List[Tuple[set, str]] = [
    ({"who", "founded", "harbor", "labs"}, "Mira Okoye"),
    ({"year", "founded", "harbor", "labs"}, "2009"),
    ({"birth", "year", "mira", "okoye"}, "1971"),
    ({"who", "founded", "ridge", "corp"}, "Jonas Weld"),
    ({"year", "founded", "ridge", "corp"}, "2004"),
    ({"birth", "year", "jonas", "weld"}, "1964"),
]


def search(query: str) -> str:
    """Tool: look up a fact in the made-up KB. Needs at least 3 shared words, else 'no result'."""
    words = set(re.findall(r"[a-z0-9]+", query.lower()))
    best, score = "no result", 0
    for kw, ans in KB:
        s = len(kw & words)
        if s > score:
            best, score = ans, s
    return best if score >= 3 else "no result"


TOOLS: Dict[str, Callable[[str], str]] = {"search": search, "calc": calc}


# ============================================================================= PART 1: workflow patterns
# In a workflow the control flow lives in this code. The model fills in text at each step.

@dataclass
class ChainResult:
    outputs: List[str]
    ok: bool
    stopped_at: Optional[int] = None      # index of the step whose gate failed


def prompt_chain(llm, steps: Sequence[str], task: str,
                 gate: Optional[Callable[[int, str], bool]] = None) -> ChainResult:
    """Prompt chaining. Each step is a template with {input}. The output of one step is the input of the next.
    gate(i, output) is a programmatic check after step i. If it returns False the chain stops there."""
    outputs: List[str] = []
    current = task
    for i, template in enumerate(steps):
        current = llm(template.format(input=current)).strip()
        outputs.append(current)
        if gate is not None and not gate(i, current):
            return ChainResult(outputs, False, i)
    return ChainResult(outputs, True)


class RoutingError(ValueError):
    pass


CLASSIFY = "Classify the input into one of: {labels}. Reply with the label only.\nInput: {input}"


def route(llm, text: str, routes: Dict[str, Tuple[Optional[Callable], str]],
          default: Optional[str] = None) -> Tuple[str, str]:
    """Routing. routes maps label -> (model or None, template with {input}). The classifier is `llm`.
    A route may name a different, cheaper or stronger model. An unknown label goes to `default`, or raises."""
    labels = list(routes)
    reply = llm(CLASSIFY.format(labels=", ".join(labels), input=text)).strip().lower()
    label = next((l for l in labels if re.search(r"\b%s\b" % re.escape(l.lower()), reply)), None)
    if label is None:
        if default is None:
            raise RoutingError("classifier said %r; known labels: %s" % (reply, labels))
        label = default
    model, template = routes[label]
    return label, (model or llm)(template.format(input=text))


def sectioning(llm, text: str, templates: Sequence[str], max_workers: int = 4) -> List[str]:
    """Parallelization, variant 1: independent subtasks run at the same time. Results keep the template order."""
    with ThreadPoolExecutor(max_workers=max_workers) as pool:
        return list(pool.map(lambda t: llm(t.format(input=text)), templates))


@dataclass
class Vote:
    winner: str
    counts: Dict[str, int]
    agreement: float


def voting(llm, prompt: str, n: int, parse: Callable[[str], str] = str.strip) -> Vote:
    """Parallelization, variant 2: ask the same thing n times and take the most common answer.
    Ties go to the answer seen first. (Run the calls in parallel in production; here they are in order.)"""
    if n < 1:
        raise ValueError("n must be at least 1")
    answers = [parse(llm(prompt)) for _ in range(n)]
    counts = Counter(answers)
    best = max(counts.values())
    winner = next(a for a in answers if counts[a] == best)
    return Vote(winner, dict(counts), best / n)


class PlanError(ValueError):
    pass


PLAN_PROMPT = ("Break the task into at most {k} independent subtasks. "
               "Reply with a JSON list of strings only.\nTask: {task}")


def parse_subtasks(reply: str) -> List[str]:
    """Accept a bare JSON list or one inside a code fence. Reject anything else."""
    m = re.search(r"\[.*\]", reply, re.S)
    if not m:
        raise PlanError("no JSON list in reply")
    try:
        items = json.loads(m.group(0))
    except json.JSONDecodeError as e:
        raise PlanError("bad JSON: %s" % e)
    if not isinstance(items, list) or not items or not all(isinstance(s, str) and s.strip() for s in items):
        raise PlanError("expected a non-empty list of non-empty strings")
    return [s.strip() for s in items]


@dataclass
class OrchestrationResult:
    subtasks: List[str]
    results: List[Tuple[str, str]]
    final: str
    truncated: bool = False
    plan_attempts: int = 1


def orchestrate(llm, task: str, worker: Optional[Callable[[str], str]] = None, max_subtasks: int = 5,
                synth_template: str = "Combine these results into one answer for: {task}\n{results}") -> OrchestrationResult:
    """Orchestrator-workers. One call plans the subtasks (the plan is NOT fixed in code), workers run them,
    one call combines the results. A bad plan is retried once. More than max_subtasks are cut off."""
    worker = worker or (lambda s: llm("Do this subtask and reply briefly: " + s))
    attempts, plan = 0, None
    for attempt in (PLAN_PROMPT, PLAN_PROMPT + "\nYour last reply was not valid. Reply with the JSON list only."):
        attempts += 1
        try:
            plan = parse_subtasks(llm(attempt.format(k=max_subtasks, task=task)))
            break
        except PlanError as e:
            last = e
    if plan is None:
        raise PlanError("no valid plan after %d tries: %s" % (attempts, last))
    truncated = len(plan) > max_subtasks
    plan = plan[:max_subtasks]
    results = [(s, worker(s)) for s in plan]
    final = llm(synth_template.format(task=task, results="\n".join("- %s: %s" % r for r in results)))
    return OrchestrationResult(plan, results, final.strip(), truncated, attempts)


@dataclass
class Round:
    number: int
    draft: str
    ok: bool
    feedback: str


@dataclass
class RefineResult:
    final: str
    converged: bool
    rounds: List[Round]


def refine(generate: Callable[[str, List[str]], str], evaluate: Callable[[str], Tuple[bool, str]],
           task: str, max_rounds: int = 3) -> RefineResult:
    """Evaluator-optimizer. generate(task, feedback_so_far) -> draft. evaluate(draft) -> (ok, feedback).
    It works only if the feedback is clear and the generator can use it. Always bounded by max_rounds."""
    if max_rounds < 1:
        raise ValueError("max_rounds must be at least 1")
    feedback: List[str] = []
    rounds: List[Round] = []
    for r in range(1, max_rounds + 1):
        draft = generate(task, list(feedback))
        ok, fb = evaluate(draft)
        rounds.append(Round(r, draft, ok, fb))
        if ok:
            return RefineResult(draft, True, rounds)
        feedback.append(fb)
    return RefineResult(rounds[-1].draft, False, rounds)


# ============================================================================= PART 2: reasoning and acting
# In an agent the model chooses the next step. The loops below differ in WHAT the model writes each turn.

@dataclass
class Step:
    thought: str
    tool: str
    arg: str
    observation: str


@dataclass
class AgentResult:
    answer: Optional[str]
    status: str                  # finished, step_limit, stuck
    steps: List[Step]


REACT_SYSTEM = ("Answer the question. Each turn write:\nThought: your reasoning\nAction: tool[argument]\n"
                "Tools: {tools}. To end, write Action: finish[answer].\n")
_ACTION = re.compile(r"^Action:\s*([A-Za-z_]\w*)\[(.*)\]\s*$", re.M)


def parse_react(reply: str) -> Tuple[str, Optional[str], str]:
    """Return (thought, tool, argument). tool is None if there is no well-formed Action line.
    Only the first Action line counts: one action per turn."""
    m = _ACTION.search(reply)
    t = re.search(r"^Thought:\s*(.*?)\s*(?=^Action:|\Z)", reply, re.M | re.S)
    thought = t.group(1).strip() if t else ""
    return (thought, m.group(1), m.group(2).strip()) if m else (thought, None, "")


def react(llm, question: str, tools: Dict[str, Callable[[str], str]], max_steps: int = 8,
          stuck_after: int = 3, memory: Sequence[str] = ()) -> AgentResult:
    """ReAct: the model writes a Thought and one Action. The harness runs the action and appends the
    Observation. The WHOLE transcript is sent again on every turn. memory holds lessons from earlier trials.
    stuck_after: stop if the same action gets the same observation this many times (Reflexion's heuristic)."""
    header = REACT_SYSTEM.format(tools=", ".join(sorted(tools)))
    if memory:
        header += "Lessons from earlier trials:\n" + "\n".join("- " + m for m in memory) + "\n"
    transcript = "Question: %s\n" % question
    steps: List[Step] = []
    seen: Counter = Counter()
    for _ in range(max_steps):
        reply = llm(header + transcript)
        thought, name, arg = parse_react(reply)
        if name is None:
            obs = "format error: write a Thought, then one line: Action: tool[argument]"
            steps.append(Step(thought, "", "", obs))
            transcript += reply.strip() + "\nObservation: " + obs + "\n"
            continue
        if name == "finish":
            steps.append(Step(thought, name, arg, ""))
            return AgentResult(arg, "finished", steps)
        fn = tools.get(name)
        if fn is None:
            obs = "unknown tool %r. Tools: %s" % (name, ", ".join(sorted(tools)))
        else:
            try:
                obs = str(fn(arg))
            except Exception as e:                                  # a tool bug becomes an observation
                obs = "error: %s: %s" % (type(e).__name__, e)
        steps.append(Step(thought, name, arg, obs))
        transcript += "Thought: %s\nAction: %s[%s]\nObservation: %s\n" % (thought, name, arg, obs)
        seen[(name, arg, obs)] += 1
        if seen[(name, arg, obs)] >= stuck_after:
            return AgentResult(None, "stuck", steps)
    return AgentResult(None, "step_limit", steps)


@dataclass
class Trial:
    number: int
    answer: Optional[str]
    ok: bool
    feedback: str
    memory_used: List[str]


@dataclass
class ReflexionResult:
    success: bool
    trials: List[Trial]
    memory: List[str]


def reflexion(run_trial: Callable[[List[str]], AgentResult],
              evaluate: Callable[[AgentResult], Tuple[bool, str]],
              reflect: Callable[[AgentResult, str], str],
              max_trials: int = 3, memory_size: int = 3) -> ReflexionResult:
    """Reflexion. After a failed trial, `reflect` writes a short lesson from the trajectory and the feedback.
    The lessons go into the next trial's prompt. Memory is a sliding window (the paper keeps 1 to 3).
    `evaluate` must be GROUNDED: tests, an environment signal, a checker. A model grading itself is weaker."""
    memory: List[str] = []
    trials: List[Trial] = []
    for n in range(1, max_trials + 1):
        used = list(memory)
        result = run_trial(used)
        ok, feedback = evaluate(result)
        trials.append(Trial(n, result.answer, ok, feedback, used))
        if ok:
            return ReflexionResult(True, trials, memory)
        memory.append(reflect(result, feedback))
        memory = memory[-memory_size:]
    return ReflexionResult(False, trials, memory)


# ----------------------------------------------------------------------------- CodeAct: actions are code
class StepLimit(Exception):
    """The code ran more lines than allowed."""


_SAFE_NAMES = ("abs all any bool dict enumerate filter float int isinstance len list map max min print range "
               "reversed round set sorted str sum tuple zip Exception ValueError KeyError TypeError "
               "IndexError ZeroDivisionError").split()
SAFE_BUILTINS = {n: getattr(builtins, n) for n in _SAFE_NAMES}


@dataclass
class Obs:
    ok: bool
    text: str


class CodeSession:
    """A persistent Python namespace. Variables survive between turns. Tools are plain functions in it.
    NOT A SECURITY SANDBOX: it only removes imports and file access from the builtins and stops runaway
    loops. Real systems run code in a container or a VM."""

    def __init__(self, tools: Optional[Dict[str, Callable]] = None, max_lines: int = 5000, max_output: int = 400):
        self.ns: Dict[str, Any] = {"__builtins__": SAFE_BUILTINS}
        self.ns.update(tools or {})
        self.max_lines, self.max_output = max_lines, max_output

    def run(self, code: str) -> Obs:
        try:
            compiled = compile(code, "<codeact>", "exec")
        except SyntaxError as e:
            return Obs(False, "SyntaxError: %s (line %s)" % (e.msg, e.lineno))
        count = [0]

        def tracer(frame, event, arg):
            if frame.f_code.co_filename != "<codeact>":
                return None

            def local(frame, event, arg):
                if event == "line":
                    count[0] += 1
                    if count[0] > self.max_lines:
                        raise StepLimit("more than %d lines executed" % self.max_lines)
                return local
            return local

        out, err, old = io.StringIO(), "", sys.gettrace()
        try:
            sys.settrace(tracer)
            with contextlib.redirect_stdout(out):
                exec(compiled, self.ns)
        except Exception as e:
            err = "%s: %s" % (type(e).__name__, e)
        finally:
            sys.settrace(old)
        text = out.getvalue()
        if len(text) > self.max_output:
            text = text[: self.max_output] + "\n[truncated %d characters]" % (len(text) - self.max_output)
        return Obs(not err, (text + err).strip())


CODEACT_SYSTEM = ("Solve the task. Each turn reply with ONE ```python block to run, or 'FINAL: answer' when done. "
                  "Tools are Python functions: {tools}. Variables stay between turns.\n")
_CODE = re.compile(r"```python\n(.*?)```", re.S)
_FINAL = re.compile(r"^FINAL:\s*(.*)$", re.M)


@dataclass
class CodeTurn:
    reply: str
    obs: Optional[Obs]


@dataclass
class CodeActResult:
    answer: Optional[str]
    status: str                  # finished, turn_limit
    turns: List[CodeTurn]


def codeact(llm, task: str, tools: Dict[str, Callable], max_turns: int = 6, max_lines: int = 5000) -> CodeActResult:
    """CodeAct: the action is a block of code. One block can loop, branch and keep results in variables."""
    session = CodeSession(tools, max_lines)
    header = CODEACT_SYSTEM.format(tools=", ".join(sorted(tools)))
    transcript = "Task: %s\n" % task
    turns: List[CodeTurn] = []
    for _ in range(max_turns):
        reply = llm(header + transcript)
        code, final = _CODE.search(reply), _FINAL.search(reply)
        if code is None and final:
            turns.append(CodeTurn(reply, None))
            return CodeActResult(final.group(1).strip(), "finished", turns)
        if code is None:
            obs = Obs(False, "format error: reply with one ```python block or FINAL: answer")
        else:
            obs = session.run(code.group(1))
        turns.append(CodeTurn(reply, obs))
        transcript += reply.strip() + "\nObservation: " + obs.text + "\n"
    return CodeActResult(None, "turn_limit", turns)


JSON_SYSTEM = ('Solve the task. Each turn reply with ONE JSON object: {{"tool": name, "args": {{...}}}} '
               'or {{"final": answer}}. Tools: {tools}.\n')


def json_actions(llm, task: str, tools: Dict[str, Callable], max_turns: int = 12) -> CodeActResult:
    """The usual tool-calling loop: one JSON call per turn. Same result type as codeact, to compare turns."""
    header = JSON_SYSTEM.format(tools=", ".join(sorted(tools)))
    transcript = "Task: %s\n" % task
    turns: List[CodeTurn] = []
    for _ in range(max_turns):
        reply = llm(header + transcript)
        try:
            call = json.loads(reply)
            if "final" in call:
                turns.append(CodeTurn(reply, None))
                return CodeActResult(str(call["final"]), "finished", turns)
            obs = Obs(True, str(tools[call["tool"]](**call.get("args", {}))))
        except Exception as e:
            obs = Obs(False, "%s: %s" % (type(e).__name__, e))
        turns.append(CodeTurn(reply, obs))
        transcript += reply.strip() + "\nObservation: " + obs.text + "\n"
    return CodeActResult(None, "turn_limit", turns)


# ----------------------------------------------------------------------------- ReWOO and LLMCompiler: plan first
@dataclass
class PlanStep:
    id: str                      # "E1"
    tool: str
    arg: str                     # may contain #E1 style references


class PlanParseError(ValueError):
    pass


_PLAN_LINE = re.compile(r"^#(E\d+)\s*=\s*([A-Za-z_]\w*)\[(.*)\]\s*$")
_REF = re.compile(r"#(E\d+)")


def parse_plan(text: str, tools: Sequence[str]) -> List[PlanStep]:
    """Lines like '#E2 = search[birth year of #E1]'. Other lines (a 'Plan:' sentence) are ignored.
    Checks: known tool, unique ids, and a reference may only point to an EARLIER step (so no cycles)."""
    steps: List[PlanStep] = []
    for line in text.splitlines():
        m = _PLAN_LINE.match(line.strip())
        if not m:
            continue
        sid, tool, arg = m.group(1), m.group(2), m.group(3).strip()
        if tool not in tools:
            raise PlanParseError("step %s uses unknown tool %r" % (sid, tool))
        if any(s.id == sid for s in steps):
            raise PlanParseError("duplicate step id %s" % sid)
        for ref in _REF.findall(arg):
            if ref not in {s.id for s in steps}:
                raise PlanParseError("step %s refers to %s, which is not an earlier step" % (sid, ref))
        steps.append(PlanStep(sid, tool, arg))
    if not steps:
        raise PlanParseError("no steps found")
    return steps


def waves(plan: Sequence[PlanStep]) -> List[List[PlanStep]]:
    """Group steps into levels. All steps in one level are independent and can run together
    (the LLMCompiler idea). Level = 1 + the deepest level among the steps it refers to."""
    level: Dict[str, int] = {}
    for s in plan:
        refs = _REF.findall(s.arg)
        level[s.id] = 1 + max((level[r] for r in refs), default=0)
    out: List[List[PlanStep]] = [[] for _ in range(max(level.values()))]
    for s in plan:
        out[level[s.id] - 1].append(s)
    return out


def execute_plan(plan: Sequence[PlanStep], tools: Dict[str, Callable[[str], str]],
                 parallel: bool = False) -> Dict[str, str]:
    """Run the plan with NO model calls. #E refs are replaced by earlier results. A failing tool
    gives an 'error:' string as evidence, so the solver can still answer or say what failed."""
    evidence: Dict[str, str] = {}

    def run_one(s: PlanStep) -> Tuple[str, str]:
        arg = _REF.sub(lambda m: evidence[m.group(1)], s.arg)
        try:
            return s.id, str(tools[s.tool](arg))
        except Exception as e:
            return s.id, "error: %s: %s" % (type(e).__name__, e)

    for level in waves(plan):
        if parallel and len(level) > 1:
            with ThreadPoolExecutor(max_workers=len(level)) as pool:
                results = list(pool.map(run_one, level))
        else:
            results = [run_one(s) for s in level]
        evidence.update(results)
    return evidence


REWOO_PLAN = ("Write a plan. One line per step: #E1 = tool[argument]. A later step may use #E1 in its argument.\n"
              "Tools: {tools}\nQuestion: {question}\n")
REWOO_SOLVE = "Use the plan and the evidence to answer.\nQuestion: {question}\n{plan}\nEvidence:\n{evidence}\nAnswer:"


@dataclass
class RewooResult:
    answer: str
    plan: List[PlanStep]
    evidence: Dict[str, str]
    waves: int


def rewoo(planner, solver, question: str, tools: Dict[str, Callable[[str], str]], parallel: bool = False) -> RewooResult:
    """ReWOO: Planner (one call) -> Worker (tools, no model) -> Solver (one call).
    Two model calls in total, however many tools run. The price: the plan is made before any observation."""
    plan_text = planner(REWOO_PLAN.format(tools=", ".join(sorted(tools)), question=question))
    plan = parse_plan(plan_text, list(tools))
    evidence = execute_plan(plan, tools, parallel=parallel)
    ev_text = "\n".join("#%s: %s" % (k, v) for k, v in evidence.items())
    plan_view = "\n".join("#%s = %s[%s]" % (s.id, s.tool, s.arg) for s in plan)
    answer = solver(REWOO_SOLVE.format(question=question, plan=plan_view, evidence=ev_text)).strip()
    return RewooResult(answer, plan, evidence, len(waves(plan)))


# ============================================================================= PART 3: context management
# One rule drives this part: keep the full record, and build each model input from it.

@dataclass(frozen=True)
class Event:
    id: int
    parent: Optional[int]
    kind: str                    # system, user, assistant, observation, compaction, branch_summary
    text: str
    first_kept: Optional[int] = None     # compaction only: the first event that stays in the context


_ROLE = {"system": "system", "user": "user", "assistant": "assistant", "observation": "tool",
         "compaction": "user", "branch_summary": "user"}


class Session:
    """An append-only event log with parent links. A linear chat is a chain. A branch is a second child of
    some event. `leaf` is the current position. Nothing is ever edited or removed, so the full record
    stays available for audit, replay and recall, whatever the model is shown."""

    def __init__(self):
        self.events: List[Event] = []
        self.leaf: Optional[int] = None
        self.observers: List[Callable[[Event], None]] = []

    def append(self, kind: str, text: str, parent: Optional[int] = -1, first_kept: Optional[int] = None) -> Event:
        if kind not in _ROLE:
            raise ValueError("unknown event kind %r" % kind)
        p = self.leaf if parent == -1 else parent
        if p is not None and not (0 <= p < len(self.events)):
            raise ValueError("no such parent %r" % p)
        e = Event(len(self.events), p, kind, text, first_kept)
        self.events.append(e)
        self.leaf = e.id
        for fn in self.observers:                      # hooks: monitors, loggers
            fn(e)
        return e

    def get(self, event_id: int) -> Event:
        return self.events[event_id]

    def path(self, leaf: Optional[int] = None) -> List[Event]:
        """Events from the root to `leaf` (default: the current leaf)."""
        cur = self.leaf if leaf is None else leaf
        out: List[Event] = []
        while cur is not None:
            e = self.events[cur]
            out.append(e)
            cur = e.parent
        return out[::-1]

    def to_jsonl(self) -> str:
        return "\n".join(json.dumps(e.__dict__, sort_keys=True) for e in self.events)

    @classmethod
    def from_jsonl(cls, text: str) -> "Session":
        s = cls()
        for line in text.splitlines():
            if line.strip():
                d = json.loads(line)
                if d["id"] != len(s.events):
                    raise ValueError("event ids must be 0, 1, 2, ... in order")
                s.events.append(Event(**d))
                s.leaf = d["id"]
        return s


def recall(session: Session, event_id: int) -> str:
    """Get the full text of any event, even one that was masked or compacted. This is what makes
    masking 'restorable': the record is still there, only the model's view of it is shorter."""
    return session.get(event_id).text


def _body(session: Session, leaf: Optional[int]) -> List[Event]:
    """The events the model would see before masking: system events, the latest compaction summary,
    the events from first_kept onward, and anything after the summary."""
    path = session.path(leaf)
    cidx = max((i for i, e in enumerate(path) if e.kind == "compaction"), default=None)
    if cidx is None:
        return path
    c = path[cidx]
    ids = [e.id for e in path]
    start = ids.index(c.first_kept) if c.first_kept in ids else cidx + 1
    head = [e for e in path[:cidx] if e.kind == "system"]
    kept = [e for e in path[start:cidx] if e.kind != "compaction"]
    return head + [c] + kept + path[cidx + 1:]


@dataclass(frozen=True)
class Message:
    role: str
    text: str
    src: int                     # id of the event this message came from


def mask_text(e: Event) -> str:
    return "[output omitted: %d tokens; recall id %d]" % (count_tokens(e.text), e.id)


def build_context(session: Session, leaf: Optional[int] = None, mask_window: Optional[int] = None,
                  mask_stride: int = 1) -> List[Message]:
    """Build the model input from the log. mask_window=M keeps the full text of the M most recent
    observations and replaces older ones with a short placeholder. Thoughts and actions are never touched.
    mask_stride=S makes the boundary move in steps of S observations, so the start of the masked region
    changes only every S turns and a prompt cache (see the simulation) stays valid in between.
    A placeholder is used only if it is shorter than the text it replaces."""
    if mask_stride < 1:
        raise ValueError("mask_stride must be at least 1")
    body = _body(session, leaf)
    obs_ids = [e.id for e in body if e.kind == "observation"]
    if mask_window is None:
        keep_full = set(obs_ids)
    else:
        n_old = max(0, len(obs_ids) - mask_window) // mask_stride * mask_stride if mask_window > 0 else len(obs_ids)
        keep_full = set(obs_ids[n_old:])
    out = []
    for e in body:
        text = e.text
        if e.kind == "observation" and e.id not in keep_full:
            ph = mask_text(e)
            if count_tokens(ph) < count_tokens(text):
                text = ph
        out.append(Message(_ROLE[e.kind], text, e.id))
    return out


def context_tokens(msgs: Sequence[Message]) -> int:
    return sum(count_tokens(m.text) for m in msgs)


def context_contains(msgs: Sequence[Message], needle: str) -> bool:
    return any(needle in m.text for m in msgs)


# ----------------------------------------------------------------------------- compaction
def extractive_summary(msgs: Sequence[Message], max_tokens: int = 150, keep_errors: bool = True) -> str:
    """A MOCK summarizer: it copies, it does not understand. A real one would be a model call.
    Format follows the idea of a structured summary: Goal, Progress, Files, Critical context."""
    goal = next((m.text for m in msgs if m.role == "user" and not m.text.startswith("## ")), "")
    goal = " ".join(goal.split()[:25])
    actions = [m.text.strip().splitlines()[0][:60] for m in msgs if m.role == "assistant"]
    read, modified = [], []
    for m in msgs:
        if m.role != "assistant":
            continue
        for f in re.findall(r"[\w./-]+\.\w{1,4}\b", m.text):
            target = modified if (">" in m.text or "sed -i" in m.text) else read
            if f not in target:
                target.append(f)
    errors = [ln.strip() for m in msgs if m.role == "tool" for ln in m.text.splitlines() if "error" in ln.lower()]
    prior = [m.text for m in msgs if m.text.startswith("## Goal")]
    for p in prior:                                   # carry files and errors across repeated compactions
        for ln in p.splitlines():
            if ln.startswith("read: "):
                read[:0] = [f for f in ln[6:].split(", ") if f and f not in read]
            elif ln.startswith("modified: "):
                modified[:0] = [f for f in ln[10:].split(", ") if f and f not in modified]
            elif ln.startswith("- ") and "error" in ln.lower():
                errors.append(ln[2:])
    for n in range(len(actions), -1, -1):
        shown = actions[len(actions) - n:] if n else []
        parts = ["## Goal\n" + goal, "## Progress\n" + "\n".join("- " + a for a in shown),
                 "## Files\nread: %s\nmodified: %s" % (", ".join(read[:6]), ", ".join(modified[:6]))]
        if keep_errors and errors:
            parts.append("## Critical context\n" + "\n".join("- " + e[:80] for e in errors[:2]))
        if prior:
            parts.insert(1, "## Earlier summary\n" + " ".join(prior[-1].split()[:20]) + " ...")
        text = "\n".join(parts)
        if count_tokens(text) <= max_tokens or n == 0:
            break
    return text


def should_compact(used_tokens: int, window: int, reserve: int) -> bool:
    """Compact when used > window - reserve. The reserve leaves room for the next reply and the summary."""
    return used_tokens > window - reserve


def compact(session: Session, summarizer: Callable[[Sequence[Message]], str] = extractive_summary,
            keep_recent_tokens: int = 800) -> Optional[Event]:
    """Summarize the old part of the current context and append a compaction event.
    Walk backwards from the newest event until keep_recent_tokens are covered. The cut point must be a
    user or assistant event, never an observation, so a tool result stays with the call that made it.
    Returns None if there is nothing old enough to summarize."""
    body = _body(session, None)
    head = [i for i, e in enumerate(body) if e.kind == "system"]
    first = (head[-1] + 1) if head else 0
    acc, cut = 0, None
    for i in range(len(body) - 1, first - 1, -1):
        acc += count_tokens(body[i].text)
        if acc >= keep_recent_tokens:
            cut = i
            break
    if cut is None:
        return None
    while cut < len(body) and body[cut].kind == "observation":      # move forward to a valid boundary
        cut += 1
    if cut >= len(body) or cut <= first:
        return None
    old = [Message(_ROLE[e.kind], e.text, e.id) for e in body[first:cut]]
    summary = summarizer(old)
    return session.append("compaction", summary, first_kept=body[cut].id)


def switch_branch(session: Session, new_leaf: int, summarizer: Callable[[Sequence[Message]], str] = extractive_summary) -> Event:
    """Move to another branch. The work left behind is summarized and attached to the new branch,
    so what was learned there is not lost. The old branch stays in the log."""
    old = session.path(session.leaf)
    new = session.path(new_leaf)
    common = 0
    while common < min(len(old), len(new)) and old[common].id == new[common].id:
        common += 1
    abandoned = [Message(_ROLE[e.kind], e.text, e.id) for e in old[common:]]
    session.leaf = new_leaf
    return session.append("branch_summary", summarizer(abandoned) if abandoned else "(nothing left behind)")


# ----------------------------------------------------------------------------- isolation
@dataclass
class IsolationResult:
    report: str
    child: Session
    child_tokens: int            # all text in the child's log
    parent_added_tokens: int     # what the parent gained: the call plus the report


def isolate(parent: Session, task: str, worker: Callable[[Session], str]) -> IsolationResult:
    """Context isolation. The worker runs in its OWN session. The parent log only gets the delegation
    and the report. The child's long exploration never enters the parent's context."""
    child = Session()
    child.append("system", "You are a worker. Do one task. Reply with a short report.")
    child.append("user", task)
    report = worker(child)
    call = parent.append("assistant", "delegate[%s]" % task)
    obs = parent.append("observation", report)
    return IsolationResult(report, child, sum(count_tokens(e.text) for e in child.events),
                           count_tokens(call.text) + count_tokens(obs.text))


# ----------------------------------------------------------------------------- a cost simulation
@dataclass
class Strategy:
    name: str
    mask_window: Optional[int] = None
    mask_stride: int = 1
    compact: bool = False
    isolate_first: int = 0       # run this many first turns in a sub-session


@dataclass
class SimResult:
    name: str
    calls: int
    input_tokens: int            # summed over every model call, parent and child
    peak_context: int
    final_context: int
    cached_tokens: int           # tokens that were an exact prefix of the previous call's input
    cost: float
    needle_in_context: bool
    compactions: int
    per_call: List[int] = field(default_factory=list)     # input tokens of each parent call, for plotting


PRICE_IN, PRICE_CACHED = 3.0, 0.30          # made-up dollars per million tokens (10 times cheaper when cached)
NEEDLE = "the release key is K-4471"


def _filler(n: int, tag: str) -> str:
    return " ".join("%s%d" % (tag, i) for i in range(n))


def make_trajectory(turns: int = 40, system_tokens: int = 200) -> Tuple[str, List[Tuple[str, str]]]:
    """A made-up coding-agent run: (assistant action, observation) pairs. Observation sizes cycle through
    a fixed list, so the run is the same every time. Turn 3 holds the needle fact."""
    sizes = [60, 900, 150, 1400, 80, 500, 300, 1100, 40, 700]
    traj = []
    for t in range(turns):
        action = "run: cat src/mod%d.py" % t if t % 3 else "run: pytest tests/test_%d.py" % t
        n = sizes[t % len(sizes)]
        obs = _filler(n, "t%d_" % t)
        if t == 3:
            obs = "NOTE: " + NEEDLE + ". " + obs
        if t == 4:
            obs = "Error: test_parse failed at line 12. " + obs
        traj.append((action, obs))
    return _filler(system_tokens, "sys"), traj


def _simulate_turns(sess: Session, traj, strat: Strategy, window: int, reserve: int, keep_recent: int):
    inputs, prev, cached, peak, compactions, calls, series = 0, [], 0, 0, 0, 0, []
    for action, obs in traj:
        ctx = build_context(sess, mask_window=strat.mask_window, mask_stride=strat.mask_stride)
        used = context_tokens(ctx)
        inputs += used
        series.append(used)
        calls += 1
        peak = max(peak, used)
        key = [(m.src, m.text) for m in ctx]
        n = 0
        while n < min(len(prev), len(key)) and prev[n] == key[n]:
            n += 1
        cached += sum(count_tokens(m.text) for m in ctx[:n])
        prev = key
        sess.append("assistant", action)
        sess.append("observation", obs)
        if strat.compact and should_compact(context_tokens(build_context(sess, mask_window=strat.mask_window, mask_stride=strat.mask_stride)), window, reserve):
            if compact(sess, keep_recent_tokens=keep_recent):
                compactions += 1
    return inputs, calls, cached, peak, compactions, series


def simulate(strat: Strategy, turns: int = 40, window: int = 4000, reserve: int = 1000, keep_recent: int = 900,
             report_tokens: int = 60) -> SimResult:
    """Run the made-up trajectory under one strategy and measure what the model would be sent."""
    system, traj = make_trajectory(turns)
    parent = Session()
    parent.append("system", system)
    parent.append("user", "Fix the failing parser test in the repository.")
    child_inputs = child_calls = 0
    if strat.isolate_first:
        first, traj = traj[: strat.isolate_first], traj[strat.isolate_first:]

        def worker(child: Session) -> str:
            nonlocal child_inputs, child_calls
            i, c, _, _, _, _ = _simulate_turns(child, first, Strategy("child"), window, reserve, keep_recent)
            child_inputs, child_calls = i, c
            return extractive_summary([Message(_ROLE[e.kind], e.text, e.id) for e in child.events], report_tokens)
        isolate(parent, "Explore the repository and report where the parser lives.", worker)
    inputs, calls, cached, peak, comps, series = _simulate_turns(parent, traj, strat, window, reserve, keep_recent)
    final = build_context(parent, mask_window=strat.mask_window, mask_stride=strat.mask_stride)
    total_in = inputs + child_inputs
    cost = ((total_in - cached) * PRICE_IN + cached * PRICE_CACHED) / 1e6
    return SimResult(strat.name, calls + child_calls, total_in, peak, context_tokens(final), cached, cost,
                     context_contains(final, NEEDLE), comps, series)


# ============================================================================= PART 4: the demo
QUESTION = "How old was the founder of Harbor Labs when Harbor Labs was founded?"


def age_policy(company: str, buggy: bool = False) -> RuleLM:
    """A scripted ReAct policy for the founder-age question. With buggy=True it subtracts in the wrong
    order, UNLESS a lesson that mentions 'order' is in the prompt header. A mock of an actor that can learn."""
    def fn(prompt: str) -> str:
        obs = re.findall(r"^Observation: (.*)$", prompt, re.M)
        fixed = "order" in prompt.split("Question:")[0]
        n = len(obs)
        if n == 0:
            return "Thought: I need the founder of %s.\nAction: search[who founded %s]" % (company, company)
        if n == 1:
            return "Thought: Now I need %s's birth year.\nAction: search[birth year of %s]" % (obs[0], obs[0])
        if n == 2:
            return "Thought: Now the year %s was founded.\nAction: search[year %s was founded]" % (company, company)
        if n == 3:
            a, b = obs[2], obs[1]
            if buggy and not fixed:
                a, b = b, a
            return "Thought: Subtract.\nAction: calc[%s - %s]" % (a, b)
        return "Thought: I have the answer.\nAction: finish[%s]" % obs[3]
    return RuleLM(fn)


def phone_tools() -> Dict[str, Callable]:
    """Made-up prices and exchange rates. India comes out cheapest by a small margin."""
    price = {"USA": 799.0, "Japan": 118000.0, "Germany": 749.0, "India": 66000.0}
    rate = {"USA": 1.0, "Japan": 0.0068, "Germany": 1.08, "India": 0.012}

    def get_price(model: str, country: str) -> float:
        if country not in price:
            raise ValueError("unknown country %r" % country)
        return price[country]

    def get_rate(country: str) -> float:
        if country not in rate:
            raise ValueError("unknown country %r" % country)
        return rate[country]
    return {"get_price": get_price, "get_rate": get_rate}


COUNTRIES = ["USA", "Japan", "Germany", "India"]
PHONE_TASK = "Which country is cheapest for the phone X-1, in USD? Countries: USA, Japan, Germany, India."


def json_phone_policy(prompt: str) -> str:
    obs = re.findall(r"^Observation: (.*)$", prompt, re.M)
    n = len(obs)
    if n < 8:
        c = COUNTRIES[n // 2]
        call = ({"tool": "get_price", "args": {"model": "X-1", "country": c}} if n % 2 == 0
                else {"tool": "get_rate", "args": {"country": c}})
        return json.dumps(call)
    usd = {c: float(obs[2 * i]) * float(obs[2 * i + 1]) for i, c in enumerate(COUNTRIES)}
    return json.dumps({"final": min(usd, key=usd.get)})


_BUGGY = '''Run it in one go.
```python
prices = {}
for c in ["USA", "Japan", "Germany", "India"]:
    prices[c] = get_price("X-1", c)
usd = {}
for c in ["USA", "Japan", "Deutschland", "India"]:
    usd[c] = prices.get(c, 0) * get_rate(c)
print(min(usd, key=usd.get))
```'''
_FIXED = '''The name was wrong. The prices are still in memory.
```python
usd = {c: prices[c] * get_rate(c) for c in prices}
print(min(usd, key=usd.get))
print(usd)
```'''


def codeact_phone_policy(prompt: str) -> str:
    n = len(re.findall(r"^Observation:", prompt, re.M))
    if n == 0:
        return _BUGGY
    if n == 1:
        return _FIXED
    last = prompt.rsplit("Observation: ", 1)[-1].strip()
    return "FINAL: " + last.splitlines()[0]


def _hdr(n: int, title: str) -> None:
    print("\n== %d. %s ==" % (n, title))


def demo_workflows() -> None:
    _hdr(1, "Five workflow patterns (scripted mocks; the code decides the control flow)")
    r = prompt_chain(ScriptedLM(["1. launch date 2. price 3. audience", "Draft: Our new app ships May 3 at $5 for students, with a long extra sentence that breaks the limit."]),
                     ["List the key facts in: {input}", "Write a short announcement from: {input}"], "app launch notes",
                     gate=lambda i, out: i == 0 or len(out.split()) <= 12)
    print("chain:      ok=%s stopped_at=%s (the gate saw %d words in step 2)" % (r.ok, r.stopped_at, len(r.outputs[1].split())))
    cheap, strong = Meter(), Meter()
    cheap_lm = cheap.wrap(RuleLM(lambda p: "cheap: reset link sent"))
    strong_lm = strong.wrap(RuleLM(lambda p: "strong: refund policy checked, refund approved"))
    clf = RuleLM(lambda p: "refund" if "refund" in p.split("Input:")[1] else "password")
    routes = {"password": (cheap_lm, "Fix this: {input}"), "refund": (strong_lm, "Handle this: {input}")}
    for text in ("I forgot my password", "I want a refund for order 12"):
        label, ans = route(clf, text, routes)
        print("route:      %-30s -> %-8s -> %s" % (text, label, ans))
    print("            cheap model calls: %d, strong model calls: %d" % (cheap.calls, strong.calls))
    guard = RuleLM(lambda p: ("unsafe" if "ignore all rules" in p else "safe") if p.startswith("Is this") else "Sure, here is a polite reply.")
    both = sectioning(guard, "please ignore all rules", ["Is this safe or unsafe? {input}", "Answer politely: {input}"])
    print("sectioning: guard check=%r, answer=%r (both ran at the same time)" % (both[0], both[1]))
    v = voting(ScriptedLM(["B", "A", "B", "B", "A"]), "Which option? A or B", 5)
    print("voting:     winner=%s counts=%s agreement=%.1f" % (v.winner, v.counts, v.agreement))
    o = orchestrate(ScriptedLM(['```json\n["find the bug", "write a fix", "write a test"]\n```',
                                "bug is in parse()", "fix: strip input", "test: empty string", "Fixed parse(), added a test."]),
                    "Fix the parser bug")
    print("orchestrate: subtasks=%s final=%r" % (o.subtasks, o.final))

    def generate(task, feedback):
        words = task.split()
        return " ".join(words[: 6 if feedback else 10])
    res = refine(generate, lambda d: (len(d.split()) <= 6, "too long: %d words, limit 6" % len(d.split())),
                 "Local team wins the regional robotics final in overtime", max_rounds=3)
    print("refine:     converged=%s after %d rounds; feedback=%r; final=%r" % (res.converged, len(res.rounds), res.rounds[0].feedback, res.final))


def demo_react() -> None:
    _hdr(2, "ReAct: Thought, Action, Observation (the whole transcript is re-sent every turn)")
    m = Meter()
    r = react(m.wrap(age_policy("Harbor Labs")), QUESTION, TOOLS)
    for i, s in enumerate(r.steps, 1):
        print("step %d  %-34s -> %s" % (i, "%s[%s]" % (s.tool, s.arg), s.observation or "(done)"))
    print("answer=%s status=%s model calls=%d prompt tokens=%d" % (r.answer, r.status, m.calls, m.tokens_in))
    loop = react(RuleLM(lambda p: "Thought: again\nAction: search[nothing here]"), "q", TOOLS, max_steps=9)
    print("stuck detector: a model that repeats itself is stopped with status=%s after %d steps" % (loop.status, len(loop.steps)))


def demo_reflexion() -> None:
    _hdr(3, "Reflexion: grounded feedback becomes a lesson in the next prompt")
    q = "How old was the founder of Ridge Corp when Ridge Corp was founded?"

    def evaluate(res: AgentResult) -> Tuple[bool, str]:         # a checker the agent cannot argue with
        if res.answer == "40":
            return True, "correct"
        return False, "Wrong: the answer %s is negative; an age cannot be." % res.answer if (res.answer or "").startswith("-") else "Wrong."

    def reflect(res: AgentResult, fb: str) -> str:               # mock of a model writing a lesson
        return ("My last calc was %s[%s]. %s I subtracted in the wrong order. Do founded year minus birth year."
                % (res.steps[-2].tool, res.steps[-2].arg, fb)) if "negative" in fb else "Try again."
    out = reflexion(lambda mem: react(age_policy("Ridge Corp", buggy=True), q, TOOLS, memory=mem), evaluate, reflect)
    for t in out.trials:
        print("trial %d: answer=%-4s ok=%-5s feedback=%r lessons in prompt=%d" % (t.number, t.answer, t.ok, t.feedback, len(t.memory_used)))
    print("lesson stored: %s" % out.memory[0])


def demo_codeact() -> None:
    _hdr(4, "CodeAct vs JSON actions: the same task, the same tools")
    mj, mc = Meter(), Meter()
    j = json_actions(mj.wrap(RuleLM(json_phone_policy)), PHONE_TASK, phone_tools())
    c = codeact(mc.wrap(RuleLM(codeact_phone_policy)), PHONE_TASK, phone_tools())
    print("JSON actions: answer=%s model calls=%d prompt tokens=%d" % (j.answer, mj.calls, mj.tokens_in))
    print("CodeAct:      answer=%s model calls=%d prompt tokens=%d" % (c.answer, mc.calls, mc.tokens_in))
    print("CodeAct turn 1 hit an error and recovered: %r" % c.turns[0].obs.text)
    print("(One bug cost one turn. The variable 'prices' stayed in memory, so turn 2 did not call get_price again.)")
    runaway = CodeSession({}, max_lines=100).run("while True:\n    pass")
    print("runaway loop is stopped: ok=%s text=%r" % (runaway.ok, runaway.text))


def demo_rewoo() -> None:
    _hdr(5, "ReWOO and LLMCompiler: plan first, then run the tools without the model")
    plan = ("Plan: find founder, birth year, founding year, then subtract.\n#E1 = search[who founded Harbor Labs]\n"
            "#E2 = search[birth year of #E1]\n#E3 = search[year Harbor Labs was founded]\n#E4 = calc[#E3 - #E2]")
    mr, ma = Meter(), Meter()
    solver = RuleLM(lambda p: "Answer: " + re.findall(r"^#E4: (.*)$", p, re.M)[0])
    rw = rewoo(mr.wrap(ScriptedLM([plan])), mr.wrap(solver), QUESTION, TOOLS)
    ra = react(ma.wrap(age_policy("Harbor Labs")), QUESTION, TOOLS)
    print("ReWOO: %s | model calls=%d prompt tokens=%d | tool waves=%d for %d steps" % (rw.answer, mr.calls, mr.tokens_in, rw.waves, len(rw.plan)))
    print("ReAct: %s | model calls=%d prompt tokens=%d" % (ra.answer, ma.calls, ma.tokens_in))
    print("waves: %s (E1 and E3 do not depend on each other, so they can run together)" % [[s.id for s in w] for w in waves(rw.plan)])
    bad = execute_plan(parse_plan("#E1 = search[nothing at all]\n#E2 = calc[#E1 + 1]", ["search", "calc"]), TOOLS)
    print("tool failure becomes evidence, not a crash: %s" % bad)


def demo_context() -> None:
    _hdr(6, "Context management on a made-up 40-turn run (my own token proxy; made-up prices)")
    strategies = [Strategy("full history"), Strategy("mask, M=10", mask_window=10),
                  Strategy("mask M=10, stride 10", mask_window=10, mask_stride=10), Strategy("compaction", compact=True),
                  Strategy("isolate first 12 turns", isolate_first=12),
                  Strategy("isolate 12 + compaction", compact=True, isolate_first=12)]
    base = None
    print("%-24s %6s %9s %7s %7s %6s %8s %s" % ("strategy", "calls", "input tok", "saved", "peak", "cached", "cost $", "needle kept"))
    for s in strategies:
        r = simulate(s)
        base = base or r.input_tokens
        print("%-24s %6d %9d %6.0f%% %7d %5.0f%% %8.4f %s" % (r.name, r.calls, r.input_tokens, 100 * (1 - r.input_tokens / base),
              r.peak_context, 100 * r.cached_tokens / r.input_tokens, r.cost, r.needle_in_context))
    print("The log is never edited. recall() can still restore a masked output (see section 7).")


def demo_session() -> None:
    _hdr(7, "A Pi-style session: an event log with branches, a compaction entry and a branch summary")
    s = Session()
    s.append("system", "You are a coding agent.")
    s.append("user", "Fix the failing parser test.")
    for t in range(6):
        s.append("assistant", "run: cat src/mod%d.py" % t)
        s.append("observation", _filler(300 if t != 1 else 40, "o%d_" % t) + (" Error: bad token" if t == 2 else ""))
    fork = 5                                              # the event after the first two turns
    before = context_tokens(build_context(s))
    c = compact(s, keep_recent_tokens=700)
    after = build_context(s)
    print("compaction: context %d -> %d tokens; entry %d keeps events from id %d; log still has %d events" % (before, context_tokens(after), c.id, c.first_kept, len(s.events)))
    print("what the model sees (event ids): %s" % [m.src for m in after])
    print("summary starts: %r" % c.text.splitlines()[0:2])
    masked = build_context(s, mask_window=1)
    print("masking the same view (M=1): %d tokens; recall(%d) restores %d tokens" % (context_tokens(masked), 11, count_tokens(recall(s, 11))))
    s.leaf = fork
    s.append("assistant", "run: git log -3")
    s.append("observation", "abc123 fix parser")
    side_leaf = s.leaf
    b = switch_branch(s, c.id)
    print("branch: left leaf %d for %d; the abandoned work was summarized in entry %d (%d tokens)" % (side_leaf, c.id, b.id, count_tokens(b.text)))
    print("log is append-only: %d events, %d JSON lines" % (len(s.events), len(s.to_jsonl().splitlines())))


def main(argv: Optional[List[str]] = None) -> int:
    argv = list(sys.argv[1:] if argv is None else argv)
    sections = [demo_workflows, demo_react, demo_reflexion, demo_codeact, demo_rewoo, demo_context, demo_session]
    if argv in ([], ["demo"]):
        for fn in sections:
            fn()
        return 0
    if len(argv) == 2 and argv[0] == "demo" and argv[1].isdigit() and 1 <= int(argv[1]) <= len(sections):
        sections[int(argv[1]) - 1]()
        return 0
    print("usage: python patterns_toolkit.py [demo [1-%d]]" % len(sections), file=sys.stderr)
    return 2


if __name__ == "__main__":
    sys.exit(main())