import { convexTest } from "convex-test";
import { anyApi, makeFunctionReference } from "convex/server";
import { describe, expect, test } from "vitest";

import type { Id } from "./_generated/dataModel";
import schema from "./schema";

// Execution-only module set: avoids loading workPlanning (and its env read) so
// these tests run without Convex env vars.
const modules = {
  "./_generated/api.ts": () => import("./_generated/api"),
  "./_generated/server.ts": () => import("./_generated/server"),
  "./authz.ts": () => import("./authz"),
  "./workExecution.ts": () => import("./workExecution"),
};
const identity = { tokenIdentifier: "https://convex.test|exec-user" };
const api = anyApi;
const claimAttemptRef = makeFunctionReference<
  "mutation",
  { attemptId: Id<"workAttempts">; leaseMs: number; owner: string },
  { number: number } | null
>("workExecution:claimAttempt");

const makeTest = () => convexTest({ modules, schema });
type TestT = ReturnType<typeof makeTest>;

type Scenario =
  | "success"
  | "transient-failure-then-success"
  | "needs-input"
  | "permanent-failure"
  | "cancelled";

const seedReadyWork = async (options: { secondSlice?: boolean } = {}) => {
  const t = convexTest({ modules, schema });
  const { workId } = await t.withIdentity(identity).run(async (ctx) => {
    const organizationId = await ctx.db.insert("organizations", {
      createdAt: 1,
      createdBy: identity.tokenIdentifier,
      kind: "personal",
      name: "Personal",
    });
    await ctx.db.insert("organizationMembers", {
      createdAt: 1,
      organizationId,
      role: "owner",
      userId: identity.tokenIdentifier,
    });
    const projectId = await ctx.db.insert("projects", {
      createdAt: 1,
      name: "Zopu",
      normalizedSourceUrl: "https://example.com/zopu",
      organizationId,
      repositoryPath: "puter/zopu",
      sourceHost: "example.com",
      sourceUrl: "https://example.com/zopu",
      updatedAt: 1,
    });
    const id = await ctx.db.insert("works", {
      createdAt: 1,
      definitionApprovalVersion: 1,
      definitionVersion: 1,
      designApprovalVersion: 1,
      designVersion: 1,
      objective: "Prove execution",
      organizationId,
      projectId,
      status: "ready",
      title: "Executable work",
      updatedAt: 1,
    });
    await ctx.db.insert("designPackets", {
      createdAt: 1,
      createdBy: "test",
      definitionVersion: 1,
      payloadJson: "{}",
      status: "current",
      version: 1,
      workId: id,
    });
    await ctx.db.insert("workSlices", {
      createdAt: 1,
      designVersion: 1,
      objective: "execute the slice",
      observableBehavior: "terminal run",
      ordinal: 0,
      payloadJson: "{}",
      sliceId: "slice-1",
      status: "ready",
      title: "first",
      workId: id,
    });
    if (options.secondSlice) {
      await ctx.db.insert("workSlices", {
        createdAt: 1,
        designVersion: 1,
        objective: "execute the second slice",
        observableBehavior: "second terminal run",
        ordinal: 1,
        payloadJson: "{}",
        sliceId: "slice-2",
        status: "planned",
        title: "second",
        workId: id,
      });
    }
    return { workId: id };
  });
  return { t, workId };
};

const runAttempt = async (
  t: TestT,
  attemptId: Id<"workAttempts">,
  scenario: Scenario
) => t.action(api.workExecution.executeFakeAttempt, { attemptId, scenario });

const attemptNumber = async (
  t: TestT,
  runId: Id<"workRuns">,
  number: number
): Promise<Id<"workAttempts">> => {
  const attempt = await t.run(async (ctx) =>
    ctx.db
      .query("workAttempts")
      .withIndex("by_runId_and_number", (q) =>
        q.eq("runId", runId).eq("number", number)
      )
      .unique()
  );
  if (!attempt) {
    throw new Error(`attempt ${number} not found`);
  }
  return attempt._id;
};

