dcrouter/ts/storage/classes.storagemanager.ts

import * as plugins from '../plugins.js';
import { logger } from '../logger.js';

// Promisify filesystem operations
const readFile = plugins.util.promisify(plugins.fs.readFile);
const writeFile = plugins.util.promisify(plugins.fs.writeFile);
const unlink = plugins.util.promisify(plugins.fs.unlink);
const rename = plugins.util.promisify(plugins.fs.rename);
const readdir = plugins.util.promisify(plugins.fs.readdir);

/**
 * Storage configuration interface
 */
export interface IStorageConfig {
  /** Filesystem path for storage */
  fsPath?: string;
  /** Custom read function */
  readFunction?: (key: string) => Promise<string>;
  /** Custom write function */
  writeFunction?: (key: string, value: string) => Promise<void>;
}

/**
 * Storage backend type
 */
export type StorageBackend = 'filesystem' | 'custom' | 'memory';

/**
 * Central storage manager for DcRouter
 * Provides unified key-value storage with multiple backend support
 */
export class StorageManager {
  private backend: StorageBackend;
  private memoryStore: Map<string, string> = new Map();
  private config: IStorageConfig;
  private fsBasePath?: string;
  
  constructor(config?: IStorageConfig) {
    this.config = config || {};
    
    // Check if both fsPath and custom functions are provided
    if (config?.fsPath && (config?.readFunction || config?.writeFunction)) {
      console.warn(
        '⚠️  WARNING: Both fsPath and custom read/write functions are configured.\n' +
        '   Using custom read/write functions. fsPath will be ignored.'
      );
    }
    
    // Determine backend based on configuration
    if (config?.readFunction && config?.writeFunction) {
      this.backend = 'custom';
    } else if (config?.fsPath) {
      // Set up internal read/write functions for filesystem
      this.backend = 'custom'; // Use custom backend with internal functions
      this.fsBasePath = plugins.path.resolve(config.fsPath);
      this.ensureDirectory(this.fsBasePath);
      
      // Set up internal filesystem read/write functions
      this.config.readFunction = async (key: string) => {
        return this.fsRead(key);
      };
      this.config.writeFunction = async (key: string, value: string) => {
        await this.fsWrite(key, value);
      };
    } else {
      this.backend = 'memory';
      this.showMemoryWarning();
    }
    
    logger.log('info', `StorageManager initialized with ${this.backend} backend`);
  }
  
  /**
   * Show warning when using memory backend
   */
  private showMemoryWarning(): void {
    console.warn(
      '⚠️  WARNING: StorageManager is using in-memory storage.\n' +
      '   Data will be lost when the process restarts.\n' +
      '   Configure storage.fsPath or storage functions for persistence.'
    );
  }
  
  /**
   * Ensure directory exists for filesystem backend
   */
  private async ensureDirectory(dirPath: string): Promise<void> {
    try {
      await plugins.smartfile.fs.ensureDir(dirPath);
    } catch (error) {
      logger.log('error', `Failed to create storage directory: ${error.message}`);
      throw error;
    }
  }
  
  /**
   * Validate and sanitize storage key
   */
  private validateKey(key: string): string {
    if (!key || typeof key !== 'string') {
      throw new Error('Storage key must be a non-empty string');
    }
    
    // Ensure key starts with /
    if (!key.startsWith('/')) {
      key = '/' + key;
    }
    
    // Remove any dangerous path elements
    key = key.replace(/\.\./g, '').replace(/\/+/g, '/');
    
    return key;
  }
  
  /**
   * Convert key to filesystem path
   */
  private keyToPath(key: string): string {
    if (!this.fsBasePath) {
      throw new Error('Filesystem base path not configured');
    }
    
    // Remove leading slash and convert to path
    const relativePath = key.substring(1);
    return plugins.path.join(this.fsBasePath, relativePath);
  }
  
  /**
   * Internal filesystem read function
   */
  private async fsRead(key: string): Promise<string> {
    const filePath = this.keyToPath(key);
    try {
      const content = await readFile(filePath, 'utf8');
      return content;
    } catch (error) {
      if (error.code === 'ENOENT') {
        return null;
      }
      throw error;
    }
  }
  
  /**
   * Internal filesystem write function
   */
  private async fsWrite(key: string, value: string): Promise<void> {
    const filePath = this.keyToPath(key);
    const dir = plugins.path.dirname(filePath);
    
    // Ensure directory exists
    await plugins.smartfile.fs.ensureDir(dir);
    
    // Write atomically with temp file
    const tempPath = `${filePath}.tmp`;
    await writeFile(tempPath, value, 'utf8');
    await rename(tempPath, filePath);
  }
  
  /**
   * Get value by key
   */
  async get(key: string): Promise<string | null> {
    key = this.validateKey(key);
    
    try {
      switch (this.backend) {
        
        case 'custom': {
          if (!this.config.readFunction) {
            throw new Error('Read function not configured');
          }
          try {
            return await this.config.readFunction(key);
          } catch (error) {
            // Assume null if read fails (key doesn't exist)
            return null;
          }
        }
        
        case 'memory': {
          return this.memoryStore.get(key) || null;
        }
        
        default:
          throw new Error(`Unknown backend: ${this.backend}`);
      }
    } catch (error) {
      logger.log('error', `Storage get error for key ${key}: ${error.message}`);
      throw error;
    }
  }
  
