import type { AiProvider, ChatMessage } from "@stc/contracts";
import type { ToolSet } from "./tools";

export type AssistantEvent =
  | { type: "text"; text: string }
  | { type: "tool"; name: string }
  | { type: "error"; message: string }
  | { type: "done"; text: string };

/**
 * Runs the model with tools until it answers without calling a tool, or until maxSteps model calls.
 * Mutates `messages` by appending the tool round trips, so callers can inspect them.
 */
export async function* runAssistant(o: {
  provider: AiProvider;
  messages: ChatMessage[];
  tools: ToolSet;
  signal?: AbortSignal;
  maxSteps?: number;
}): AsyncIterable<AssistantEvent> {
  const maxSteps = o.maxSteps ?? 6;
  let answer = "";
  for (let step = 0; step < maxSteps; step++) {
    let text = "";
    const calls: { id: string; name: string; input: unknown }[] = [];
    try {
      for await (const ev of o.provider.chat(o.messages, { tools: o.tools.specs, signal: o.signal })) {
        if (ev.text) { text += ev.text; yield { type: "text", text: ev.text }; }
        if (ev.toolCall) calls.push(ev.toolCall);
      }
    } catch (e) {
      yield { type: "error", message: (e as Error).message };
      return;
    }
    answer += text;
    if (!calls.length) {
      yield { type: "done", text: answer };
      return;
    }
    if (text) o.messages.push({ role: "assistant", content: text });
    for (const c of calls) o.messages.push({ role: "assistant", content: JSON.stringify(c.input ?? {}), toolCallId: c.id, name: c.name });
    for (const c of calls) {
      yield { type: "tool", name: c.name };
      let result: string;
      try { result = await o.tools.run(c.name, c.input); } catch (e) { result = `Tool error: ${(e as Error).message}`; }
      o.messages.push({ role: "tool", content: result, toolCallId: c.id, name: c.name });
    }
  }
  const note = "\n\n(I stopped after several lookups. Ask again to continue.)";
  yield { type: "text", text: note };
  yield { type: "done", text: answer + note };
}