import { Collection, Filter } from "mongodb";

import { getDb } from "../index";
import {
  MODELS_COLLECTION,
  MODEL_CONFIG_DEFAULTS_CONTEXT_WINDOW,
} from "./models.constants";
import { ModelDoc, ModelType } from "./models.types";

/** Raw typed handle to the `models` collection. */
export const getModelsCollection = (): Collection<ModelDoc> => {
  return getDb().collection<ModelDoc>(MODELS_COLLECTION);
};

/** Find models by an arbitrary filter. */
export const findModels = async (
  filter: Filter<ModelDoc> = {},
): Promise<ModelDoc[]> => {
  return await getModelsCollection().find(filter).toArray();
};

/**
 * Get a single model by its `modelId`, optionally scoped to a `type`.
 * @param modelId e.g. "deepgram-stt-phonecall"
 * @param type optional narrowing filter, since `modelId` isn't guaranteed unique across types
 */
export const getModelByModelId = async (
  modelId: string,
  type?: ModelType,
): Promise<ModelDoc | null> => {
  const filter = (type ? { modelId, type } : { modelId }) as Filter<ModelDoc>;
  return await getModelsCollection().findOne(filter);
};

/**
 * Token context window for a model (`configDefaults.contextWindow`).
 * `null` when the model is missing or the value is not a positive number.
 */
export const getModelContextWindow = async (
  modelId: string,
): Promise<number | null> => {
  const doc = await getModelByModelId(modelId);
  const value = doc?.configDefaults?.[MODEL_CONFIG_DEFAULTS_CONTEXT_WINDOW];
  if (typeof value !== "number" || !Number.isFinite(value) || value <= 0) {
    return null;
  }
  return Math.floor(value);
};

/**
 * Get all models for a given `provider`, optionally scoped to a `type`.
 */
export const getModelsByProvider = async (
  provider: string,
  type?: ModelType,
): Promise<ModelDoc[]> => {
  const filter = (type ? { provider, type } : { provider }) as Filter<ModelDoc>;
  return await findModels(filter);
};

/** Get all models of a given `type`. */
export const getModelsByType = async (type: ModelType): Promise<ModelDoc[]> => {
  return await findModels({ type } as Filter<ModelDoc>);
};

/** Get all models of a given `type` where `available` is true. */
export const getAvailableModelsByType = async (
  type: ModelType,
): Promise<ModelDoc[]> => {
  return await findModels({ type, available: true } as Filter<ModelDoc>);
};