import { and, eq, lt, sql, count } from 'drizzle-orm';
import type { DrizzleClient } from '../client.js';
import { llmJobs } from '../schema.js';
import type { JobCompletionMeta } from '@/server/ai/kinds/types.js';

export type LlmJobType =
  | 'DEAL_MODERATION'
  | 'IMAGE_APPROVAL'
  | 'REVIEW_PRESCORING'
  | 'TRANSLATION'
  | 'VENDOR_VIOLATION';
export type EnqueueLlmJobInput = {
  id?: string;
  jobType: LlmJobType;
  targetId: string;
  targetType: 'DEAL' | 'IMAGE' | 'REVIEW_REMOVAL' | 'VENDOR';
  inputPayload: Record<string, unknown>;
  e2eRunId?: string;
  queueChainOverride?: Array<{ llmProviderId: string; model: string }>;
};

function isRecord(value: unknown): value is Record<string, unknown> {
  return typeof value === 'object' && value !== null && !Array.isArray(value);
}

export async function enqueue(db: DrizzleClient, input: EnqueueLlmJobInput): Promise<string> {
  const rows = await db
    .insert(llmJobs)
    .values({ ...input, status: 'PENDING', retryCount: 0 })
    .returning({ id: llmJobs.id });
  const id = rows[0]?.id;
  if (!id) throw new Error('Failed to insert llm_job row');
  return id;
}

export async function enqueueVendorAdmissionReviewOnce(
  db: DrizzleClient,
  vendorId: string,
): Promise<string | null> {
  const existing = await db
    .select({ n: count() })
    .from(llmJobs)
    .where(
      and(
        eq(llmJobs.jobType, 'VENDOR_VIOLATION'),
        eq(llmJobs.targetType, 'VENDOR'),
        eq(llmJobs.targetId, vendorId),
        eq(llmJobs.status, 'PENDING'),
      ),
    );
  if ((existing[0]?.n ?? 0) > 0) return null;
  return enqueue(db, {
    jobType: 'VENDOR_VIOLATION',
    targetType: 'VENDOR',
    targetId: vendorId,
    inputPayload: { vendorId },
  });
}

export async function claim(db: DrizzleClient, jobId: string): Promise<boolean> {
  return (
    (
      await db
        .update(llmJobs)
        .set({ status: 'RUNNING', startedAt: new Date() })
        .where(and(eq(llmJobs.id, jobId), eq(llmJobs.status, 'PENDING')))
        .returning({ id: llmJobs.id })
    ).length > 0
  );
}

export async function fail(
  db: DrizzleClient,
  jobId: string,
  error: string,
  terminal = true,
): Promise<void> {
  await db
    .update(llmJobs)
    .set({ status: 'FAILED', lastError: error, ...(terminal ? { completedAt: new Date() } : {}) })
    .where(eq(llmJobs.id, jobId));
}

export async function complete(
  db: DrizzleClient,
  jobId: string,
  result: unknown,
  meta?: JobCompletionMeta,
): Promise<void> {
  await db
    .update(llmJobs)
    .set({
      status: 'COMPLETED',
      outputPayload: isRecord(result) ? result : {},
      completedAt: new Date(),
      ...(meta?.modelName !== undefined ? { modelName: meta.modelName } : {}),
      ...(meta?.promptTokens !== undefined ? { promptTokens: meta.promptTokens } : {}),
      ...(meta?.completionTokens !== undefined ? { completionTokens: meta.completionTokens } : {}),
      ...(meta?.totalTokens !== undefined ? { totalTokens: meta.totalTokens } : {}),
      ...(meta?.costUsd !== undefined ? { costUsd: meta.costUsd } : {}),
    })
    .where(eq(llmJobs.id, jobId));
}

export async function resetStaleRunningJobs(db: DrizzleClient): Promise<void> {
  await db
    .update(llmJobs)
    .set({ status: 'PENDING', startedAt: sql`NULL` })
    .where(
      and(eq(llmJobs.status, 'RUNNING'), lt(llmJobs.startedAt, sql`now() - interval '10 minutes'`)),
    );
}

export async function markLlmJobCompleted(db: DrizzleClient, jobId: string): Promise<void> {
  await db.update(llmJobs).set({ completedAt: new Date() }).where(eq(llmJobs.id, jobId));
}

export async function markLlmJobsFailedAndCompleted(
  db: DrizzleClient,
  jobId: string,
  lastError: string,
): Promise<void> {
  await db
    .update(llmJobs)
    .set({
      status: 'FAILED',
      lastError,
      completedAt: new Date(),
    })
    .where(eq(llmJobs.id, jobId));
}

export async function setE2eLlmJobFixtureState(
  db: DrizzleClient,
  runId: string,
  jobId: string,
  fixture: {
    status: 'FAILED' | 'COMPLETED';
    retryCount?: number;
    modelName?: string;
    promptTokens?: number;
    completionTokens?: number;
    totalTokens?: number;
    costUsd?: string;
  },
): Promise<void> {
  await db
    .update(llmJobs)
    .set({
      ...fixture,
      lastError: fixture.status === 'FAILED' ? 'E2E fixture failure' : null,
      completedAt: fixture.status === 'COMPLETED' ? new Date() : null,
      outputPayload: fixture.status === 'COMPLETED' ? { e2eFixture: true } : null,
    })
    .where(and(eq(llmJobs.id, jobId), eq(llmJobs.e2eRunId, runId)));
}