stock-bot/apps/stock/data-ingestion/src/handlers/qm/shared/session-manager.ts

/**
 * QM Session Manager - Centralized session state management
 */

import type { CacheProvider } from '@stock-bot/cache';
import type { Logger } from '@stock-bot/types';
import { QM_SESSION_IDS, SESSION_CONFIG } from './config';
import type { CachedSession, QMSession } from './types';

export class QMSessionManager {
  private static instance: QMSessionManager | null = null;
  private sessionCache: Record<string, QMSession[]> = {};
  private isInitialized = false;
  private cacheProvider: CacheProvider | null = null;
  private logger: Logger | null = null;

  private constructor() {
    // Initialize session cache with known session IDs
    Object.values(QM_SESSION_IDS).forEach(sessionId => {
      this.sessionCache[sessionId] = [];
    });
  }

  static getInstance(): QMSessionManager {
    if (!QMSessionManager.instance) {
      QMSessionManager.instance = new QMSessionManager();
    }
    return QMSessionManager.instance;
  }
  
  /**
   * Reset the singleton instance (for testing only)
   */
  static resetInstance(): void {
    if (QMSessionManager.instance?.logger) {
      QMSessionManager.instance.logger.warn('Resetting QMSessionManager instance - this should only be used for testing');
    }
    QMSessionManager.instance = null;
  }

  /**
   * Set the cache provider for persistence
   */
  setCacheProvider(cache: CacheProvider): void {
    this.cacheProvider = cache;
  }

  /**
   * Set the logger
   */
  setLogger(logger: Logger): void {
    this.logger = logger;
    this.logger.trace('Logger set for QMSessionManager');
  }

  /**
   * Initialize with cache provider and logger
   */
  initialize(cache?: CacheProvider, logger?: Logger): void {
    if (cache) {
      this.setCacheProvider(cache);
    }
    if (logger) {
      this.setLogger(logger);
    }
  }

  /**
   * Get a random session for the given session ID
   */
  async getSession(sessionId: string): Promise<QMSession | null> {
    // Always load fresh data from cache
    await this.loadFromCache();
    
    const sessions = this.sessionCache[sessionId];
    if (!sessions || sessions.length === 0) {
      return null;
    }

    // Filter out sessions with excessive failures
    const validSessions = sessions.filter(
      session => session.failedCalls <= SESSION_CONFIG.MAX_FAILED_CALLS
    );
    if (validSessions.length === 0) {
      return null;
    }

    return validSessions[Math.floor(Math.random() * validSessions.length)];
  }

  /**
   * Get a specific session by UUID
   */
  getSessionByUuid(sessionId: string, uuid: string): QMSession | null {
    const sessions = this.sessionCache[sessionId] || [];
    return sessions.find(s => s.uuid === uuid) || null;
  }

  /**
   * Add a session to the cache
   */
  async addSession(sessionId: string, session: QMSession): Promise<void> {
    // Load latest from cache first to avoid overwriting other service's changes
    await this.loadFromCache();
    
    if (!this.sessionCache[sessionId]) {
      this.sessionCache[sessionId] = [];
    }
    this.sessionCache[sessionId].push(session);
    
    this.logger?.debug(`Added session ${session.uuid} to ${sessionId}`, { 
      sessionId, 
      uuid: session.uuid 
    });
    
    // Sync to cache immediately
    await this.syncToCache();
  }

  /**
   * Get all sessions for a session ID
   */
  getSessions(sessionId: string): QMSession[] {
    return this.sessionCache[sessionId] || [];
  }

  /**
   * Get session count for all session IDs
   */
  getSessionCount(): number {
    return Object.values(this.sessionCache).reduce((total, sessions) => total + sessions.length, 0);
  }

