"""
clankerDiffusion — tool-using agent (ReAct loop) with mid-response RAG.

The model emits <tool name="...">args</tool> blocks. We parse + execute
them locally (tools.py) and feed the <result> back, looping until the
model stops calling tools. Each turn alternates AR <-> diffusion, so the
agent is literally "sometimes normal, sometimes diffusion".

Knowledge injection happens in TWO ways, both mid-response:
  * model-driven : the model calls <tool name="retrieve">q</tool>; we run the
    retriever and feed <result> back, then keep generating.
  * controller-driven : the RAG controller watches the partial response, builds
    a query from the user question + what was just said, and if it finds
    relevant passages it injects <context>...</context> straight into the
    stream and the model continues -- no need for the model to ask.
"""
import argparse
import infer
from tools import execute_tool, parse_tool_calls
from rag import default_kb


SYSTEM = ("You are clanker, a helpful assistant that can THINK, USE TOOLS, and "
          "READ CONTEXT. When you need to compute, inspect files, or look up "
          "knowledge, emit a tool call like "
          "<tool name=\"calc\">17 * 23</tool>, "
          "<tool name=\"python\">print(2**10)</tool>, "
          "<tool name=\"read_file\">notes.txt</tool>, "
          "<tool name=\"list_dir\">.</tool>, or "
          "<tool name=\"retrieve\">the capital of France</tool>. "
          "Available tools: calc(expr), python(code), read_file(path), "
          "list_dir(path), retrieve(query). If <context>...</context> is "
          "provided, use it to answer. After a <result>...</result> appears, "
          "continue and give the final answer. You may reason in <think>...</think>.")


def run_agent(prompt, max_turns=8, mode_cycle=("diff", "ar"),
              max_new=200, gen_len=160, steps=24,
              rag=True, rag_k=3, rag_budget=2):
    model, tok = infer._MODEL, infer._TOK
    kb = default_kb()
    injected = set()

    ctx = [tok.bos_id] + tok.encode(
        f"<system>{SYSTEM}</system><user>{prompt}</user><assistant>")

    print(f"\n=== USER ===\n{prompt}\n")
    for turn in range(max_turns):
        mode = mode_cycle[turn % len(mode_cycle)]
        prev_len = len(ctx)
        print(f"--- turn {turn} [{mode}] ---", flush=True)
        if mode == "ar":
            ids = infer.generate_ar(model, tok, ctx, max_new=max_new)
        else:
            ids = infer.generate_diff(model, tok, ctx, gen_len=gen_len, steps=steps)
        new_ids = ids[prev_len:]              # ONLY the newly generated tokens
        ctx = ids                             # full updated context (no dup)
        text_new = tok.decode(new_ids)
        print(text_new, flush=True)

        # 1) model-driven tool calls (incl. retrieve) -> mid-response injection
        calls = parse_tool_calls(text_new)
        if calls:
            for name, arg in calls:
                print(f"  [tool] {name}({arg[:120]})", flush=True)
                res = execute_tool(name, arg)
                print(f"  [result] {res[:300]}", flush=True)
                ctx = ctx + tok.encode(f"<result>{res}</result>")
            continue                          # keep going within the same answer

        # 2) controller-driven RAG: inject context mid-response, then continue
        if rag and len(injected) < rag_budget * rag_k:
            q = prompt + " " + text_new[-300:]
            fresh = [p for p in kb.retrieve(q, k=rag_k) if p not in injected]
            if fresh:
                injected.update(fresh)
                block = "<context>\n" + "\n---\n".join(fresh) + "\n</context>"
                ctx = ctx + tok.encode(block)
                print(f"  [rag] injected {len(fresh)} passage(s) mid-response",
                      flush=True)
                continue

        break                                 # nothing to do -> answer is final

    final = tok.decode(ctx)
    i = final.rfind("<assistant>")
    print("\n=== clanker (final) ===\n" + (final[i + 10:] if i >= 0 else final))
    return final


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("prompt", nargs="?", default="What is 17 * 23, and what files are in the current folder?")
    ap.add_argument("--ckpt", default=None)
    ap.add_argument("--turns", type=int, default=8)
    ap.add_argument("--no-rag", action="store_true")
    a = ap.parse_args()
    infer.init(ckpt_path=a.ckpt)
    run_agent(a.prompt, max_turns=a.turns, rag=not a.no_rag)