import { Global, Module } from '@nestjs/common';
import {
  createMemory,
  type Memory,
  createMemoryAdapter,
  type MemoryAdapter,
  type MemoryConfig,
} from '@mosaicstack/memory';
import type { Db } from '@mosaicstack/db';
import type { StorageAdapter } from '@mosaicstack/storage';
import type { MosaicConfig } from '@mosaicstack/config';
import { MOSAIC_CONFIG } from '../config/config.module.js';
import { DB, STORAGE_ADAPTER } from '../database/database.module.js';
import { MEMORY } from './memory.tokens.js';
import { MemoryController } from './memory.controller.js';
import { EmbeddingService } from './embedding.service.js';

export const MEMORY_ADAPTER = 'MEMORY_ADAPTER';

function buildMemoryConfig(config: MosaicConfig, storageAdapter: StorageAdapter): MemoryConfig {
  if (config.memory.type === 'keyword') {
    return { type: 'keyword', storage: storageAdapter };
  }
  return { type: config.memory.type };
}

@Global()
@Module({
  providers: [
    {
      provide: MEMORY,
      useFactory: (db: Db): Memory => createMemory(db),
      inject: [DB],
    },
    {
      provide: MEMORY_ADAPTER,
      useFactory: (config: MosaicConfig, storageAdapter: StorageAdapter): MemoryAdapter =>
        createMemoryAdapter(buildMemoryConfig(config, storageAdapter)),
      inject: [MOSAIC_CONFIG, STORAGE_ADAPTER],
    },
    EmbeddingService,
  ],
  controllers: [MemoryController],
  exports: [MEMORY, MEMORY_ADAPTER, EmbeddingService],
})
export class MemoryModule {}