  /**
   * Clean up failed sessions
   */
  async cleanupFailedSessions(): Promise<number> {
    // Always load latest from cache first
    await this.loadFromCache();
    
    let removedCount = 0;

    Object.keys(this.sessionCache).forEach(sessionId => {
      const initialCount = this.sessionCache[sessionId].length;
      this.sessionCache[sessionId] = this.sessionCache[sessionId].filter(
        session => session.failedCalls <= SESSION_CONFIG.MAX_FAILED_CALLS
      );
      removedCount += initialCount - this.sessionCache[sessionId].length;
    });

    // Sync back to cache after cleanup
    await this.syncToCache();
    
    return removedCount;
  }

  /**
   * Check if more sessions are needed for a session ID
   */
  async needsMoreSessions(sessionId: string): Promise<boolean> {
    // Always load fresh data from cache
    await this.loadFromCache();
    
    const sessions = this.sessionCache[sessionId] || [];
    const validSessions = sessions.filter(
      session => session.failedCalls <= SESSION_CONFIG.MAX_FAILED_CALLS
    );
    return validSessions.length < SESSION_CONFIG.MAX_SESSIONS;
  }

  /**
   * Check if session ID is at capacity
   */
  isAtCapacity(sessionId: string): boolean {
    const sessions = this.sessionCache[sessionId] || [];
    return sessions.length >= SESSION_CONFIG.MAX_SESSIONS;
  }

  /**
   * Get session cache statistics
   */
  getStats() {
    const stats: Record<string, { total: number; valid: number; failed: number }> = {};

    Object.entries(this.sessionCache).forEach(([sessionId, sessions]) => {
      const validSessions = sessions.filter(
        session => session.failedCalls <= SESSION_CONFIG.MAX_FAILED_CALLS
      );
      const failedSessions = sessions.filter(
        session => session.failedCalls > SESSION_CONFIG.MAX_FAILED_CALLS
      );

      stats[sessionId] = {
        total: sessions.length,
        valid: validSessions.length,
        failed: failedSessions.length,
      };
    });

    return stats;
  }

  /**
   * Mark manager as initialized (deprecated - we always load from cache now)
   */
  setInitialized(initialized: boolean = true): void {
    this.isInitialized = initialized;
  }

  /**
   * Check if manager is initialized (deprecated - we always load from cache now)
   */
  getInitialized(): boolean {
    return this.isInitialized;
  }

  /**
   * Load sessions from cache
   */
  async loadFromCache(): Promise<void> {
    if (!this.cacheProvider) {
      this.logger?.warn('No cache provider available for loading sessions');
      return;
    }

    try {
      this.logger?.trace('Loading sessions from cache...');
      
      // Load sessions for each session type
      for (const [sessionType, sessionId] of Object.entries(QM_SESSION_IDS)) {
        const listKey = `qm:sessions:${sessionType.toLowerCase()}:list`;
        const sessionIds = await this.cacheProvider.get<string[]>(listKey);

        this.logger?.trace(`Loading ${sessionType} sessions`, { sessionIds });

        if (sessionIds && Array.isArray(sessionIds)) {
          const sessions: QMSession[] = [];

          for (const id of sessionIds) {
            const sessionKey = `qm:sessions:${sessionType.toLowerCase()}:${id}`;
            const cachedSession = await this.cacheProvider.get<CachedSession>(sessionKey);

            if (cachedSession) {
              const session = {
                uuid: cachedSession.uuid,
                proxy: cachedSession.proxy,
                headers: cachedSession.headers,
                successfulCalls: cachedSession.successfulCalls || 0,
                failedCalls: cachedSession.failedCalls || 0,
                lastUsed: new Date(cachedSession.lastUsed),
                createdAt: cachedSession.createdAt ? new Date(cachedSession.createdAt) : new Date(),
              };
              
              this.logger?.trace(`Loaded session ${id}`, { 
                uuid: session.uuid, 
                successfulCalls: session.successfulCalls, 
                failedCalls: session.failedCalls 
              });
              
              sessions.push(session);
            }
          }

          this.sessionCache[sessionId] = sessions;
          this.logger?.debug(`Loaded ${sessions.length} sessions for ${sessionType}`);
        } else {
          this.logger?.trace(`No sessions found for ${sessionType}`);
        }
      }
    } catch (error) {
      this.logger?.error('Failed to load sessions from cache', { error });
    }
  }

