/**
 * v5 InferenceTask repository — durable queue for inference work.
 *
 * Workers (mcplocal SSE sessions) drain rows whose `poolName` matches the
 * pool keys they own when they bind. The state machine is:
 *
 *     pending  ──claim──>  claimed  ──first chunk──>  running
 *                                 \                       │
 *                                  \                      ▼
 *                                   ──complete/fail──> completed | error
 *
 * Plus orthogonal terminal transitions:
 *
 *     pending|claimed|running ──cancel──> cancelled
 *     claimed|running ──worker disconnect──> pending  (revert + re-queue)
 *
 * The DB is the source of truth. mcpd's in-flight HTTP handlers wait via
 * an in-process EventEmitter (see TaskService) — single-instance for now;
 * pg LISTEN/NOTIFY is the v6 swap when scaling horizontally.
 */
import type {
  PrismaClient,
  InferenceTask,
  InferenceTaskStatus,
} from '@prisma/client';
// `Prisma` is imported as a value because we reference Prisma.JsonNull at
// runtime (the sentinel for explicit JSON null on a nullable column).
import { Prisma } from '@prisma/client';

export interface CreateInferenceTaskInput {
  poolName: string;
  llmName: string;
  model: string;
  tier?: string | null;
  requestBody: Record<string, unknown>;
  streaming: boolean;
  ownerId: string;
  agentId?: string | null;
}

export interface IInferenceTaskRepository {
  create(data: CreateInferenceTaskInput): Promise<InferenceTask>;
  findById(id: string): Promise<InferenceTask | null>;
  /** Pending rows for one or more pool keys, oldest first (FIFO). */
  findPendingForPools(poolNames: string[], limit?: number): Promise<InferenceTask[]>;
  /** Tasks held by a worker session — used by unbindSession to revert on disconnect. */
  findHeldBy(claimedBy: string): Promise<InferenceTask[]>;
  /** List for the async API; filters are AND-combined. */
  list(filter: {
    ownerId?: string;
    status?: InferenceTaskStatus | InferenceTaskStatus[];
    poolName?: string;
    agentId?: string;
    limit?: number;
  }): Promise<InferenceTask[]>;

  // ── State transitions ──
  // Each transition is conditional ("compare-and-swap" via updateMany +
  // affected-row count) so two workers racing to claim the same row both
  // succeed at the DB level — one gets affected=1 and a populated row,
  // the other gets affected=0 and we fall through to the next candidate.

  /** pending → claimed; returns the updated row, or null if the row was no longer pending. */
  tryClaim(id: string, claimedBy: string, claimedAt: Date): Promise<InferenceTask | null>;
  /** claimed → running; idempotent — returns the row even if already running. */
  markRunning(id: string, at: Date): Promise<InferenceTask | null>;
  /** {claimed,running} → completed; only allowed from a non-terminal state. */
  markCompleted(id: string, responseBody: Record<string, unknown> | null, at: Date): Promise<InferenceTask | null>;
  /** any non-terminal → error. Records `errorMessage`. */
  markError(id: string, errorMessage: string, at: Date): Promise<InferenceTask | null>;
  /** any non-terminal → cancelled. */
  markCancelled(id: string, at: Date): Promise<InferenceTask | null>;
  /** {claimed,running} → pending; clears `claimedBy`. Used on worker disconnect. */
  revertToPending(id: string): Promise<InferenceTask | null>;

  // ── GC ──

  /** Pending rows older than the cutoff — the GC sweep flips these to error. */
  findStalePending(cutoff: Date): Promise<InferenceTask[]>;
  /** Completed/error/cancelled rows older than the cutoff — GC deletes them. */
  findExpiredTerminal(cutoff: Date): Promise<InferenceTask[]>;
  /** Bulk delete by id — used by the GC sweep after collecting expired ids. */
  deleteMany(ids: string[]): Promise<number>;
}

const NON_TERMINAL: InferenceTaskStatus[] = ['pending', 'claimed', 'running'];
const HELD: InferenceTaskStatus[] = ['claimed', 'running'];

export class InferenceTaskRepository implements IInferenceTaskRepository {
  constructor(private readonly prisma: PrismaClient) {}

  async create(data: CreateInferenceTaskInput): Promise<InferenceTask> {
    return this.prisma.inferenceTask.create({
      data: {
        poolName: data.poolName,
        llmName: data.llmName,
        model: data.model,
        tier: data.tier ?? null,
        requestBody: data.requestBody as Prisma.InputJsonValue,
        streaming: data.streaming,
        ownerId: data.ownerId,
        agentId: data.agentId ?? null,
      },
    });
  }

  async findById(id: string): Promise<InferenceTask | null> {
    return this.prisma.inferenceTask.findUnique({ where: { id } });
  }

