import { and, eq, inArray } from 'drizzle-orm';
import { z } from 'zod';
import type { DrizzleClient } from '@/server/db/client';
import { llmJobs, llmProviders } from '@/server/db/schema';
import { enqueueLlmJob } from '@/server/ai/llm';
import { runOwnedE2eLlmJob } from '@/server/cron/process-llm-jobs';
import * as providerQueries from '@/server/db/queries/llm-providers';
import { getModelCache, upsertModelCache } from '@/server/db/queries/llm-model-cache';
import { setE2eLlmJobFixtureState } from '@/server/db/queries/llm-jobs';
import type { FactoryOperationHandler, FactoryStore } from './core';

const createJobInput = z.object({
  key: z.uuid(),
  jobType: z.enum(['DEAL_MODERATION', 'TRANSLATION']),
  targetId: z.uuid(),
  inputPayload: z
    .record(z.string(), z.unknown())
    .refine(
      (value) => new TextEncoder().encode(JSON.stringify(value)).byteLength <= 8_192,
      'LLM input payload exceeds 8192 bytes',
    ),
  chain: z
    .array(
      z.object({
        llmProviderId: z.uuid(),
        model: z.string().regex(/^TEST(?:-[A-Z]+)?$/),
      }),
    )
    .min(1)
    .max(3),
  fixture: z
    .object({
      status: z.enum(['FAILED', 'COMPLETED']),
      retryCount: z.number().int().min(0).max(10).optional(),
      modelName: z.string().min(1).max(120).optional(),
      promptTokens: z.number().int().nonnegative().optional(),
      completionTokens: z.number().int().nonnegative().optional(),
      totalTokens: z.number().int().nonnegative().optional(),
      costUsd: z
        .string()
        .regex(/^\d+\.\d{6}$/)
        .optional(),
    })
    .optional(),
});
const processJobInput = z.object({ jobId: z.uuid() });
const createProviderInput = z.object({
  key: z.uuid(),
  name: z.string().min(1).max(120),
  models: z
    .array(z.object({ id: z.string().min(1).max(120), label: z.string().min(1).max(160) }))
    .min(1)
    .max(50),
});