  /**
   * Set value by key
   */
  async set(key: string, value: string): Promise<void> {
    key = this.validateKey(key);
    
    if (typeof value !== 'string') {
      throw new Error('Storage value must be a string');
    }
    
    try {
      switch (this.backend) {
        case 'filesystem': {
          const filePath = this.keyToPath(key);
          const dirPath = plugins.path.dirname(filePath);
          
          // Ensure directory exists
          await plugins.smartfile.fs.ensureDir(dirPath);
          
          // Write atomically
          const tempPath = filePath + '.tmp';
          await writeFile(tempPath, value, 'utf8');
          await rename(tempPath, filePath);
          break;
        }
        
        case 'custom': {
          if (!this.config.writeFunction) {
            throw new Error('Write function not configured');
          }
          await this.config.writeFunction(key, value);
          break;
        }
        
        case 'memory': {
          this.memoryStore.set(key, value);
          break;
        }
        
        default:
          throw new Error(`Unknown backend: ${this.backend}`);
      }
    } catch (error) {
      logger.log('error', `Storage set error for key ${key}: ${error.message}`);
      throw error;
    }
  }
  
  /**
   * Delete value by key
   */
  async delete(key: string): Promise<void> {
    key = this.validateKey(key);
    
    try {
      switch (this.backend) {
        case 'filesystem': {
          const filePath = this.keyToPath(key);
          try {
            await unlink(filePath);
          } catch (error) {
            if (error.code !== 'ENOENT') {
              throw error;
            }
          }
          break;
        }
        
        case 'custom': {
          // Try to delete by setting empty value
          if (this.config.writeFunction) {
            await this.config.writeFunction(key, '');
          }
          break;
        }
        
        case 'memory': {
          this.memoryStore.delete(key);
          break;
        }
        
        default:
          throw new Error(`Unknown backend: ${this.backend}`);
      }
    } catch (error) {
      logger.log('error', `Storage delete error for key ${key}: ${error.message}`);
      throw error;
    }
  }
  
  /**
   * List keys by prefix
   */
  async list(prefix?: string): Promise<string[]> {
    prefix = prefix ? this.validateKey(prefix) : '/';
    
    try {
      switch (this.backend) {
        case 'custom': {
          // If we have fsBasePath, this is actually filesystem backend
          if (this.fsBasePath) {
            const basePath = this.keyToPath(prefix);
            const keys: string[] = [];
            
            const walkDir = async (dir: string, baseDir: string): Promise<void> => {
              try {
                const entries = await readdir(dir, { withFileTypes: true });
                
                for (const entry of entries) {
                  const fullPath = plugins.path.join(dir, entry.name);
                  
                  if (entry.isDirectory()) {
                    await walkDir(fullPath, baseDir);
                  } else if (entry.isFile()) {
                    // Convert path back to key
                    const relativePath = plugins.path.relative(this.fsBasePath!, fullPath);
                    const key = '/' + relativePath.replace(/\\/g, '/');
                    if (key.startsWith(prefix)) {
                      keys.push(key);
                    }
                  }
                }
              } catch (error) {
                if (error.code !== 'ENOENT') {
                  throw error;
                }
              }
            };
            
            await walkDir(basePath, basePath);
            return keys.sort();
          } else {
            // True custom backends need to implement their own listing
            logger.log('warn', 'List operation not supported for custom backend');
            return [];
          }
        }
        
        case 'memory': {
          const keys: string[] = [];
          for (const key of this.memoryStore.keys()) {
            if (key.startsWith(prefix)) {
              keys.push(key);
            }
          }
          return keys.sort();
        }
        
        default:
          throw new Error(`Unknown backend: ${this.backend}`);
      }
    } catch (error) {
      logger.log('error', `Storage list error for prefix ${prefix}: ${error.message}`);
      throw error;
    }
  }
  
  /**
   * Check if key exists
   */
  async exists(key: string): Promise<boolean> {
    key = this.validateKey(key);
    
    try {
      const value = await this.get(key);
      return value !== null;
    } catch (error) {
      return false;
    }
  }
  
  /**
   * Get storage backend type
   */
  getBackend(): StorageBackend {
    // If we're using custom backend with fsBasePath, report it as filesystem
    if (this.backend === 'custom' && this.fsBasePath) {
      return 'filesystem' as StorageBackend;
    }
    return this.backend;
  }
  
  /**
   * JSON helper: Get and parse JSON value
   */
  async getJSON<T = any>(key: string): Promise<T | null> {
    const value = await this.get(key);
    if (value === null) {
      return null;
    }
    
    try {
      return JSON.parse(value) as T;
    } catch (error) {
      logger.log('error', `Failed to parse JSON for key ${key}: ${error.message}`);
      throw error;
    }
  }
  
  /**
   * JSON helper: Set value as JSON
   */
  async setJSON(key: string, value: any): Promise<void> {
    const jsonString = JSON.stringify(value, null, 2);
    await this.set(key, jsonString);
  }
}