import { agentTasks, type AgentTask, type Database } from '@pig/db';
import { and, asc, desc, eq, isNull, lt, lte, or, sql } from 'drizzle-orm';

export type FailureDisposition = 'retry' | 'failed' | 'lease_lost';

export class AgentTaskQueue {
  constructor(
    private readonly db: Database,
    private readonly workerId: string,
    private readonly leaseSeconds: number,
  ) {}

  async claimNext(now = new Date()): Promise<AgentTask | null> {
    return this.db.transaction(async (tx) => {
      const [candidate] = await tx
        .select()
        .from(agentTasks)
        .where(
          and(
            isNull(agentTasks.finishedAt),
            lte(agentTasks.dueAt, now),
            or(isNull(agentTasks.leasedUntil), lt(agentTasks.leasedUntil, now)),
            sql`${agentTasks.attempts} < ${agentTasks.maxAttempts}`,
          ),
        )
        .orderBy(desc(agentTasks.priority), asc(agentTasks.dueAt), asc(agentTasks.createdAt))
        .limit(1)
        .for('update', { skipLocked: true });

      if (!candidate) return null;

      const [claimed] = await tx
        .update(agentTasks)
        .set({
          leasedBy: this.workerId,
          leasedUntil: leaseUntil(now, this.leaseSeconds),
          startedAt: candidate.startedAt ?? now,
          attempts: sql`${agentTasks.attempts} + 1`,
          error: null,
        })
        .where(eq(agentTasks.id, candidate.id))
        .returning();
      return claimed ?? null;
    });
  }

  async renew(taskId: string, now = new Date()): Promise<boolean> {
    const rows = await this.db
      .update(agentTasks)
      .set({ leasedUntil: leaseUntil(now, this.leaseSeconds) })
      .where(
        and(
          eq(agentTasks.id, taskId),
          eq(agentTasks.leasedBy, this.workerId),
          isNull(agentTasks.finishedAt),
        ),
      )
      .returning({ id: agentTasks.id });
    return rows.length === 1;
  }

  async succeed(taskId: string, now = new Date()): Promise<boolean> {
    const rows = await this.db
      .update(agentTasks)
      .set({
        finishedAt: now,
        outcome: 'succeeded',
        leasedBy: null,
        leasedUntil: null,
        error: null,
      })
      .where(
        and(
          eq(agentTasks.id, taskId),
          eq(agentTasks.leasedBy, this.workerId),
          isNull(agentTasks.finishedAt),
        ),
      )
      .returning({ id: agentTasks.id });
    return rows.length === 1;
  }

  async fail(task: AgentTask, error: string, now = new Date()): Promise<FailureDisposition> {
    const retry = task.attempts < task.maxAttempts;
    const rows = await this.db
      .update(agentTasks)
      .set(
        retry
          ? {
              dueAt: new Date(now.getTime() + retryBackoffMs(task.attempts)),
              leasedBy: null,
              leasedUntil: null,
              error,
            }
          : {
              finishedAt: now,
              outcome: 'failed',
              leasedBy: null,
              leasedUntil: null,
              error,
            },
      )
      .where(
        and(
          eq(agentTasks.id, task.id),
          eq(agentTasks.leasedBy, this.workerId),
          isNull(agentTasks.finishedAt),
        ),
      )
      .returning({ id: agentTasks.id });

    if (rows.length === 0) return 'lease_lost';
    return retry ? 'retry' : 'failed';
  }
}

export function retryBackoffMs(attempts: number): number {
  return Math.min(60 * 60_000, 60_000 * 2 ** Math.max(0, attempts - 1));
}

function leaseUntil(now: Date, leaseSeconds: number): Date {
  return new Date(now.getTime() + leaseSeconds * 1_000);
}
