Momento/memento-note/lib/ai/factory.ts

import { OpenAIProvider } from './providers/openai';
import { OllamaProvider } from './providers/ollama';
import { CustomOpenAIProvider } from './providers/custom-openai';
import { AIProvider } from './types';

type ProviderType = 'ollama' | 'openai' | 'custom' | 'deepseek' | 'openrouter' | 'mistral' | 'zai' | 'lmstudio';

// --- Provider defaults ---
const PROVIDER_DEFAULTS: Record<string, { baseUrl: string; model: string; embeddingModel: string }> = {
  deepseek: {
    baseUrl: 'https://api.deepseek.com/v1',
    model: 'deepseek-chat',
    embeddingModel: '',
  },
  openrouter: {
    baseUrl: 'https://openrouter.ai/api/v1',
    model: 'openai/gpt-4o-mini',
    embeddingModel: 'openai/text-embedding-3-small',
  },
  mistral: {
    baseUrl: 'https://api.mistral.ai/v1',
    model: 'mistral-small-latest',
    embeddingModel: 'mistral-embed',
  },
  zai: {
    baseUrl: 'https://api.zukijourney.com/v1',
    model: 'gpt-4o-mini',
    embeddingModel: 'text-embedding-3-small',
  },
  lmstudio: {
    baseUrl: 'http://localhost:1234/v1',
    model: '',
    embeddingModel: '',
  },
};

function createOllamaProvider(config: Record<string, string>, modelName: string, embeddingModelName: string, baseUrlOverride?: string): OllamaProvider {
  let baseUrl = baseUrlOverride || config?.OLLAMA_BASE_URL || process.env.OLLAMA_BASE_URL

  // Only use localhost as fallback for local development (not in Docker)
  if (!baseUrl && process.env.NODE_ENV !== 'production') {
    baseUrl = 'http://localhost:11434'
  }

  if (!baseUrl) {
    throw new Error('OLLAMA_BASE_URL is required when using Ollama provider')
  }

  // Ensure baseUrl doesn't end with /api, we'll add it in OllamaProvider
  if (baseUrl.endsWith('/api')) {
    baseUrl = baseUrl.slice(0, -4);
  }

  return new OllamaProvider(baseUrl, modelName, embeddingModelName);
}

function createOpenAIProvider(config: Record<string, string>, modelName: string, embeddingModelName: string): OpenAIProvider {
  const apiKey = config?.OPENAI_API_KEY || process.env.OPENAI_API_KEY || '';

  if (!apiKey) {
    throw new Error('OPENAI_API_KEY is required when using OpenAI provider');
  }

  return new OpenAIProvider(apiKey, modelName, embeddingModelName);
}

function createCustomOpenAIProvider(config: Record<string, string>, modelName: string, embeddingModelName: string): CustomOpenAIProvider {
  const apiKey = config?.CUSTOM_OPENAI_API_KEY || process.env.CUSTOM_OPENAI_API_KEY || '';
  const baseUrl = config?.CUSTOM_OPENAI_BASE_URL || process.env.CUSTOM_OPENAI_BASE_URL || '';

  if (!apiKey) {
    throw new Error('CUSTOM_OPENAI_API_KEY is required when using Custom OpenAI provider');
  }

  if (!baseUrl) {
    throw new Error('CUSTOM_OPENAI_BASE_URL is required when using Custom OpenAI provider');
  }

  return new CustomOpenAIProvider(apiKey, baseUrl, modelName, embeddingModelName);
}

function createDeepSeekProvider(config: Record<string, string>, modelName: string, embeddingModelName: string): CustomOpenAIProvider {
  const apiKey = config?.DEEPSEEK_API_KEY || process.env.DEEPSEEK_API_KEY || '';
  if (!apiKey) throw new Error('DEEPSEEK_API_KEY is required when using DeepSeek provider');
  const defaults = PROVIDER_DEFAULTS.deepseek;
  return new CustomOpenAIProvider(apiKey, defaults.baseUrl, modelName || defaults.model, embeddingModelName || defaults.embeddingModel);
}

function createOpenRouterProvider(config: Record<string, string>, modelName: string, embeddingModelName: string): CustomOpenAIProvider {
  const apiKey = config?.OPENROUTER_API_KEY || process.env.OPENROUTER_API_KEY || '';
  if (!apiKey) throw new Error('OPENROUTER_API_KEY is required when using OpenRouter provider');
  const defaults = PROVIDER_DEFAULTS.openrouter;
  return new CustomOpenAIProvider(apiKey, defaults.baseUrl, modelName || defaults.model, embeddingModelName || defaults.embeddingModel);
}

