Codewalkers/apps/server/db/repositories/drizzle/account.ts

/**
 * Drizzle Account Repository Adapter
 *
 * Implements AccountRepository interface using Drizzle ORM.
 * Handles round-robin selection via lastUsedAt ordering
 * and automatic exhaustion expiry.
 */

import { eq, and, asc, lte } from 'drizzle-orm';
import { nanoid } from 'nanoid';
import type { DrizzleDatabase } from '../../index.js';
import { accounts, agents, type Account } from '../../schema.js';
import type { AccountRepository, CreateAccountData } from '../account-repository.js';

export class DrizzleAccountRepository implements AccountRepository {
  constructor(private db: DrizzleDatabase) {}

  async create(data: CreateAccountData): Promise<Account> {
    const id = nanoid();
    const now = new Date();

    const [created] = await this.db.insert(accounts).values({
      id,
      email: data.email,
      provider: data.provider ?? 'claude',
      configJson: data.configJson ?? null,
      credentials: data.credentials ?? null,
      isExhausted: false,
      exhaustedUntil: null,
      lastUsedAt: null,
      sortOrder: 0,
      createdAt: now,
      updatedAt: now,
    }).returning();

    return created;
  }

  async findById(id: string): Promise<Account | null> {
    const result = await this.db
      .select()
      .from(accounts)
      .where(eq(accounts.id, id))
      .limit(1);
    return result[0] ?? null;
  }

  async findByEmail(email: string): Promise<Account | null> {
    const result = await this.db
      .select()
      .from(accounts)
      .where(eq(accounts.email, email))
      .limit(1);
    return result[0] ?? null;
  }

  async findByProvider(provider: string): Promise<Account[]> {
    return this.db
      .select()
      .from(accounts)
      .where(eq(accounts.provider, provider));
  }

  async findNextAvailable(provider: string): Promise<Account | null> {
    await this.clearExpiredExhaustion();

    const result = await this.db
      .select()
      .from(accounts)
      .where(
        and(
          eq(accounts.provider, provider),
          eq(accounts.isExhausted, false),
        ),
      )
      .orderBy(asc(accounts.lastUsedAt))
      .limit(1);

    return result[0] ?? null;
  }

  async markExhausted(id: string, until: Date): Promise<Account> {
    const now = new Date();
    const [updated] = await this.db
      .update(accounts)
      .set({
        isExhausted: true,
        exhaustedUntil: until,
        updatedAt: now,
      })
      .where(eq(accounts.id, id))
      .returning();

    if (!updated) {
      throw new Error(`Account not found: ${id}`);
    }

    return updated;
  }

  async markAvailable(id: string): Promise<Account> {
    const now = new Date();
    const [updated] = await this.db
      .update(accounts)
      .set({
        isExhausted: false,
        exhaustedUntil: null,
        updatedAt: now,
      })
      .where(eq(accounts.id, id))
      .returning();

    if (!updated) {
      throw new Error(`Account not found: ${id}`);
    }

    return updated;
  }

  async updateLastUsed(id: string): Promise<Account> {
    const now = new Date();
    const [updated] = await this.db
      .update(accounts)
      .set({ lastUsedAt: now, updatedAt: now })
      .where(eq(accounts.id, id))
      .returning();

    if (!updated) {
      throw new Error(`Account not found: ${id}`);
    }

    return updated;
  }

  async clearExpiredExhaustion(): Promise<number> {
    const now = new Date();

    const cleared = await this.db
      .update(accounts)
      .set({
        isExhausted: false,
        exhaustedUntil: null,
        updatedAt: now,
      })
      .where(
        and(
          eq(accounts.isExhausted, true),
          lte(accounts.exhaustedUntil, now),
        ),
      )
      .returning({ id: accounts.id });

    return cleared.length;
  }

  async findAll(): Promise<Account[]> {
    return this.db.select().from(accounts);
  }

  async updateCredentials(id: string, credentials: string): Promise<Account> {
    const now = new Date();
    const [updated] = await this.db
      .update(accounts)
      .set({ credentials, updatedAt: now })
      .where(eq(accounts.id, id))
      .returning();

    if (!updated) {
      throw new Error(`Account not found: ${id}`);
    }

    return updated;
  }

  async updateAccountAuth(id: string, configJson: string, credentials: string): Promise<Account> {
    const now = new Date();
    const [updated] = await this.db
      .update(accounts)
      .set({ configJson, credentials, updatedAt: now })
      .where(eq(accounts.id, id))
      .returning();

    if (!updated) {
      throw new Error(`Account not found: ${id}`);
    }

    return updated;
  }

  async delete(id: string): Promise<void> {
    // Manually nullify agent FK — the migration lacks ON DELETE SET NULL
    await this.db
      .update(agents)
      .set({ accountId: null })
      .where(eq(agents.accountId, id));

    const [deleted] = await this.db.delete(accounts).where(eq(accounts.id, id)).returning();

    if (!deleted) {
      throw new Error(`Account not found: ${id}`);
    }
  }
}