import { and, eq, lte, or, sql } from 'drizzle-orm';
import type { DrizzleClient } from '@/server/db/client.js';
import { deals, outbox, translationJobs } from '@/server/db/schema.js';

export async function findTranslationDeal(db: DrizzleClient, dealId: string) {
  return db
    .select({ sourceLanguage: deals.sourceLanguage })
    .from(deals)
    .where(eq(deals.id, dealId))
    .limit(1);
}

export async function insertTranslationJob(
  db: DrizzleClient,
  input: {
    dealId: string;
    targetLocale: string;
    forceFresh: boolean;
    scheduledAt: Date | null;
  },
) {
  return db
    .insert(translationJobs)
    .values({ ...input, status: 'PENDING' })
    .onConflictDoNothing()
    .returning({ id: translationJobs.id });
}

export async function insertTranslationEnqueuedOutbox(db: DrizzleClient, jobId: string) {
  return db
    .insert(outbox)
    .values({
      aggregateType: 'translation_job',
      aggregateId: jobId,
      eventType: 'TRANSLATION_JOB_ENQUEUED',
      payload: { jobId },
    })
    .returning({ id: outbox.id });
}

export async function failTranslationJob(
  db: DrizzleClient,
  jobId: string,
  error: string,
): Promise<void> {
  await db
    .update(translationJobs)
    .set({ status: 'FAILED', lastError: error, finishedAt: sql`now()` })
    .where(eq(translationJobs.id, jobId));
}

export async function claimTranslationJob(db: DrizzleClient, jobId: string) {
  const [row] = await db
    .update(translationJobs)
    .set({ status: 'RUNNING', attempt: sql`${translationJobs.attempt} + 1`, startedAt: new Date() })
    .where(
      and(
        eq(translationJobs.id, jobId),
        or(eq(translationJobs.status, 'PENDING'), eq(translationJobs.status, 'PENDING_BUDGET')),
        or(
          sql`${translationJobs.scheduledAt} IS NULL`,
          lte(translationJobs.scheduledAt, new Date()),
        ),
      ),
    )
    .returning({
      id: translationJobs.id,
      deal_id: translationJobs.dealId,
      target_locale: translationJobs.targetLocale,
      attempt: translationJobs.attempt,
      force_fresh: translationJobs.forceFresh,
    });
  return row ?? null;
}

export async function completeTranslationJob(
  db: DrizzleClient,
  jobId: string,
  result: { costUsd: number | null; inputTokens: number | null; outputTokens: number | null },
): Promise<void> {
  await db
    .update(translationJobs)
    .set({
      status: 'DONE',
      costUsd: result.costUsd === null ? null : result.costUsd.toFixed(6),
      inputTokens: result.inputTokens,
      outputTokens: result.outputTokens,
      finishedAt: new Date(),
    })
    .where(eq(translationJobs.id, jobId));
}

export async function markTranslationBudgetExhausted(
  db: DrizzleClient,
  jobId: string,
): Promise<void> {
  await db
    .update(translationJobs)
    .set({ status: 'PENDING_BUDGET', lastError: 'Budget exhausted', finishedAt: new Date() })
    .where(eq(translationJobs.id, jobId));
}

export async function rescheduleTranslationJob(
  db: DrizzleClient,
  jobId: string,
  error: string,
  delaySeconds: number,
): Promise<void> {
  await db
    .update(translationJobs)
    .set({
      status: 'PENDING',
      lastError: error,
      startedAt: null,
      scheduledAt: sql`now() + (${delaySeconds} || ' seconds')::interval`,
    })
    .where(eq(translationJobs.id, jobId));
}

export async function getTranslationJobDetails(db: DrizzleClient, jobId: string) {
  return db
    .select({
      dealId: translationJobs.dealId,
      targetLocale: translationJobs.targetLocale,
      attempt: translationJobs.attempt,
    })
    .from(translationJobs)
    .where(eq(translationJobs.id, jobId));
}

export async function insertTranslationFailureOutbox(
  db: DrizzleClient,
  input: {
    jobId: string;
    dealId: string | null;
    targetLocale: string | null;
    error: string;
    attempt: number | null;
  },
) {
  return db
    .insert(outbox)
    .values({
      aggregateType: 'translation_job',
      aggregateId: input.jobId,
      eventType: 'TRANSLATION_JOB_FAILED',
      payload: input,
    })
    .returning({ id: outbox.id });
}