// src/models/ai-provider.ts
// Copyright (C) 2026 Robert Colbert <rob.colbert@openplatform.us>
// All Rights Reserved

import { Schema, model } from "mongoose";
import { nanoid } from "nanoid";

import {
  IAiModel,
  IAiModelCapabilities,
  IAiModelSettings,
  IAiProvider,
} from "@gadget/api";

export const AiModelCapabilitiesSchema = new Schema<IAiModelCapabilities>(
  {
    canCallTools: { type: Boolean, default: false },
    hasVision: { type: Boolean, default: false },
    hasEmbedding: { type: Boolean, default: false },
    hasThinking: { type: Boolean, default: false },
    isInstructTuned: { type: Boolean, default: false },
  },
  { _id: false },
);

export const AiModelSettingsSchema = new Schema<IAiModelSettings>(
  {
    temperature: { type: Number },
    topP: { type: Number },
    topK: { type: Number },
    numCtx: { type: Number },
    numPredict: { type: Number },
    maxCompletionTokens: { type: Number },
  },
  { _id: false },
);

export const AiModelSchema = new Schema<IAiModel>(
  {
    id: { type: String, required: true },
    name: { type: String, required: true },
    parameterCount: { type: Number },
    parameterLabel: { type: String },
    contextWindow: { type: Number },
    capabilities: {
      type: AiModelCapabilitiesSchema,
      default: () => ({
        canCallTools: false,
        hasVision: false,
        hasEmbedding: false,
        hasThinking: false,
        isInstructTuned: false,
      }),
    },
    settings: {
      type: AiModelSettingsSchema,
      default: undefined,
    },
  },
  { _id: false },
);

export const AiProviderSchema = new Schema<IAiProvider>({
  _id: { type: String, default: () => nanoid() },
  name: { type: String, required: true },
  apiType: { type: String, enum: ["ollama", "openai"], required: true },
  baseUrl: { type: String, required: true },
  apiKey: { type: String, required: false, default: "" },
  enabled: { type: Boolean, default: true, required: true },
  models: { type: [AiModelSchema], default: [], required: true },
  lastModelRefresh: { type: Date, default: Date.now },
});

AiProviderSchema.index({ name: 1 }, { unique: true });

export const AiProvider = model<IAiProvider>("AiProvider", AiProviderSchema);
export default AiProvider;
