import { beforeEach, describe, expect, it, vi } from "vitest";
import { spawnSubagentDirect } from "./subagent-spawn.js";

type TestAgentConfig = {
  id?: string;
  workspace?: string;
  subagents?: {
    allowAgents?: string[];
  };
};

type TestConfig = {
  agents?: {
    list?: TestAgentConfig[];
  };
};

const hoisted = vi.hoisted(() => ({
  callGatewayMock: vi.fn(),
  configOverride: {} as Record<string, unknown>,
  registerSubagentRunMock: vi.fn(),
}));

vi.mock("../gateway/call.js", () => ({
  callGateway: (opts: unknown) => hoisted.callGatewayMock(opts),
}));

vi.mock("../config/config.js", async (importOriginal) => {
  const actual = await importOriginal<typeof import("../config/config.js")>();
  return {
    ...actual,
    loadConfig: () => hoisted.configOverride,
  };
});

vi.mock("@mariozechner/pi-ai/oauth", () => ({
  getOAuthApiKey: () => "",
  getOAuthProviders: () => [],
}));

vi.mock("./subagent-registry.js", () => ({
  countActiveRunsForSession: () => 0,
  registerSubagentRun: (args: unknown) => hoisted.registerSubagentRunMock(args),
}));

vi.mock("./subagent-announce.js", () => ({
  buildSubagentSystemPrompt: () => "system-prompt",
}));

vi.mock("./subagent-depth.js", () => ({
  getSubagentDepthFromSessionStore: () => 0,
}));

vi.mock("./model-selection.js", () => ({
  resolveSubagentSpawnModelSelection: () => undefined,
}));

vi.mock("./sandbox/runtime-status.js", () => ({
  resolveSandboxRuntimeStatus: () => ({ sandboxed: false }),
}));

vi.mock("../plugins/hook-runner-global.js", () => ({
  getGlobalHookRunner: () => ({ hasHooks: () => false }),
}));

vi.mock("../utils/delivery-context.js", () => ({
  normalizeDeliveryContext: (value: unknown) => value,
}));

vi.mock("./tools/sessions-helpers.js", () => ({
  resolveMainSessionAlias: () => ({ mainKey: "main", alias: "main" }),
  resolveInternalSessionKey: ({ key }: { key?: string }) => key ?? "agent:main:main",
  resolveDisplaySessionKey: ({ key }: { key?: string }) => key ?? "agent:main:main",
}));

vi.mock("./agent-scope.js", () => ({
  resolveAgentConfig: (cfg: TestConfig, agentId: string) =>
    cfg.agents?.list?.find((entry) => entry.id === agentId),
  resolveAgentWorkspaceDir: (cfg: TestConfig, agentId: string) =>
    cfg.agents?.list?.find((entry) => entry.id === agentId)?.workspace ??
    `/tmp/workspace-${agentId}`,
}));

function createConfigOverride(overrides?: Record<string, unknown>) {
  return {
    session: {
      mainKey: "main",
      scope: "per-sender",
    },
    agents: {
      list: [
        {
          id: "main",
          workspace: "/tmp/workspace-main",
        },
      ],
    },
    ...overrides,
  };
}

function setupGatewayMock() {
  hoisted.callGatewayMock.mockImplementation(
    async (opts: { method?: string; params?: Record<string, unknown> }) => {
      if (opts.method === "sessions.patch") {
        return { ok: true };
      }
      if (opts.method === "sessions.delete") {
        return { ok: true };
      }
      if (opts.method === "agent") {
        return { runId: "run-1" };
      }
      return {};
    },
  );
}

function getRegisteredRun() {
  return hoisted.registerSubagentRunMock.mock.calls.at(0)?.[0] as
    | Record<string, unknown>
    | undefined;
}

describe("spawnSubagentDirect workspace inheritance", () => {
  beforeEach(() => {
    hoisted.callGatewayMock.mockClear();
    hoisted.registerSubagentRunMock.mockClear();
    hoisted.configOverride = createConfigOverride();
    setupGatewayMock();
  });

  it("uses the target agent workspace for cross-agent spawns", async () => {
    hoisted.configOverride = createConfigOverride({
      agents: {
        list: [
          {
            id: "main",
            workspace: "/tmp/workspace-main",
            subagents: {
              allowAgents: ["ops"],
            },
          },
          {
            id: "ops",
            workspace: "/tmp/workspace-ops",
          },
        ],
      },
    });

    const result = await spawnSubagentDirect(
      {
        task: "inspect workspace",
        agentId: "ops",
      },
      {
        agentSessionKey: "agent:main:main",
        agentChannel: "telegram",
        agentAccountId: "123",
        agentTo: "456",
        workspaceDir: "/tmp/requester-workspace",
      },
    );

    expect(result.status).toBe("accepted");
    expect(getRegisteredRun()).toMatchObject({
      workspaceDir: "/tmp/workspace-ops",
    });
  });

  it("preserves the inherited workspace for same-agent spawns", async () => {
    const result = await spawnSubagentDirect(
      {
        task: "inspect workspace",
        agentId: "main",
      },
      {
        agentSessionKey: "agent:main:main",
        agentChannel: "telegram",
        agentAccountId: "123",
        agentTo: "456",
        workspaceDir: "/tmp/requester-workspace",
      },
    );

    expect(result.status).toBe("accepted");
    expect(getRegisteredRun()).toMatchObject({
      workspaceDir: "/tmp/requester-workspace",
    });
  });
});