  /**
   * Sync sessions to cache
   */
  async syncToCache(): Promise<void> {
    if (!this.cacheProvider) {
      return;
    }

    try {
      for (const [sessionType, sessionId] of Object.entries(QM_SESSION_IDS)) {
        const sessions = this.sessionCache[sessionId] || [];
        const sessionIds: string[] = [];

        // Clear old sessions first
        const oldListKey = `qm:sessions:${sessionType.toLowerCase()}:list`;
        const oldSessionIds = await this.cacheProvider.get<string[]>(oldListKey);
        if (oldSessionIds && Array.isArray(oldSessionIds)) {
          // Delete old session entries
          for (const oldId of oldSessionIds) {
            const oldSessionKey = `qm:sessions:${sessionType.toLowerCase()}:${oldId}`;
            await this.cacheProvider.del(oldSessionKey);
          }
        }

        // Store each session with stable IDs
        for (let i = 0; i < sessions.length; i++) {
          const session = sessions[i];
          const id = `${sessionType.toLowerCase()}_${i}`;
          const sessionKey = `qm:sessions:${sessionType.toLowerCase()}:${id}`;

          const cachedSession: CachedSession = {
            ...session,
            lastUsed: session.lastUsed instanceof Date ? session.lastUsed.toISOString() : session.lastUsed,
            createdAt: session.createdAt instanceof Date ? session.createdAt.toISOString() : session.createdAt,
            id,
            sessionType,
          } as any;

          this.logger?.trace(`Saving session ${id}`, { 
            uuid: cachedSession.uuid, 
            successfulCalls: cachedSession.successfulCalls, 
            failedCalls: cachedSession.failedCalls 
          });

          await this.cacheProvider.set(sessionKey, cachedSession, 86400); // 24 hour TTL
          sessionIds.push(id);
        }

        // Store the list of session IDs
        await this.cacheProvider.set(oldListKey, sessionIds, 86400);
      }

      // Store stats
      const statsKey = 'qm:sessions:stats';
      await this.cacheProvider.set(statsKey, this.getStats(), 3600);
      
      this.logger?.trace('Session sync to cache completed');
    } catch (error) {
      this.logger?.error('Failed to sync sessions to cache', { error });
    }
  }

  /**
   * Increment failed calls for a session
   */
  async incrementFailedCalls(sessionId: string, sessionUuid: string): Promise<void> {
    // Load latest from cache first
    await this.loadFromCache();
    
    // Find session by UUID
    const sessions = this.sessionCache[sessionId] || [];
    const session = sessions.find(s => s.uuid === sessionUuid);
    
    if (session) {
      session.failedCalls++;
      session.lastUsed = new Date();
      
      this.logger?.debug(`Incremented failed calls for session`, {
        sessionUuid,
        failedCalls: session.failedCalls,
        sessionId
      });
      
      // Sync to cache after update
      await this.syncToCache();
    } else {
      this.logger?.warn(`Session not found`, { sessionUuid, sessionId });
    }
  }

  /**
   * Increment successful calls for a session
   */
  async incrementSuccessfulCalls(sessionId: string, sessionUuid: string): Promise<void> {
    // Load latest from cache first
    await this.loadFromCache();
    
    // Find session by UUID
    const sessions = this.sessionCache[sessionId] || [];
    const session = sessions.find(s => s.uuid === sessionUuid);
    
    if (session) {
      session.successfulCalls++;
      session.lastUsed = new Date();
      
      this.logger?.debug(`Incremented successful calls for session`, {
        sessionUuid,
        successfulCalls: session.successfulCalls,
        sessionId
      });
      
      // Sync to cache after update
      await this.syncToCache();
    } else {
      this.logger?.warn(`Session not found`, { sessionUuid, sessionId });
    }
  }
}