Alfred/src/server/db/models.ts

import type Database from "bun:sqlite";

import { desc, eq, like, or, sql } from "drizzle-orm";

import type { CreateModelRequest, Model, ModelCapability, UpdateModelRequest } from "../../shared/api";

import { paginateQuery, wrap } from "./connection";
import { models, providers } from "./schema";

export function createModel(
  raw: Database,
  request: CreateModelRequest,
): { error: string; status: number } | { model: Model } {
  const db = wrap(raw);

  const provider = db.select().from(providers).where(eq(providers.id, request.providerId)).get();
  if (!provider) return { error: "供应商不存在", status: 400 };

  const name = request.name.trim();
  if (!name) return { error: "模型名称不能为空", status: 400 };

  const modelId = request.modelId.trim();
  if (!modelId) return { error: "模型 ID 不能为空", status: 400 };

  const capabilities = request.capabilities;
  if (!capabilities || capabilities.length === 0) {
    return { error: "至少选择一个能力标签", status: 400 };
  }

  const id = crypto.randomUUID();
  const now = new Date().toISOString();

  try {
    db.insert(models)
      .values({
        capabilities: JSON.stringify(capabilities),
        contextLength: request.contextLength ?? null,
        createdAt: now,
        id,
        maxOutputTokens: request.maxOutputTokens ?? null,
        modelId,
        name,
        providerId: request.providerId,
        updatedAt: now,
      })
      .run();
  } catch (e: unknown) {
    const msg = e instanceof Error ? e.message : String(e);
    if (msg.includes("UNIQUE constraint")) {
      return { error: "该供应商下模型 ID 已存在", status: 409 };
    }
    throw e;
  }

  const row = db.select().from(models).where(eq(models.id, id)).get();
  return { model: toModel(row!) };
}

export function deleteModel(raw: Database, id: string): { error: string; status: number } | { success: true } {
  const db = wrap(raw);
  const existing = db.select().from(models).where(eq(models.id, id)).get();
  if (!existing) return { error: "模型不存在", status: 404 };

  db.delete(models).where(eq(models.id, id)).run();
  return { success: true };
}

export function getModel(raw: Database, id: string): { error: string; status: number } | { model: Model } {
  const db = wrap(raw);
  const row = db.select().from(models).where(eq(models.id, id)).get();

  if (!row) return { error: "模型不存在", status: 404 };
  return { model: toModel(row) };
}

export function getModelsByProviderId(raw: Database, providerId: string): number {
  const db = wrap(raw);
  const result = db
    .select({ count: sql<number>`count(*)` })
    .from(models)
    .where(eq(models.providerId, providerId))
    .get();
  return Number(result?.count ?? 0);
}

export function listModels(
  raw: Database,
  options: { keyword?: string; page: number; pageSize: number; providerId?: string },
): { items: Model[]; page: number; pageSize: number; total: number } {
  const conditions = [];

  if (options.providerId) {
    conditions.push(eq(models.providerId, options.providerId));
  }

  if (options.keyword) {
    const pattern = `%${options.keyword}%`;
    conditions.push(or(like(models.name, pattern), like(models.modelId, pattern))!);
  }

  return paginateQuery(raw, models, {
    conditions,
    mapRow: toModel,
    orderBy: () => desc(models.createdAt),
    page: options.page,
    pageSize: options.pageSize,
  });
}

export function updateModel(
  raw: Database,
  id: string,
  request: UpdateModelRequest,
): { error: string; status: number } | { model: Model } {
  const db = wrap(raw);
  const existing = db.select().from(models).where(eq(models.id, id)).get();
  if (!existing) return { error: "模型不存在", status: 404 };

  const updates: Partial<typeof models.$inferInsert> = {
    updatedAt: new Date().toISOString(),
  };

  const name = request.name?.trim();
  if (name === "") return { error: "模型名称不能为空", status: 400 };
  if (name !== undefined && name !== existing.name) {
    updates.name = name;
  }

  const modelId = request.modelId?.trim();
  if (modelId === "") return { error: "模型 ID 不能为空", status: 400 };
  if (modelId !== undefined) {
    updates.modelId = modelId;
  }

  if (request.providerId !== undefined) {
    const provider = db.select().from(providers).where(eq(providers.id, request.providerId)).get();
    if (!provider) return { error: "供应商不存在", status: 400 };
    updates.providerId = request.providerId;
  }

  if (request.capabilities !== undefined) {
    if (request.capabilities.length === 0) {
      return { error: "至少选择一个能力标签", status: 400 };
    }
    updates.capabilities = JSON.stringify(request.capabilities);
  }

  if (request.contextLength !== undefined) {
    updates.contextLength = request.contextLength;
  }

  if (request.maxOutputTokens !== undefined) {
    updates.maxOutputTokens = request.maxOutputTokens;
  }

  if (Object.keys(updates).length === 1 && updates.updatedAt) {
    return { model: toModel(existing) };
  }

  try {
    db.update(models).set(updates).where(eq(models.id, id)).run();
  } catch (e: unknown) {
    const msg = e instanceof Error ? e.message : String(e);
    if (msg.includes("UNIQUE constraint")) {
      return { error: "该供应商下模型 ID 已存在", status: 409 };
    }
    throw e;
  }

  const updated = db.select().from(models).where(eq(models.id, id)).get();
  return { model: toModel(updated!) };
}

function toModel(row: typeof models.$inferSelect): Model {
  return {
    capabilities: JSON.parse(row.capabilities) as ModelCapability[],
    contextLength: row.contextLength,
    createdAt: row.createdAt,
    id: row.id,
    maxOutputTokens: row.maxOutputTokens,
    modelId: row.modelId,
    name: row.name,
    providerId: row.providerId,
    updatedAt: row.updatedAt,
  };
}