import { mkdtempSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import type { AgentProfile } from "../agent/profile-registry.js";
import { afterEach, describe, expect, it, vi } from "vitest";
import type { ThreadId, Turn, TurnId } from "@best-agent/protocol";
import { CoordinatorExecutor } from "./coordinator-executor.js";

function createProfile(name: string): AgentProfile {
  return {
    allowedTools: [],
    description: `${name} profile`,
    name,
    source: "user",
  };
}

function createTurn(): Turn {
  return {
    createdAt: new Date().toISOString(),
    id: "turn-1" as TurnId,
    status: "running",
    threadId: "thread-1" as ThreadId,
  };
}

describe("CoordinatorExecutor JSON retry", () => {
  const scratchDirs: string[] = [];

  afterEach(() => {
    while (scratchDirs.length > 0) {
      const dir = scratchDirs.pop();
      if (dir) {
        rmSync(dir, { force: true, recursive: true });
      }
    }
  });

  it("retries plan parsing once when initial response is invalid JSON", async () => {
    const responses = [
      "not-json",
      '{"plannedIterations":1,"workers":[{"id":"worker-1","profile":"explore","ownershipMode":"read-only","objective":"inspect code"}]}',
      '{"decision":"complete","finalOutput":"plan retry recovered"}',
    ];
    const executeCoordinatorModelTurn = vi.fn(async (input: unknown) => {
      void input;
      return responses.shift() ?? "";
    });
    const planProfile = createProfile("plan");
    const profiles = new Map<string, AgentProfile>([
      ["plan", planProfile],
      ["explore", createProfile("explore")],
    ]);
    const executor = new CoordinatorExecutor({
      cleanupScratchRun: async () => {},
      createScratchRun: () => {
        const scratchDir = mkdtempSync(join(tmpdir(), "best-agent-coordinator-"));
        scratchDirs.push(scratchDir);
        return { scratchDir, workflowRunId: "wf-1" };
      },
      defaultOwnershipModeForProfile: () => "read-only",
      emitCoordinatorComplete: () => {},
      emitCoordinatorPlan: () => {},
      emitCoordinatorSynthesize: () => {},
      executeCoordinatorModelTurn,
      executeWorker: async (input) => ({
        childThreadId: "child-1" as ThreadId,
        objective: input.workerObjective,
        profileName: input.profile.name,
        status: "completed",
        workerId: input.workerId,
      }),
      resolveProfile: (profileName) => profiles.get(profileName),
    });

    const output = await executor.execute({
      activeSkillIds: [],
      hostProfile: createProfile("host"),
      prompt: "workflow:coordinate fix runtime",
      turn: createTurn(),
    });

    expect(output).toBe("plan retry recovered");
    expect(executeCoordinatorModelTurn).toHaveBeenCalledTimes(3);
    const retryInput = executeCoordinatorModelTurn.mock.calls[1]?.[0] as
      | { extraMessages: Array<{ content: string }> }
      | undefined;
    expect(retryInput?.extraMessages).toHaveLength(2);
    expect(retryInput?.extraMessages[1]?.content).toContain("invalid JSON");
  });

  it("retries synthesis parsing once when initial response is invalid JSON", async () => {
    const responses = [
      '{"plannedIterations":1,"workers":[{"id":"worker-1","profile":"explore","ownershipMode":"read-only","objective":"inspect code"}]}',
      "```json\n{invalid}\n```",
      '{"decision":"complete","finalOutput":"synthesis retry recovered"}',
    ];
    const executeCoordinatorModelTurn = vi.fn(async (input: unknown) => {
      void input;
      return responses.shift() ?? "";
    });
    const planProfile = createProfile("plan");
    const profiles = new Map<string, AgentProfile>([
      ["plan", planProfile],
      ["explore", createProfile("explore")],
    ]);
    const executor = new CoordinatorExecutor({
      cleanupScratchRun: async () => {},
      createScratchRun: () => {
        const scratchDir = mkdtempSync(join(tmpdir(), "best-agent-coordinator-"));
        scratchDirs.push(scratchDir);
        return { scratchDir, workflowRunId: "wf-1" };
      },
      defaultOwnershipModeForProfile: () => "read-only",
      emitCoordinatorComplete: () => {},
      emitCoordinatorPlan: () => {},
      emitCoordinatorSynthesize: () => {},
      executeCoordinatorModelTurn,
      executeWorker: async (input) => ({
        childThreadId: "child-2" as ThreadId,
        objective: input.workerObjective,
        profileName: input.profile.name,
        status: "completed",
        workerId: input.workerId,
      }),
      resolveProfile: (profileName) => profiles.get(profileName),
    });

    const output = await executor.execute({
      activeSkillIds: [],
      hostProfile: createProfile("host"),
      prompt: "workflow:coordinate fix runtime",
      turn: createTurn(),
    });

    expect(output).toBe("synthesis retry recovered");
    expect(executeCoordinatorModelTurn).toHaveBeenCalledTimes(3);
    const retryInput = executeCoordinatorModelTurn.mock.calls[2]?.[0] as
      | { extraMessages: Array<{ content: string }> }
      | undefined;
    expect(retryInput?.extraMessages).toHaveLength(2);
    expect(retryInput?.extraMessages[1]?.content).toContain("invalid JSON");
  });
});
