from agent.core.parser import TOOL_PATTERN
from agent.tools import dispatch

AGENT_SYSTEM = (
    "You are The Abyss, an AI assistant with acess to tools.\n"
    "WHen you need a tool, output EXACTLY this format on its own line:\n"
    "  [SEARCH: your query]\n"
    "  [CALC: expression]\n"
    "  [CODE: python code]\n"
    "  [FILE: read name | write name content | list]\n"
    "After the observation is show, continue your answer naturally.\n"
    "Never fabricate tool results."
)   

MAX_STEPS = 6


def run_agent(messages: list, generate_fn, **gen_kwargs):
    """
    Runs the ReAct loop. Yields (kind, value) tuples:
      ("token", str)       – stream to user (now emitted incrementally)
      ("tool_call", str)   – tool being called
      ("observation", str) – tool result
      ("done", str)        – full assembled response (for DB storage)
    """
    if any(m.get("role") == "system" for m in messages):
        msgs = list(messages)
        msgs[0] = {
            "role": "system",
            "content": AGENT_SYSTEM + "\n\n" + msgs[0]["content"],
        }
    else:
        msgs = [{"role": "system", "content": AGENT_SYSTEM}] + list(messages)

    assembled: list[str] = []

    for _step in range(MAX_STEPS):
        buffer = ""
        flushed = 0          # chars of buffer already streamed as tokens
        used_tool = False

        for token in generate_fn(msgs, **gen_kwargs):
            buffer += token
            match = TOOL_PATTERN.search(buffer)
            if match:
                before = buffer[flushed:match.start()]
                if before:
                    yield ("token", before)
                    assembled.append(before)

                tool_name = match.group(1)
                tool_args = match.group(2)
                call_str = f"[{tool_name}: {tool_args}]"
                yield ("tool_call", call_str)

                result = dispatch(tool_name, tool_args)
                yield ("observation", result)

                msgs = msgs + [
                    {"role": "assistant", "content": buffer[:match.start()] + call_str},
                    {
                        "role": "user",
                        "content": f"[OBSERVATION: {result}]\nContinue your response.",
                    },
                ]
                used_tool = True
                break

            # No tool yet: stream everything that can't start a "[TOOL ...]" call.
            # A call begins with '[', so hold back from the last '[' onward.
            pending = buffer.rfind("[", flushed)
            safe_upto = pending if pending != -1 else len(buffer)
            if safe_upto > flushed:
                chunk = buffer[flushed:safe_upto]
                yield ("token", chunk)
                assembled.append(chunk)
                flushed = safe_upto

        if not used_tool:
            rest = buffer[flushed:]
            if rest:
                yield ("token", rest)
                assembled.append(rest)
            break

    yield ("done", "".join(assembled))