platform/plugins/offscale/db/queries.ts

import { eq, and, desc } from 'drizzle-orm';
import { db } from 'officerdb/db';
import { headscaleServers } from './schema';
import { encryptSecret, decryptSecret } from 'officerdb/crypto';

// Headscale server registry access for the officer-headscale sidecar. Callers deal in PLAINTEXT —
// encryption to/from at-rest ciphertext happens here, so the sidecar's route handlers never touch crypto.
// See ../crypto.ts and ../schema/headscale.ts.
//
// Two return types on purpose:
//   HeadscaleServer            — safe to serialize to the browser. Has NO api key field at all.
//   HeadscaleServerCredentials — url + decrypted key, for the sidecar's own upstream calls. Never returned
//                                by a route handler.
// The `serverCols` projection is what enforces that: `select()` without it would leak the ciphertext column
// into every list response the moment someone forgot to strip it.

export type HeadscaleServer = {
  id: number;
  name: string;
  url: string;
  version: string | null;
  sshHost: string | null;
  isActive: boolean;
  lastSeenAt: Date | null;
  createdAt: Date;
};

export type HeadscaleServerCredentials = { id: number; name: string; url: string; apiKey: string };

const serverCols = {
  id: headscaleServers.id,
  name: headscaleServers.name,
  url: headscaleServers.url,
  version: headscaleServers.version,
  sshHost: headscaleServers.sshHost,
  isActive: headscaleServers.isActive,
  lastSeenAt: headscaleServers.lastSeenAt,
  createdAt: headscaleServers.createdAt,
};

/** Every server the owner has registered, active first then newest. Never includes the API key. */
export async function listHeadscaleServers(userId: number): Promise<HeadscaleServer[]> {
  return db
    .select(serverCols)
    .from(headscaleServers)
    .where(eq(headscaleServers.userId, userId))
    .orderBy(desc(headscaleServers.isActive), desc(headscaleServers.createdAt));
}

/** The currently selected server with its key decrypted, or null when none is registered/active. */
export async function getActiveHeadscaleCredentials(userId: number): Promise<HeadscaleServerCredentials | null> {
  const [row] = await db
    .select()
    .from(headscaleServers)
    .where(and(eq(headscaleServers.userId, userId), eq(headscaleServers.isActive, true)));
  if (!row) return null;
  return { id: row.id, name: row.name, url: row.url, apiKey: decryptSecret('headscale', row.apiKey) };
}

/** One server's credentials by id — for probing a specific server rather than the active one. */
export async function getHeadscaleCredentials(userId: number, id: number): Promise<HeadscaleServerCredentials | null> {
  const [row] = await db
    .select()
    .from(headscaleServers)
    .where(and(eq(headscaleServers.userId, userId), eq(headscaleServers.id, id)));
  if (!row) return null;
  return { id: row.id, name: row.name, url: row.url, apiKey: decryptSecret('headscale', row.apiKey) };
}

type CreateHeadscaleServerParams = {
  userId: number;
  name: string;
  url: string;
  apiKey: string;
  version: string | null;
  /** Optional SSH target for the console. Null when the owner hasn't set one. */
  sshHost: string | null;
  /** Make it the active server. True for the first registration, so the UI is never left with none selected. */
  activate: boolean;
};

/** Register a server. The key is encrypted before write; the returned row carries no key. */
export async function createHeadscaleServer(params: CreateHeadscaleServerParams): Promise<HeadscaleServer> {
  const { userId, name, url, apiKey, version, sshHost, activate } = params;
  return db.transaction(async (tx) => {
    if (activate) {
      await tx
        .update(headscaleServers)
        .set({ isActive: false, updatedAt: new Date() })
        .where(and(eq(headscaleServers.userId, userId), eq(headscaleServers.isActive, true)));
    }
    const [row] = await tx
      .insert(headscaleServers)
      .values({
        userId,
        name,
        url,
        apiKey: encryptSecret('headscale', apiKey),
        version,
        sshHost,
        isActive: activate,
        lastSeenAt: version ? new Date() : null,
      })
      .returning(serverCols);
    return row!;
  });
}

// `sshHost: null` clears the console target; omitting the field leaves it alone. The two must stay
// distinguishable, which is why this is `string | null` and not `string`.
type UpdateHeadscaleServerParams = { name?: string; url?: string; apiKey?: string; sshHost?: string | null };

/** Edit a registration. Omitted fields are left alone; a supplied key is re-encrypted. */
export async function updateHeadscaleServer(
  userId: number,
  id: number,
  params: UpdateHeadscaleServerParams,
): Promise<HeadscaleServer | null> {
  const set: Record<string, unknown> = { updatedAt: new Date() };
  if (params.name !== undefined) set.name = params.name;
  if (params.url !== undefined) set.url = params.url;
  if (params.apiKey !== undefined) set.apiKey = encryptSecret('headscale', params.apiKey);
  if (params.sshHost !== undefined) set.sshHost = params.sshHost;

  const [row] = await db
    .update(headscaleServers)
    .set(set)
    .where(and(eq(headscaleServers.userId, userId), eq(headscaleServers.id, id)))
    .returning(serverCols);
  return row ?? null;
}

/** Select a server. Clearing the others first keeps the one-active partial index satisfied. */
export async function setActiveHeadscaleServer(userId: number, id: number): Promise<HeadscaleServer | null> {
  return db.transaction(async (tx) => {
    await tx
      .update(headscaleServers)
      .set({ isActive: false, updatedAt: new Date() })
      .where(and(eq(headscaleServers.userId, userId), eq(headscaleServers.isActive, true)));
    const [row] = await tx
      .update(headscaleServers)
      .set({ isActive: true, updatedAt: new Date() })
      .where(and(eq(headscaleServers.userId, userId), eq(headscaleServers.id, id)))
      .returning(serverCols);
    return row ?? null;
  });
}

/**
 * Drop a registration. If it was the active one, the newest survivor is promoted — otherwise deleting the
 * active server would leave the UI with servers registered but none selected, which reads as "not
 * configured" and is a confusing place to land.
 */
export async function deleteHeadscaleServer(userId: number, id: number): Promise<boolean> {
  return db.transaction(async (tx) => {
    const [deleted] = await tx
      .delete(headscaleServers)
      .where(and(eq(headscaleServers.userId, userId), eq(headscaleServers.id, id)))
      .returning({ id: headscaleServers.id, wasActive: headscaleServers.isActive });
    if (!deleted) return false;

    if (deleted.wasActive) {
      const [next] = await tx
        .select({ id: headscaleServers.id })
        .from(headscaleServers)
        .where(eq(headscaleServers.userId, userId))
        .orderBy(desc(headscaleServers.createdAt))
        .limit(1);
      if (next) {
        await tx
          .update(headscaleServers)
          .set({ isActive: true, updatedAt: new Date() })
          .where(eq(headscaleServers.id, next.id));
      }
    }
    return true;
  });
}

/** Record a successful reachability probe: the version observed and when we last reached the server. */
export async function recordHeadscaleProbe(userId: number, id: number, version: string | null): Promise<void> {
  await db
    .update(headscaleServers)
    .set({ version, lastSeenAt: new Date(), updatedAt: new Date() })
    .where(and(eq(headscaleServers.userId, userId), eq(headscaleServers.id, id)));
}