import { convexTest } from 'convex-test';
import { describe, expect, test, vi } from 'vitest';

import type { Id } from '../../convex/_generated/dataModel.js';
import { internal } from '../../convex/_generated/api.js';
import schema from '../../convex/schema';

const modules = import.meta.glob('../../convex/**/*.*s');

type ActiveStatus = 'claimed' | 'preparing' | 'running' | 'checks_running';

const seedJob = async (
  t: ReturnType<typeof convexTest>,
  args: {
    status?: ActiveStatus;
    claimedAt?: number;
    lastHeartbeatAt?: number;
    workspaceStatus?: string;
    completedAt?: number;
    threadStatus?: 'running' | 'resolved' | 'ignored' | 'waiting_for_user';
    threadResolvedAt?: number;
  },
) =>
  await t.mutation(async (ctx) => {
    const now = Date.now();
    const ownerId = await ctx.db.insert('users', {
      email: 'owner@example.com',
      name: 'Owner',
    });
    const spoonId = await ctx.db.insert('spoons', {
      ownerId,
      name: 'Recovery Spoon',
      provider: 'github',
      upstreamOwner: 'upstream',
      upstreamRepo: 'recovery',
      upstreamDefaultBranch: 'main',
      upstreamUrl: 'https://github.com/upstream/recovery',
      forkOwner: 'team',
      forkRepo: 'recovery-spoon',
      forkUrl: 'https://github.com/team/recovery-spoon',
      visibility: 'private',
      maintenanceMode: 'watch',
      syncCadence: 'daily',
      productionRefStrategy: 'default_branch',
      status: 'active',
      createdAt: now,
      updatedAt: now,
    });
    const requestId = await ctx.db.insert('agentRequests', {
      spoonId,
      ownerId,
      prompt: 'Recover the worker job',
      status: 'running',
      createdAt: now,
      updatedAt: now,
    });
    const threadId = await ctx.db.insert('threads', {
      ownerId,
      spoonId,
      title: 'Recover the worker job',
      source: 'user_request',
      status: args.threadStatus ?? 'running',
      priority: 'normal',
      relatedAgentRequestId: requestId,
      ...(args.threadResolvedAt === undefined
        ? {}
        : { resolvedAt: args.threadResolvedAt }),
      createdAt: now,
      updatedAt: now,
    });
    const jobId = await ctx.db.insert('agentJobs', {
      spoonId,
      ownerId,
      agentRequestId: requestId,
      threadId,
      status: args.status ?? 'running',
      prompt: 'Recover the worker job',
      runtime: 'opencode',
      workspaceStatus: args.workspaceStatus ?? 'active',
      baseBranch: 'main',
      workBranch: 'spoon/recover-worker-job',
      forkOwner: 'team',
      forkRepo: 'recovery-spoon',
      forkUrl: 'https://github.com/team/recovery-spoon',
      upstreamOwner: 'upstream',
      upstreamRepo: 'recovery',
      selectedSecretIds: [],
      model: 'gpt-5.1-codex',
      reasoningEffort: 'medium',
      claimedBy: 'w1',
      claimedAt: args.claimedAt ?? now,
      ...(args.lastHeartbeatAt === undefined
        ? {}
        : { lastHeartbeatAt: args.lastHeartbeatAt }),
      ...(args.completedAt === undefined
        ? {}
        : { completedAt: args.completedAt }),
      createdAt: now,
      updatedAt: now,
    });
    await ctx.db.patch(requestId, { agentJobId: jobId });
    await ctx.db.patch(threadId, { latestAgentJobId: jobId });
    return { jobId, requestId, threadId };
  });

const readState = async (
  t: ReturnType<typeof convexTest>,
  ids: {
    jobId: Id<'agentJobs'>;
    requestId: Id<'agentRequests'>;
    threadId: Id<'threads'>;
  },
) =>
  await t.run(async (ctx) => ({
    job: await ctx.db.get(ids.jobId),
    request: await ctx.db.get(ids.requestId),
    thread: await ctx.db.get(ids.threadId),
  }));