function createMistralProvider(config: Record<string, string>, modelName: string, embeddingModelName: string): CustomOpenAIProvider {
  const apiKey = config?.MISTRAL_API_KEY || process.env.MISTRAL_API_KEY || '';
  if (!apiKey) throw new Error('MISTRAL_API_KEY is required when using Mistral provider');
  const defaults = PROVIDER_DEFAULTS.mistral;
  return new CustomOpenAIProvider(apiKey, defaults.baseUrl, modelName || defaults.model, embeddingModelName || defaults.embeddingModel);
}

function createZAIProvider(config: Record<string, string>, modelName: string, embeddingModelName: string): CustomOpenAIProvider {
  const apiKey = config?.ZAI_API_KEY || process.env.ZAI_API_KEY || '';
  if (!apiKey) throw new Error('ZAI_API_KEY is required when using Z.AI provider');
  const defaults = PROVIDER_DEFAULTS.zai;
  return new CustomOpenAIProvider(apiKey, defaults.baseUrl, modelName || defaults.model, embeddingModelName || defaults.embeddingModel);
}

function createLMStudioProvider(config: Record<string, string>, modelName: string, embeddingModelName: string): CustomOpenAIProvider {
  const baseUrl = config?.LMSTUDIO_BASE_URL || process.env.LMSTUDIO_BASE_URL || PROVIDER_DEFAULTS.lmstudio.baseUrl;
  // LM Studio doesn't require an API key, but the CustomOpenAI provider needs one
  // Use a dummy key if not provided
  const apiKey = config?.LMSTUDIO_API_KEY || process.env.LMSTUDIO_API_KEY || 'lm-studio';
  return new CustomOpenAIProvider(apiKey, baseUrl, modelName, embeddingModelName);
}

function getProviderInstance(providerType: ProviderType, config: Record<string, string>, modelName: string, embeddingModelName: string, ollamaBaseUrl?: string): AIProvider {
  switch (providerType) {
    case 'ollama':
      return createOllamaProvider(config, modelName, embeddingModelName, ollamaBaseUrl);
    case 'openai':
      return createOpenAIProvider(config, modelName, embeddingModelName);
    case 'custom':
      return createCustomOpenAIProvider(config, modelName, embeddingModelName);
    case 'deepseek':
      return createDeepSeekProvider(config, modelName, embeddingModelName);
    case 'openrouter':
      return createOpenRouterProvider(config, modelName, embeddingModelName);
    case 'mistral':
      return createMistralProvider(config, modelName, embeddingModelName);
    case 'zai':
      return createZAIProvider(config, modelName, embeddingModelName);
    case 'lmstudio':
      return createLMStudioProvider(config, modelName, embeddingModelName);
    default:
      return createOllamaProvider(config, modelName, embeddingModelName, ollamaBaseUrl);
  }
}

// Resolve the effective provider type and config keys for a given provider
// Returns { providerType, apiKeyConfigKey, baseUrlConfigKey }
function getProviderConfigKeys(providerType: string): { apiKeyConfigKey: string; baseUrlConfigKey: string } {
  switch (providerType) {
    case 'deepseek':   return { apiKeyConfigKey: 'DEEPSEEK_API_KEY', baseUrlConfigKey: '' };
    case 'openrouter': return { apiKeyConfigKey: 'OPENROUTER_API_KEY', baseUrlConfigKey: '' };
    case 'mistral':    return { apiKeyConfigKey: 'MISTRAL_API_KEY', baseUrlConfigKey: '' };
    case 'zai':        return { apiKeyConfigKey: 'ZAI_API_KEY', baseUrlConfigKey: '' };
    case 'lmstudio':   return { apiKeyConfigKey: 'LMSTUDIO_API_KEY', baseUrlConfigKey: 'LMSTUDIO_BASE_URL' };
    case 'openai':     return { apiKeyConfigKey: 'OPENAI_API_KEY', baseUrlConfigKey: '' };
    case 'custom':     return { apiKeyConfigKey: 'CUSTOM_OPENAI_API_KEY', baseUrlConfigKey: 'CUSTOM_OPENAI_BASE_URL' };
    default:           return { apiKeyConfigKey: '', baseUrlConfigKey: 'OLLAMA_BASE_URL' };
  }
}

