LightningLotto/telegram_bot/src/services/state.ts

import Redis from 'ioredis';
import config from '../config';
import { logger } from './logger';
import {
  TelegramUser,
  UserState,
  AwaitingPaymentData,
} from '../types';

const STATE_PREFIX = 'tg_user:';
const PURCHASE_PREFIX = 'tg_purchase:';
const USER_PURCHASES_PREFIX = 'tg_user_purchases:';
const STATE_TTL = 60 * 60 * 24 * 30; // 30 days

class StateManager {
  private redis: Redis | null = null;
  private memoryStore: Map<string, string> = new Map();
  private useRedis: boolean = false;

  async init(): Promise<void> {
    if (config.redis.url) {
      try {
        this.redis = new Redis(config.redis.url);
        
        this.redis.on('error', (error) => {
          logger.error('Redis connection error', { error: error.message });
        });

        this.redis.on('connect', () => {
          logger.info('Connected to Redis');
        });

        // Test connection
        await this.redis.ping();
        this.useRedis = true;
        logger.info('State manager initialized with Redis');
      } catch (error) {
        logger.warn('Failed to connect to Redis, falling back to in-memory store', {
          error: (error as Error).message,
        });
        this.redis = null;
        this.useRedis = false;
      }
    } else {
      logger.info('State manager initialized with in-memory store');
      this.useRedis = false;
    }
  }

  private async get(key: string): Promise<string | null> {
    if (this.useRedis && this.redis) {
      return await this.redis.get(key);
    }
    return this.memoryStore.get(key) || null;
  }

  private async set(key: string, value: string, ttl?: number): Promise<void> {
    if (this.useRedis && this.redis) {
      if (ttl) {
        await this.redis.setex(key, ttl, value);
      } else {
        await this.redis.set(key, value);
      }
    } else {
      this.memoryStore.set(key, value);
    }
  }

  private async del(key: string): Promise<void> {
    if (this.useRedis && this.redis) {
      await this.redis.del(key);
    } else {
      this.memoryStore.delete(key);
    }
  }

  private async lpush(key: string, value: string): Promise<void> {
    if (this.useRedis && this.redis) {
      await this.redis.lpush(key, value);
      await this.redis.ltrim(key, 0, 99); // Keep last 100 purchases
    } else {
      const existing = this.memoryStore.get(key);
      const list = existing ? JSON.parse(existing) : [];
      list.unshift(value);
      if (list.length > 100) list.pop();
      this.memoryStore.set(key, JSON.stringify(list));
    }
  }

  private async lrange(key: string, start: number, stop: number): Promise<string[]> {
    if (this.useRedis && this.redis) {
      return await this.redis.lrange(key, start, stop);
    }
    const existing = this.memoryStore.get(key);
    if (!existing) return [];
    const list = JSON.parse(existing);
    return list.slice(start, stop + 1);
  }

  /**
   * Get or create user
   */
  async getUser(telegramId: number): Promise<TelegramUser | null> {
    const key = `${STATE_PREFIX}${telegramId}`;
    const data = await this.get(key);
    if (!data) return null;
    
    try {
      const user = JSON.parse(data);
      return {
        ...user,
        createdAt: new Date(user.createdAt),
        updatedAt: new Date(user.updatedAt),
      };
    } catch (error) {
      logger.error('Failed to parse user data', { telegramId, error });
      return null;
    }
  }

  /**
   * Create or update user
   */
  async saveUser(user: TelegramUser): Promise<void> {
    const key = `${STATE_PREFIX}${user.telegramId}`;
    user.updatedAt = new Date();
    await this.set(key, JSON.stringify(user), STATE_TTL);
    logger.debug('User saved', { telegramId: user.telegramId, state: user.state });
  }

  /**
   * Create new user
   */
  async createUser(
    telegramId: number,
    username?: string,
    firstName?: string,
    lastName?: string
  ): Promise<TelegramUser> {
    const user: TelegramUser = {
      telegramId,
      username,
      firstName,
      lastName,
      state: 'awaiting_lightning_address',
      createdAt: new Date(),
      updatedAt: new Date(),
    };
    await this.saveUser(user);
    logger.info('New user created', { telegramId, username });
    return user;
  }

  /**
   * Update user state
   */
  async updateUserState(
    telegramId: number,
    state: UserState,
    stateData?: Record<string, any>
  ): Promise<void> {
    const user = await this.getUser(telegramId);
    if (!user) {
      logger.warn('Attempted to update state for non-existent user', { telegramId });
      return;
    }
    user.state = state;
    user.stateData = stateData;
    await this.saveUser(user);
  }

  /**
   * Update user's lightning address
   */
  async updateLightningAddress(
    telegramId: number,
    lightningAddress: string
  ): Promise<void> {
    const user = await this.getUser(telegramId);
    if (!user) {
      logger.warn('Attempted to update address for non-existent user', { telegramId });
      return;
    }
    user.lightningAddress = lightningAddress;
    user.state = 'idle';
    user.stateData = undefined;
    await this.saveUser(user);
  }

  /**
   * Store a ticket purchase for a user
   */
  async storePurchase(
    telegramId: number,
    purchaseId: string,
    data: AwaitingPaymentData
  ): Promise<void> {
    // Store purchase data
    const purchaseKey = `${PURCHASE_PREFIX}${purchaseId}`;
    await this.set(purchaseKey, JSON.stringify({
      telegramId,
      ...data,
      createdAt: new Date().toISOString(),
    }), STATE_TTL);

    // Add to user's purchase list
    const userPurchasesKey = `${USER_PURCHASES_PREFIX}${telegramId}`;
    await this.lpush(userPurchasesKey, purchaseId);
  }

  /**
   * Get purchase data
   */
  async getPurchase(purchaseId: string): Promise<(AwaitingPaymentData & { telegramId: number }) | null> {
    const key = `${PURCHASE_PREFIX}${purchaseId}`;
    const data = await this.get(key);
    if (!data) return null;
    
    try {
      return JSON.parse(data);
    } catch (error) {
      logger.error('Failed to parse purchase data', { purchaseId, error });
      return null;
    }
  }

  /**
   * Get user's recent purchase IDs
   */
  async getUserPurchaseIds(
    telegramId: number,
    limit: number = 10
  ): Promise<string[]> {
    const key = `${USER_PURCHASES_PREFIX}${telegramId}`;
    return await this.lrange(key, 0, limit - 1);
  }

  /**
   * Clear user state data (keeping lightning address)
   */
  async clearUserStateData(telegramId: number): Promise<void> {
    const user = await this.getUser(telegramId);
    if (!user) return;
    user.state = 'idle';
    user.stateData = undefined;
    await this.saveUser(user);
  }

  /**
   * Shutdown
   */
  async close(): Promise<void> {
    if (this.redis) {
      await this.redis.quit();
      logger.info('Redis connection closed');
    }
  }
}

export const stateManager = new StateManager();
export default stateManager;