import assert from "node:assert/strict";
import { streamText, tool } from "ai";
import { z } from "zod";
import type { HarnessProgressUpdate } from "@/lib/mogplex-api/harness-progress";
import { MockLanguageModelV4 } from "ai/test";
import { runNativeMogplexAgent } from "@/lib/mogplex-api/native-run";
import type { SkillSelectionInput } from "@/lib/decisions/skills";
import { createRunGuidanceSession } from "@/lib/slack/run-guidance-session";
import type { RunGuidance } from "@/lib/slack/run-guidance-store";
import {
  buildAiCall,
  buildRunRow,
} from "../unit/helpers/mogplex-api-runs-fixtures";

export async function exercise(
  mode:
    | "success"
    | "error"
    | "cancelled"
    | "unauthorized"
    | "lease_failure"
    | "tool_progress"
    | "guidance"
    | "workspace_guidance"
    | "agent",
  response = "Fixed the header.",
  metadata: Record<string, unknown> = {}
) {
  let call = buildAiCall({ model: "harness:mogplex", metadata });
  const run = buildRunRow({
    harness: "mogplex",
    metadata,
    ...(mode === "agent" ? { agent_id: "agent-1" } : {}),
    ...(mode === "tool_progress" || mode === "guidance"
      ? {
          metadata: {
            slack_guidance_enabled: mode === "guidance",
            slackRunControls: {
              teamId: "T1",
              channelId: "D1",
              messageTs: "1.2",
            },
          },
        }
      : {}),
  });
  // Workspace guidance steers a run that has no Slack thread at all.
  const steered = mode === "guidance" || mode === "workspace_guidance";
  const slackCardQueues: string[] = [];
  const progress: HarnessProgressUpdate[] = [];
  let guidanceRows: RunGuidance[] = [];
  const guidanceSteps: number[] = [];
  const guidanceReplies: string[] = [];
  let modelStep = 0;
  const controller = new AbortController();
  const events: string[] = [];
  const finalReports: string[] = [];
  const toolEvidence: unknown[] = [];
  let cleaned = false;
  let closed = false;
  let executionLeaseAcquired = false;
  let systemSuffix: string | null | undefined;
  let invokedModel: string | undefined;
  let attachedSkillIds: readonly string[] | undefined;
  const skillChecks: SkillSelectionInput[] = [];
  let skillCheckStartedBeforeStream = false;
  let skillCheckSettled = false;
  const model = new MockLanguageModelV4({
    doStream: async () => ({
      stream: new ReadableStream({
        start(sink) {
          if (mode === "error") {
            sink.enqueue({
              type: "error",
              error: new Error("Provider disconnected"),
            });
          } else if (
            (mode === "tool_progress" || steered) &&
            modelStep++ === 0
          ) {
            sink.enqueue({
              type: "tool-call",
              toolCallId: "progress-1",
              toolName: "report_progress",
              input: JSON.stringify({
                phase: "Verifying",
                summary: "Updated the header layout.",
                next: "Run the regression tests.",
              }),
            });
            sink.enqueue({
              type: "tool-call",
              toolCallId: "command-1",
              toolName: "terminal_exec",
              input: JSON.stringify({ command: "pnpm test" }),
            });
            sink.enqueue({
              type: "finish",
              finishReason: { unified: "tool-calls", raw: "tool-calls" },
              usage: {
                inputTokens: {
                  total: 12,
                  noCache: 12,
                  cacheRead: 0,
                  cacheWrite: 0,
                },
                outputTokens: { total: 5, text: 5, reasoning: 0 },
              },
            });
          } else {
            sink.enqueue({ type: "text-start", id: "text" });
            for (const delta of response) {
              sink.enqueue({ type: "text-delta", id: "text", delta });
            }
            sink.enqueue({ type: "text-end", id: "text" });
            sink.enqueue({
              type: "finish",
              finishReason: { unified: "stop", raw: "stop" },
              usage: {
                inputTokens: {
                  total: 12,
                  noCache: 12,
                  cacheRead: 0,
                  cacheWrite: 0,
                },
                outputTokens: { total: 5, text: 5, reasoning: 0 },
              },
            });
          }
          sink.close();
        },
      }),
    }),
  });
  let caught: unknown;
  let result: { output: string } | undefined;
  try {
    result = await runNativeMogplexAgent(
      run,
      { recordId: "sandbox-record-1", sandboxId: "sbx_123" },
      {
        loadCall: async () => call,
        resolveAgent: async (input) => {
          assert.deepEqual(input, { agentId: "agent-1", userId: run.user_id });
          return {
            id: "agent-1",
            name: "NEXTJS-REVIEWER",
            slug: "nextjs-reviewer",
            model: null,
            systemPrompt: "Review App Router code carefully.",
            skills: [
              {
                id: "skill-1",
                name: "RSC Audit",
                description: null,
                content: "Look for use client.",
              },
            ],
            rules: [
              { id: "rule-1", name: "No any", content: "Never use any." },
            ],
            preset: false,
            teamId: null,
            ownerUserId: run.user_id,
          };
        },
        ensureExecutionLease: async (_run, _sandbox, teamId) => {
          assert.equal(teamId, "team-1");
          if (mode === "lease_failure") throw new Error("Lease refused");
          executionLeaseAcquired = true;
        },
        loadContext: async () => {
          if (mode === "unauthorized")
            throw new Error("Active sandbox not found for this agent run");
          return {
            userId: run.user_id,
            repoId: run.repo_id,
            repoFullName: "example/app",
            repoOwner: "example",
            repoName: "app",
            repoBranch: "fix/header",
            repoBaseBranch: "main",
            sandboxId: "sandbox-record-1",
            teamId: "team-1",
            conversationId: null,
            workspaceSessionId: null,
            surface: "chat",
            enableTools: true,
            toolExecutionIdempotencyKey: call.id,
          };
        },
        resolveModel: async (_userId, requested, surface) => {
          if (metadata.run_origin === "slack") assert.equal(surface, "slack");
          return requested ?? "test/native-model";
        },
        buildMessages: async () => [
          { role: "user", parts: [{ type: "text", text: run.prompt }] },
        ],
        createControl: async () => {
          if (mode === "cancelled") controller.abort();
          return {
            signal: controller.signal,
            isCancelled: () => mode === "cancelled",
            async close() {
              closed = true;
            },
          };
        },
        observeSkills: async (input) => {
          skillChecks.push(input);
          // Slower than the whole run, so only an await can see it settle.
          await new Promise((resolve) => setTimeout(resolve, 40));
          skillCheckSettled = true;
        },
        createStream: async (input) => {
          invokedModel = input.resolvedModel;
          systemSuffix = input.systemSuffix;
          attachedSkillIds = input.attachedSkillIds;
          // The real stream reports the catalog once it has resolved it.
          input.onSkillsResolved?.({ invoked: [], available: [] });
          skillCheckStartedBeforeStream = skillChecks.length > 0;
          assert.equal(
            skillCheckSettled,
            false,
            "the skill check must never be awaited in front of the run"
          );
          assert.equal(
            input.context.sandboxExecution?.retryOnSandboxLoss,
            false
          );
          assert.equal(
            executionLeaseAcquired,
            true,
            "reserve VM lifetime before starting the agent"
          );
          return {
            result: streamText({
              model,
              prompt: run.prompt,
              prepareStep: async ({ messages, stepNumber }) => ({
                messages: input.prepareMessages
                  ? await input.prepareMessages(messages, stepNumber)
                  : messages,
              }),
              abortSignal: input.abortSignal,
              tools: {
                ...input.additionalTools,
                terminal_exec: tool({
                  inputSchema: z.object({ command: z.string() }),
                  execute: async () => {
                    if (steered)
                      guidanceRows = [
                        {
                          id: "00000000-0000-4000-8000-000000000005",
                          run_id: run.id,
                          user_id: run.user_id,
                          ai_call_id: run.ai_call_id,
                          body: "Keep the desktop header unchanged.",
                          status: "received",
                          attachments: null,
                          created_at: new Date(0).toISOString(),
                          delivered_step: null,
                        },
                      ];
                    return {
                      exitCode: 1,
                      stdout: "",
                      stderr: "test failure",
                    };
                  },
                }),
              },
              stopWhen: () => false,
              ...input.hooks,
            }),
            connections: [],
            cleanup: async () => {
              cleaned = true;
            },
          };
        },
        createProgress: () => ({
          async report(update) {
            progress.push(update);
          },
          async flush() {},
        }),
        createGuidance: (row) =>
          createRunGuidanceSession(row, {
            load: async () => guidanceRows,
            deliver: async (receipt) => {
              assert.ok(
                JSON.stringify(model.doStreamCalls.at(-1)?.prompt).includes(
                  "Keep the desktop header unchanged."
                )
              );
              guidanceSteps.push(receipt.step);
              guidanceRows = guidanceRows.map((entry) => ({
                ...entry,
                status: "delivered",
                delivered_step: receipt.step,
              }));
              return receipt.ids.length;
            },
            recordReply: async (input) => {
              guidanceReplies.push(input.text);
              return input.ids.length;
            },
            queue: async (input) => {
              slackCardQueues.push(input.runId);
            },
          }),
        updateCall: async (_id, update) => {
          call = { ...call, ...update };
          return call;
        },
        finishCall: async (_id, update) => {
          call = { ...call, ...update };
          return call;
        },
        cancelCall: async (_id, update) => {
          call = { ...call, ...update };
          return call;
        },
        appendEvent: async (event) => {
          events.push(event.eventType);
          if (event.toolName === "terminal_exec")
            toolEvidence.push(event.payload);
          if (event.payload?.kind === "assistant_final")
            finalReports.push(event.message ?? "");
          return null;
        },
      }
    );
  } catch (error) {
    caught = error;
  }
  return {
    call,
    result,
    caught,
    events,
    finalReports,
    toolEvidence,
    cleaned,
    closed,
    model,
    progress,
    guidanceSteps,
    slackCardQueues,
    guidanceReplies,
    systemSuffix,
    invokedModel,
    attachedSkillIds,
    skillChecks,
    skillCheckStartedBeforeStream,
    skillCheckSettled,
  };
}