export function getTagsProvider(config?: Record<string, string>): AIProvider {
  const providerType = (
    config?.AI_PROVIDER_TAGS ||
    config?.AI_PROVIDER_EMBEDDING ||
    config?.AI_PROVIDER ||
    process.env.AI_PROVIDER_TAGS ||
    process.env.AI_PROVIDER_EMBEDDING ||
    process.env.AI_PROVIDER
  );

  if (!providerType) {
    console.error('[getTagsProvider] FATAL: No provider configured. Config received:', config);
    throw new Error(
      'AI_PROVIDER_TAGS is not configured. Please set it in the admin settings or environment variables. ' +
      'Options: ollama, openai, deepseek, openrouter, mistral, zai, lmstudio, custom'
    );
  }

  const provider = providerType.toLowerCase() as ProviderType;
  const modelName = config?.AI_MODEL_TAGS || process.env.AI_MODEL_TAGS || 'granite4:latest';
  const embeddingModelName = config?.AI_MODEL_EMBEDDING || process.env.AI_MODEL_EMBEDDING || 'embeddinggemma:latest';
  const ollamaBaseUrl = config?.OLLAMA_BASE_URL_TAGS || config?.OLLAMA_BASE_URL;

  return getProviderInstance(provider, config || {}, modelName, embeddingModelName, ollamaBaseUrl);
}

export function getEmbeddingsProvider(config?: Record<string, string>): AIProvider {
  const providerType = (
    config?.AI_PROVIDER_EMBEDDING ||
    config?.AI_PROVIDER_TAGS ||
    config?.AI_PROVIDER ||
    process.env.AI_PROVIDER_EMBEDDING ||
    process.env.AI_PROVIDER_TAGS ||
    process.env.AI_PROVIDER
  );

  if (!providerType) {
    console.error('[getEmbeddingsProvider] FATAL: No provider configured. Config received:', config);
    throw new Error(
      'AI_PROVIDER_EMBEDDING is not configured. Please set it in the admin settings or environment variables. ' +
      'Options: ollama, openai, deepseek, openrouter, mistral, zai, lmstudio, custom'
    );
  }

  const provider = providerType.toLowerCase() as ProviderType;
  const modelName = config?.AI_MODEL_TAGS || process.env.AI_MODEL_TAGS || 'granite4:latest';
  const embeddingModelName = config?.AI_MODEL_EMBEDDING || process.env.AI_MODEL_EMBEDDING || 'embeddinggemma:latest';
  const ollamaBaseUrl = config?.OLLAMA_BASE_URL_EMBEDDING || config?.OLLAMA_BASE_URL;

  return getProviderInstance(provider, config || {}, modelName, embeddingModelName, ollamaBaseUrl);
}

export function getAIProvider(config?: Record<string, string>): AIProvider {
  return getEmbeddingsProvider(config);
}

export function getChatProvider(config?: Record<string, string>): AIProvider {
  const providerType = (
    config?.AI_PROVIDER_CHAT ||
    config?.AI_PROVIDER_TAGS ||
    config?.AI_PROVIDER_EMBEDDING ||
    config?.AI_PROVIDER ||
    process.env.AI_PROVIDER_CHAT ||
    process.env.AI_PROVIDER_TAGS ||
    process.env.AI_PROVIDER_EMBEDDING ||
    process.env.AI_PROVIDER
  );

  if (!providerType) {
    console.error('[getChatProvider] FATAL: No provider configured. Config received:', config);
    throw new Error(
      'AI_PROVIDER_CHAT is not configured. Please set it in the admin settings or environment variables. ' +
      'Options: ollama, openai, deepseek, openrouter, mistral, zai, lmstudio, custom'
    );
  }

  const provider = providerType.toLowerCase() as ProviderType;
  const modelName = (
    config?.AI_MODEL_CHAT ||
    process.env.AI_MODEL_CHAT ||
    'granite4:latest'
  );
  const embeddingModelName = config?.AI_MODEL_EMBEDDING || process.env.AI_MODEL_EMBEDDING || 'embeddinggemma:latest';
  const ollamaBaseUrl = config?.OLLAMA_BASE_URL_CHAT || config?.OLLAMA_BASE_URL_TAGS || config?.OLLAMA_BASE_URL_EMBEDDING || config?.OLLAMA_BASE_URL;

  return getProviderInstance(provider, config || {}, modelName, embeddingModelName, ollamaBaseUrl);
}

// Export for use by admin settings form and deploy scripts
export { PROVIDER_DEFAULTS, getProviderConfigKeys };