  async findPendingForPools(poolNames: string[], limit?: number): Promise<InferenceTask[]> {
    if (poolNames.length === 0) return [];
    return this.prisma.inferenceTask.findMany({
      where: { status: 'pending', poolName: { in: poolNames } },
      orderBy: { createdAt: 'asc' },
      ...(limit !== undefined ? { take: limit } : {}),
    });
  }

  async findHeldBy(claimedBy: string): Promise<InferenceTask[]> {
    return this.prisma.inferenceTask.findMany({
      where: { claimedBy, status: { in: HELD } },
    });
  }

  async list(filter: {
    ownerId?: string;
    status?: InferenceTaskStatus | InferenceTaskStatus[];
    poolName?: string;
    agentId?: string;
    limit?: number;
  }): Promise<InferenceTask[]> {
    const where: Prisma.InferenceTaskWhereInput = {};
    if (filter.ownerId !== undefined) where.ownerId = filter.ownerId;
    if (filter.status !== undefined) {
      where.status = Array.isArray(filter.status) ? { in: filter.status } : filter.status;
    }
    if (filter.poolName !== undefined) where.poolName = filter.poolName;
    if (filter.agentId !== undefined) where.agentId = filter.agentId;
    return this.prisma.inferenceTask.findMany({
      where,
      orderBy: { createdAt: 'desc' },
      ...(filter.limit !== undefined ? { take: filter.limit } : {}),
    });
  }

  async tryClaim(id: string, claimedBy: string, claimedAt: Date): Promise<InferenceTask | null> {
    // Compare-and-swap on status='pending'. Two workers racing both run
    // this UPDATE; whichever the DB serializes first sees affected=1 and
    // gets the row, the loser sees affected=0 and falls through.
    const result = await this.prisma.inferenceTask.updateMany({
      where: { id, status: 'pending' },
      data: { status: 'claimed', claimedBy, claimedAt },
    });
    if (result.count === 0) return null;
    return this.findById(id);
  }

  async markRunning(id: string, at: Date): Promise<InferenceTask | null> {
    const result = await this.prisma.inferenceTask.updateMany({
      where: { id, status: { in: ['claimed', 'running'] } },
      data: { status: 'running', streamStartedAt: at },
    });
    if (result.count === 0) return null;
    return this.findById(id);
  }

  async markCompleted(id: string, responseBody: Record<string, unknown> | null, at: Date): Promise<InferenceTask | null> {
    // Prisma's nullable JSON column writes need an explicit
    // `Prisma.JsonNull` sentinel for null; passing literal null fails
    // typecheck because the column also accepts a plain JSON value.
    const bodyForUpdate: Prisma.InputJsonValue | typeof Prisma.JsonNull =
      responseBody === null ? Prisma.JsonNull : (responseBody as Prisma.InputJsonValue);
    const result = await this.prisma.inferenceTask.updateMany({
      where: { id, status: { in: NON_TERMINAL } },
      data: {
        status: 'completed',
        responseBody: bodyForUpdate,
        completedAt: at,
      },
    });
    if (result.count === 0) return null;
    return this.findById(id);
  }

  async markError(id: string, errorMessage: string, at: Date): Promise<InferenceTask | null> {
    const result = await this.prisma.inferenceTask.updateMany({
      where: { id, status: { in: NON_TERMINAL } },
      data: { status: 'error', errorMessage, completedAt: at },
    });
    if (result.count === 0) return null;
    return this.findById(id);
  }

  async markCancelled(id: string, at: Date): Promise<InferenceTask | null> {
    const result = await this.prisma.inferenceTask.updateMany({
      where: { id, status: { in: NON_TERMINAL } },
      data: { status: 'cancelled', completedAt: at },
    });
    if (result.count === 0) return null;
    return this.findById(id);
  }

  async revertToPending(id: string): Promise<InferenceTask | null> {
    const result = await this.prisma.inferenceTask.updateMany({
      where: { id, status: { in: HELD } },
      data: { status: 'pending', claimedBy: null, claimedAt: null, streamStartedAt: null },
    });
    if (result.count === 0) return null;
    return this.findById(id);
  }

  async findStalePending(cutoff: Date): Promise<InferenceTask[]> {
    return this.prisma.inferenceTask.findMany({
      where: { status: 'pending', createdAt: { lt: cutoff } },
    });
  }

  async findExpiredTerminal(cutoff: Date): Promise<InferenceTask[]> {
    return this.prisma.inferenceTask.findMany({
      where: {
        status: { in: ['completed', 'error', 'cancelled'] },
        completedAt: { lt: cutoff },
      },
    });
  }

  async deleteMany(ids: string[]): Promise<number> {
    if (ids.length === 0) return 0;
    const result = await this.prisma.inferenceTask.deleteMany({
      where: { id: { in: ids } },
    });
    return result.count;
  }
}
