import { and, eq, inArray, isNull, lt, notLike, or, sql } from 'drizzle-orm';
import type { DrizzleClient } from '@/server/db/client.js';
import { llmJobs } from '@/server/db/schema.js';

const RETRYABLE = ['FAILED', 'TIMED_OUT'] as const;
const CANCELLABLE = ['PENDING', 'RUNNING'] as const;
const BULK_RETRY_CAP = 500;

export async function retryLlmJobWrite(db: DrizzleClient, id: string): Promise<boolean> {
  const rows = await db
    .update(llmJobs)
    .set({
      status: 'PENDING',
      lastError: null,
      completedAt: null,
      startedAt: null,
      outputPayload: null,
      decision: null,
      flagReason: null,
      retryCount: sql`${llmJobs.retryCount} + 1`,
    })
    .where(
      and(eq(llmJobs.id, id), inArray(llmJobs.status, [...RETRYABLE]), lt(llmJobs.retryCount, 3)),
    )
    .returning({ id: llmJobs.id });
  return rows.length > 0;
}

export async function cancelLlmJobWrite(db: DrizzleClient, id: string): Promise<boolean> {
  const rows = await db
    .update(llmJobs)
    .set({ status: 'FAILED', lastError: 'Cancelled by admin', completedAt: new Date() })
    .where(and(eq(llmJobs.id, id), inArray(llmJobs.status, [...CANCELLABLE])))
    .returning({ id: llmJobs.id });
  return rows.length > 0;
}

export async function retryFailedLlmJobsWrite(
  db: DrizzleClient,
  limit = BULK_RETRY_CAP,
): Promise<string[]> {
  const targets = await db
    .select({ id: llmJobs.id })
    .from(llmJobs)
    .where(
      and(
        inArray(llmJobs.status, [...RETRYABLE]),
        or(isNull(llmJobs.lastError), notLike(llmJobs.lastError, 'Cancelled%')),
      ),
    )
    .limit(limit);
  if (targets.length === limit)
    console.warn(`[llm-jobs] bulk retry capped at ${limit}; more retryable jobs remain`);
  const changed: string[] = [];
  for (const { id } of targets) if (await retryLlmJobWrite(db, id)) changed.push(id);
  return changed;
}