export function createLlmFactoryHandlers(input: {
  db: DrizzleClient;
  store: FactoryStore;
  databaseUrl: string;
  piiKey: string;
}): Record<string, FactoryOperationHandler> {
  const { db, store, databaseUrl, piiKey } = input;
  const assertMockChain = async (
    runId: string,
    chain: Array<{ llmProviderId: string; model: string }>,
  ) => {
    const providerIds = [...new Set(chain.map((entry) => entry.llmProviderId))];
    const ownership = await Promise.all(
      providerIds.map((providerId) => store.findOwned(runId, 'llm-provider', providerId)),
    );
    if (ownership.some((owned) => !owned)) {
      throw new Error('LLM provider is not owned by factory run');
    }
    const providers = await db
      .select({ id: llmProviders.id, type: llmProviders.type, isActive: llmProviders.isActive })
      .from(llmProviders)
      .where(inArray(llmProviders.id, providerIds))
      .for('update');
    if (
      providers.length !== providerIds.length ||
      providers.some((provider) => provider.type !== 'mock' || !provider.isActive)
    ) {
      throw new Error('E2E LLM jobs require active mock providers');
    }
  };
  return {
    createLlmProvider: {
      dependencyOrder: 35,
      reserve: (raw) => ({ kind: 'llm-provider', id: createProviderInput.parse(raw).key }),
      async execute({ runId, input: raw }) {
        const parsed = createProviderInput.parse(raw);
        const provider = await providerQueries.insertLlmProvider(db, {
          id: parsed.key,
          slug: `e2e-${runId}-${parsed.key}`,
          name: parsed.name,
          type: 'mock',
          baseUrl: null,
          apiKey: 'TEST',
          piiKey,
        });
        await upsertModelCache(db, provider.id, parsed.models);
        return {
          entity: { kind: 'llm-provider', id: provider.id },
          result: { providerId: provider.id, name: provider.name, models: parsed.models, runId },
        };
      },
      async read({ runId, entity }) {
        if (!(await store.findOwned(runId, 'llm-provider', entity.entityId))) {
          throw new Error('LLM provider is not owned by factory run');
        }
        const provider = await providerQueries.getLlmProvider(db, entity.entityId);
        if (!provider) throw new Error('Owned LLM provider missing');
        return {
          ...providerQueries.toPublicLlmProvider(provider),
          cache: await getModelCache(db, provider.id),
        };
      },
      async cleanup({ runId, entity }) {
        if (!(await store.findOwned(runId, 'llm-provider', entity.entityId))) {
          throw new Error('LLM provider is not owned by factory run');
        }
        await providerQueries.deleteLlmProvider(db, entity.entityId);
      },
    },
    createLlmJob: {
      dependencyOrder: 40,
      reserve: (raw) => ({ kind: 'llm-job', id: createJobInput.parse(raw).key }),
      async execute({ runId, input: raw }) {
        const parsed = createJobInput.parse(raw);
        if (!(await store.findOwned(runId, 'deal', parsed.targetId))) {
          throw new Error('LLM target deal is not owned by factory run');
        }
        await assertMockChain(runId, parsed.chain);
        const jobId = await enqueueLlmJob(db, {
          id: parsed.key,
          jobType: parsed.jobType,
          targetId: parsed.targetId,
          targetType: 'DEAL',
          inputPayload: parsed.inputPayload,
          e2eRunId: runId,
          queueChainOverride: parsed.chain,
        });
        if (parsed.fixture) {
          await setE2eLlmJobFixtureState(db, runId, jobId, parsed.fixture);
        }
        return { entity: { kind: 'llm-job', id: jobId }, result: { jobId, runId } };
      },
      async read({ runId, entity }) {
        const [job] = await db
          .select()
          .from(llmJobs)
          .where(and(eq(llmJobs.id, entity.entityId), eq(llmJobs.e2eRunId, runId)))
          .limit(1)
          .for('update');
        if (!job) throw new Error('Owned LLM job missing');
        return job;
      },
      async cleanup({ runId, entity }) {
        const [job] = await db
          .select({ status: llmJobs.status })
          .from(llmJobs)
          .where(and(eq(llmJobs.id, entity.entityId), eq(llmJobs.e2eRunId, runId)))
          .limit(1);
        if (job && !['COMPLETED', 'FAILED', 'TIMED_OUT'].includes(job.status)) {
          throw new Error('Cannot clean up non-terminal LLM job');
        }
        await db
          .delete(llmJobs)
          .where(and(eq(llmJobs.id, entity.entityId), eq(llmJobs.e2eRunId, runId)));
      },
    },
    processLlmJob: {
      dependencyOrder: 0,
      async execute({ runId, input: raw }) {
        const { jobId } = processJobInput.parse(raw);
        if (!(await store.findOwned(runId, 'llm-job', jobId))) {
          throw new Error('LLM job is not owned by factory run');
        }
        const [job] = await db
          .select({ chain: llmJobs.queueChainOverride, status: llmJobs.status })
          .from(llmJobs)
          .where(and(eq(llmJobs.id, jobId), eq(llmJobs.e2eRunId, runId)))
          .limit(1);
        if (!job?.chain) throw new Error('Owned LLM job missing chain override');
        if (job.status !== 'PENDING') throw new Error('Owned LLM job is not pending');
        await assertMockChain(runId, job.chain);
        const result = await runOwnedE2eLlmJob(
          { DATABASE_URL: databaseUrl, PII_KEY: piiKey, db },
          runId,
          jobId,
        );
        return { result };
      },
      async read() {
        throw new Error('Process command has no entity');
      },
      async cleanup() {},
    },
  };
}