import { afterEach, describe, expect, it } from "bun:test";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { Agent } from "../src/agent/agent.ts";
import type { AgentConfig } from "../src/agent/parser.ts";
import type { StreamEvent } from "../src/agent/types.ts";
import { createMockModel, textStreamResult, toolCallStreamResult } from "./helpers/mock-provider.ts";

const testDir = join(tmpdir(), `mozart-wf-msg-${process.pid}-${Date.now()}`);

function makeConfig(name: string): AgentConfig {
  return {
    model: "test/model",
    identity: `I am ${name}.`,
    name,
    ifThen: [],
    schedule: [],
    skills: [],
    sanctums: [],
    sourcePath: `/tmp/${name}.soul`,
  };
}

function dataDir(name: string): string {
  return join(testDir, name);
}

async function collectEvents(gen: AsyncGenerator<StreamEvent>): Promise<StreamEvent[]> {
  const events: StreamEvent[] = [];
  for await (const e of gen) events.push(e);
  return events;
}

function collectText(events: StreamEvent[]): string {
  return events
    .filter((e) => e.type === "text")
    .map((e) => e.text ?? "")
    .join("");
}

describe("workflow: inter-agent messaging", () => {
  const agents: Agent[] = [];

  afterEach(() => {
    for (const a of agents) {
      try {
        a.destroy();
      } catch {}
    }
    agents.length = 0;
  });

  it("agent-a sends to agent-b and receives response", async () => {
    const modelA = createMockModel({
      streamResponses: [
        toolCallStreamResult([{ name: "send", arguments: { agent_id: "agent-b", message: "ping" } }]),
        textStreamResult("Agent B said: pong"),
      ],
    });
    const modelB = createMockModel({
      streamResponses: [textStreamResult("pong")],
    });

    const agentA = new Agent(makeConfig("agent-a"), modelA, dataDir("agent-a"));
    const agentB = new Agent(makeConfig("agent-b"), modelB, dataDir("agent-b"));
    agents.push(agentA, agentB);

    await agentA.initialize();
    await agentB.initialize();

    const router = async (fromId: string, toId: string, message: string): Promise<string> => {
      if (toId === "agent-b") {
        const events = await collectEvents(agentB.handleMessage(fromId, message, `route-${Date.now()}`));
        return collectText(events);
      }
      return `Agent '${toId}' not found`;
    };

    agentA.setRouter(router);
    agentB.setRouter(router);

    const events = await collectEvents(agentA.handleMessage("user", "Ask agent-b to respond", "conv-msg-1"));

    const sendUse = events.find((e) => e.type === "tool_use" && e.toolUse?.name === "send");
    expect(sendUse).toBeDefined();

    const sendResult = events.find((e) => e.type === "tool_result");
    expect(sendResult).toBeDefined();
    expect(sendResult!.text).toContain("pong");

    expect(events.some((e) => e.type === "done")).toBe(true);
    expect(modelB.doStreamCalls).toHaveLength(1);
  });

  it("router rejection surfaces as tool error", async () => {
    const modelA = createMockModel({
      streamResponses: [
        toolCallStreamResult([{ name: "send", arguments: { agent_id: "agent-b", message: "hello" } }]),
        textStreamResult("Could not reach the other agent."),
      ],
    });

    const agentA = new Agent(makeConfig("agent-a"), modelA, dataDir("agent-a"));
    agents.push(agentA);
    await agentA.initialize();

    agentA.setRouter(async () => {
      throw new Error("Connection refused");
    });

    const events = await collectEvents(agentA.handleMessage("user", "Send a message", "conv-msg-2"));

    const toolResult = events.find((e) => e.type === "tool_result");
    expect(toolResult).toBeDefined();
    expect(toolResult!.text).toContain("Error");

    expect(events.some((e) => e.type === "done")).toBe(true);
  });

  it("unknown agent returns an error via router response string", async () => {
    const modelA = createMockModel({
      streamResponses: [
        toolCallStreamResult([{ name: "send", arguments: { agent_id: "ghost", message: "hey" } }]),
        textStreamResult("That agent doesn't exist."),
      ],
    });

    const agentA = new Agent(makeConfig("agent-a"), modelA, dataDir("agent-a"));
    agents.push(agentA);
    await agentA.initialize();

    agentA.setRouter(async (_from, toId) => { throw new Error(`Agent '${toId}' not found`); });

    const events = await collectEvents(agentA.handleMessage("user", "Talk to ghost", "conv-msg-3"));

    const toolResult = events.find((e) => e.type === "tool_result");
    expect(toolResult).toBeDefined();
    expect(toolResult!.text).toContain("Error");
    expect(toolResult!.text).toContain("not found");

    expect(events.some((e) => e.type === "done")).toBe(true);
  });

  it("no router set returns error", async () => {
    const modelA = createMockModel({
      streamResponses: [
        toolCallStreamResult([{ name: "send", arguments: { agent_id: "agent-b", message: "hi" } }]),
        textStreamResult("No router available."),
      ],
    });

    const agentA = new Agent(makeConfig("agent-a"), modelA, dataDir("agent-a"));
    agents.push(agentA);
    await agentA.initialize();

    const events = await collectEvents(agentA.handleMessage("user", "Try sending", "conv-msg-4"));

    const toolResult = events.find((e) => e.type === "tool_result");
    expect(toolResult).toBeDefined();
    expect(toolResult!.text).toContain("Error");
    expect(toolResult!.text).toContain("No router");
  });

  it("both agents log messages to their memory stores", async () => {
    const modelA = createMockModel({
      streamResponses: [
        toolCallStreamResult([{ name: "send", arguments: { agent_id: "agent-b", message: "data-check" } }]),
        textStreamResult("Done."),
      ],
    });
    const modelB = createMockModel({
      streamResponses: [textStreamResult("ack")],
    });

    const agentA = new Agent(makeConfig("agent-a"), modelA, dataDir("agent-a"));
    const agentB = new Agent(makeConfig("agent-b"), modelB, dataDir("agent-b"));
    agents.push(agentA, agentB);

    await agentA.initialize();
    await agentB.initialize();

    agentA.setRouter(async (fromId, toId, message) => {
      if (toId === "agent-b") {
        const events = await collectEvents(agentB.handleMessage(fromId, message, "route-bi"));
        return collectText(events);
      }
      return "not found";
    });

    await collectEvents(agentA.handleMessage("user", "Send data", "conv-bi"));

    expect(agentA.memory.getMessageCount()).toBeGreaterThan(0);
    expect(agentB.memory.getMessageCount()).toBeGreaterThan(0);
  });
});