const snapshot = async (t: TestT, workId: Id<"works">, runId: Id<"workRuns">) =>
  t.run(async (ctx) => ({
    attempts: await ctx.db
      .query("workAttempts")
      .withIndex("by_runId_and_number", (q) => q.eq("runId", runId))
      .collect(),
    run: await ctx.db.get(runId),
    slice: await ctx.db
      .query("workSlices")
      .withIndex("by_workId_and_designVersion", (q) =>
        q.eq("workId", workId).eq("designVersion", 1)
      )
      .first(),
    slices: await ctx.db
      .query("workSlices")
      .withIndex("by_workId_and_designVersion", (q) =>
        q.eq("workId", workId).eq("designVersion", 1)
      )
      .collect(),
    work: await ctx.db.get(workId),
  }));

describe("simulated execution resolver", () => {
  test("success settles Work as completed and marks the slice completed", async () => {
    const { t, workId } = await seedReadyWork();
    const started = await t
      .withIdentity(identity)
      .mutation(api.workExecution.startSimulatedExecution, {
        scenario: "success",
        sliceId: "slice-1",
        workId,
      });
    await runAttempt(t, started.attemptId, "success");
    const state = await snapshot(t, workId, started.runId);
    expect(state.work?.status).toBe("completed");
    expect(state.run?.terminalClassification).toBe("Succeeded");
    expect(state.attempts).toHaveLength(1);
    expect(state.slice?.status).toBe("completed");
    const decisions = await t.run(async (ctx) =>
      ctx.db.query("resolverDecisions").collect()
    );
    expect(decisions).toHaveLength(1);
    expect(decisions[0]).toMatchObject({
      classification: "Succeeded",
      decision: "terminal",
      resultingWorkStatus: "completed",
    });
  });

  test("completes Work only after every approved slice succeeds", async () => {
    const { t, workId } = await seedReadyWork({ secondSlice: true });
    const first = await t
      .withIdentity(identity)
      .mutation(api.workExecution.startSimulatedExecution, {
        scenario: "success",
        workId,
      });
    await runAttempt(t, first.attemptId, "success");
    const afterFirst = await snapshot(t, workId, first.runId);
    expect(afterFirst.work?.status).toBe("ready");
    expect(
      afterFirst.slices.find((slice) => slice.sliceId === "slice-1")?.status
    ).toBe("completed");
    expect(
      afterFirst.slices.find((slice) => slice.sliceId === "slice-2")?.status
    ).toBe("ready");
    await expect(
      t
        .withIdentity(identity)
        .mutation(api.workExecution.retrySimulatedExecution, {
          runId: first.runId,
        })
    ).rejects.toThrow(/retryable state/u);

    const second = await t
      .withIdentity(identity)
      .mutation(api.workExecution.startSimulatedExecution, {
        scenario: "success",
        sliceId: "slice-2",
        workId,
      });
    await runAttempt(t, second.attemptId, "success");
    const completed = await snapshot(t, workId, second.runId);
    expect(completed.work?.status).toBe("completed");
    expect(
      completed.slices.every((slice) => slice.status === "completed")
    ).toBe(true);
  });

  test("allows only one owner to claim a queued attempt", async () => {
    const { t, workId } = await seedReadyWork();
    const started = await t
      .withIdentity(identity)
      .mutation(api.workExecution.startSimulatedExecution, {
        scenario: "success",
        workId,
      });
    const first = await t.mutation(claimAttemptRef, {
      attemptId: started.attemptId,
      leaseMs: 60_000,
      owner: "worker-a",
    });
    const second = await t.mutation(claimAttemptRef, {
      attemptId: started.attemptId,
      leaseMs: 60_000,
      owner: "worker-b",
    });
    expect(first?.number).toBe(1);
    expect(second).toBeNull();
  });

  test("transient failure retries within the kit budget and then succeeds", async () => {
    const { t, workId } = await seedReadyWork();
    const started = await t
      .withIdentity(identity)
      .mutation(api.workExecution.startSimulatedExecution, {
        scenario: "transient-failure-then-success",
        workId,
      });
    await runAttempt(t, started.attemptId, "transient-failure-then-success");
    // resolver auto-scheduled attempt 2; drive it manually in the test runtime
    const second = await attemptNumber(t, started.runId, 2);
    await runAttempt(t, second, "transient-failure-then-success");
    const state = await snapshot(t, workId, started.runId);
    expect(state.work?.status).toBe("completed");
    expect(state.run?.terminalClassification).toBe("Succeeded");
    expect(state.attempts).toHaveLength(2);
    expect(state.attempts[0]?.classification).toBe("RetryableFailure");
  });

  test("permanent failure settles Work as failed, not silently runnable", async () => {
    const { t, workId } = await seedReadyWork();
    const started = await t
      .withIdentity(identity)
      .mutation(api.workExecution.startSimulatedExecution, {
        scenario: "permanent-failure",
        workId,
      });
    await runAttempt(t, started.attemptId, "permanent-failure");
    const state = await snapshot(t, workId, started.runId);
    expect(state.work?.status).toBe("failed");
    expect(state.run?.terminalClassification).toBe("PermanentFailure");
  });

  test("needs-input surfaces a blocker instead of looping", async () => {
    const { t, workId } = await seedReadyWork();
    const started = await t
      .withIdentity(identity)
      .mutation(api.workExecution.startSimulatedExecution, {
        scenario: "needs-input",
        workId,
      });
    await runAttempt(t, started.attemptId, "needs-input");
    const state = await snapshot(t, workId, started.runId);
    expect(state.work?.status).toBe("needs-input");
    expect(state.run?.terminalClassification).toBe("NeedsInput");
  });

  test("cancellation terminates the attempt and returns Work to ready", async () => {
    const { t, workId } = await seedReadyWork();
    const started = await t
      .withIdentity(identity)
      .mutation(api.workExecution.startSimulatedExecution, {
        scenario: "success",
        workId,
      });
    await t
      .withIdentity(identity)
      .mutation(api.workExecution.cancelSimulatedExecution, {
        runId: started.runId,
      });
    const state = await snapshot(t, workId, started.runId);
    expect(state.work?.status).toBe("ready");
    expect(state.run?.status).toBe("cancelled");
    expect(state.slice?.status).toBe("ready");
  });

  test("a slice that is not part of the current approved Design is rejected", async () => {
    const { t, workId } = await seedReadyWork();
    await expect(
      t
        .withIdentity(identity)
        .mutation(api.workExecution.startSimulatedExecution, {
          scenario: "success",
          sliceId: "not-a-real-slice",
          workId,
        })
    ).rejects.toThrow(/current approved Design/u);
  });

  test("manual retry restarts a failed Run", async () => {
    const { t, workId } = await seedReadyWork();
    const started = await t
      .withIdentity(identity)
      .mutation(api.workExecution.startSimulatedExecution, {
        scenario: "permanent-failure",
        workId,
      });
    await runAttempt(t, started.attemptId, "permanent-failure");
    expect((await snapshot(t, workId, started.runId)).work?.status).toBe(
      "failed"
    );
    const retried = await t
      .withIdentity(identity)
      .mutation(api.workExecution.retrySimulatedExecution, {
        runId: started.runId,
      });
    await runAttempt(t, retried.attemptId as Id<"workAttempts">, "success");
    const state = await snapshot(t, workId, started.runId);
    expect(state.work?.status).toBe("completed");
    expect(retried.number).toBe(2);
  });

  test("reconciling an expired lease resumes within budget", async () => {
    const { t, workId } = await seedReadyWork();
    const started = await t
      .withIdentity(identity)
      .mutation(api.workExecution.startSimulatedExecution, {
        scenario: "success",
        workId,
      });
    // simulate a worker that claimed but died before finishing
    await t.run(async (ctx) => {
      await ctx.db.patch(started.attemptId, {
        leaseExpiresAt: Date.now() - 1000,
        leaseOwner: "dead-worker",
        status: "claimed",
      });
    });
    await t.mutation(api.workExecution.reconcileExpiredAttempts, {});
    const after = await snapshot(t, workId, started.runId);
    expect(after.attempts.find((a) => a.number === 1)?.status).toBe("terminal");
    expect(
      after.attempts.find((a) => a.number === 2 && a.status === "queued")
    ).toBeTruthy();
  });
});