describe('recoverStaleJobs', () => {
  test('times out a running job with a stale heartbeat', async () => {
    const t = convexTest(schema, modules);
    const ids = await seedJob(t, {
      lastHeartbeatAt: Date.now() - 10 * 60 * 1000,
    });

    const result = await t.mutation(internal.agentJobs.recoverStaleJobs, {});
    const state = await readState(t, ids);

    expect(result).toEqual({ recovered: 1 });
    expect(state.job?.status).toBe('timed_out');
    expect(state.job?.workspaceStatus).toBe('expired');
    expect(state.job?.completedAt).toBeTypeOf('number');
    expect(state.request?.status).toBe('failed');
    expect(state.thread?.status).toBe('failed');
    expect(state.thread?.resolvedAt).toBeTypeOf('number');
  });

  test('leaves a terminal-workspace maintenance job untouched', async () => {
    const t = convexTest(schema, modules);
    const resolvedAt = Date.now() - 5 * 60 * 1000;
    const ids = await seedJob(t, {
      status: 'running',
      workspaceStatus: 'stopped',
      completedAt: resolvedAt,
      threadStatus: 'resolved',
      threadResolvedAt: resolvedAt,
      // stale heartbeat: the worker stopped its heartbeat after applying the
      // maintenance decision, so a naive cron would time this job out
      claimedAt: Date.now() - 10 * 60 * 1000,
      lastHeartbeatAt: Date.now() - 10 * 60 * 1000,
    });

    const result = await t.mutation(internal.agentJobs.recoverStaleJobs, {});
    const state = await readState(t, ids);

    expect(result).toEqual({ recovered: 0 });
    expect(state.job?.status).toBe('running');
    expect(state.job?.workspaceStatus).toBe('stopped');
    expect(state.thread?.status).toBe('resolved');
    expect(state.thread?.resolvedAt).toBe(resolvedAt);
  });

  test('reaps an abandoned-stop job that never completed', async () => {
    const t = convexTest(schema, modules);
    // The user hit "Stop workspace" (markWorkspaceStopped) while the job was
    // still running: workspaceStatus flips to 'stopped' but job.status,
    // completedAt, and the thread are never finalized. The heartbeat then goes
    // stale, so the cron must still reap this job.
    const ids = await seedJob(t, {
      status: 'running',
      workspaceStatus: 'stopped',
      // no completedAt: this is NOT a resolved maintenance job
      threadStatus: 'running',
      claimedAt: Date.now() - 10 * 60 * 1000,
      lastHeartbeatAt: Date.now() - 10 * 60 * 1000,
    });

    const result = await t.mutation(internal.agentJobs.recoverStaleJobs, {});
    const state = await readState(t, ids);

    expect(result).toEqual({ recovered: 1 });
    expect(state.job?.status).toBe('timed_out');
    expect(state.job?.workspaceStatus).toBe('expired');
    expect(state.job?.completedAt).toBeTypeOf('number');
    expect(state.request?.status).toBe('failed');
    expect(state.thread?.status).toBe('failed');
    expect(state.thread?.resolvedAt).toBeTypeOf('number');
  });

  test('leaves a running job with a fresh heartbeat untouched', async () => {
    const t = convexTest(schema, modules);
    const ids = await seedJob(t, { lastHeartbeatAt: Date.now() });

    const result = await t.mutation(internal.agentJobs.recoverStaleJobs, {});
    const state = await readState(t, ids);

    expect(result).toEqual({ recovered: 0 });
    expect(state.job?.status).toBe('running');
    expect(state.request?.status).toBe('running');
    expect(state.thread?.status).toBe('running');
  });

  test('uses claimedAt when a claimed job has no heartbeat', async () => {
    const t = convexTest(schema, modules);
    const ids = await seedJob(t, {
      status: 'claimed',
      claimedAt: Date.now() - 10 * 60 * 1000,
    });

    const result = await t.mutation(internal.agentJobs.recoverStaleJobs, {});
    const state = await readState(t, ids);

    expect(result).toEqual({ recovered: 1 });
    expect(state.job?.status).toBe('timed_out');
  });

  test('times out a heartbeat exactly at the default threshold', async () => {
    vi.useFakeTimers();
    try {
      const now = new Date('2026-07-10T12:00:00.000Z');
      vi.setSystemTime(now);
      const t = convexTest(schema, modules);
      const ids = await seedJob(t, {
        lastHeartbeatAt: now.getTime() - 3 * 60 * 1000,
      });

      const result = await t.mutation(internal.agentJobs.recoverStaleJobs, {});
      const state = await readState(t, ids);

      expect(result).toEqual({ recovered: 1 });
      expect(state.job?.status).toBe('timed_out');
    } finally {
      vi.useRealTimers();
    }
  });

  test.each<ActiveStatus>([
    'claimed',
    'preparing',
    'running',
    'checks_running',
  ])('recovers stale %s jobs', async (status) => {
    const t = convexTest(schema, modules);
    const ids = await seedJob(t, {
      status,
      lastHeartbeatAt: Date.now() - 10 * 60 * 1000,
    });

    const result = await t.mutation(internal.agentJobs.recoverStaleJobs, {});
    const state = await readState(t, ids);

    expect(result).toEqual({ recovered: 1 });
    expect(state.job?.status).toBe('timed_out');
  });
});
