/**
 * AI Streaming Server Extension
 * Server-Sent Events (SSE) and WebSocket support for AI chat streaming
 * 
 * Features:
 * - SSE streaming for AI responses
 * - WebSocket for bi-directional chat
 * - Message queue management
 * - Rate limiting for AI APIs
 * - Response caching
 * - Conversation context
 * - Token-based streaming
 */

import { Request, Response, NextFunction } from 'express';
import { EventEmitter } from 'events';

export interface AIStreamConfig {
  /** SSE endpoint path */
  ssePath?: string;
  /** WebSocket path */
  wsPath?: string;
  /** Rate limit (requests per minute) */
  rateLimit?: number;
  /** Cache TTL (seconds) */
  cacheTTL?: number;
  /** Max concurrent streams */
  maxConcurrentStreams?: number;
  /** Enable debug logging */
  debug?: boolean;
}

export interface StreamMessage {
  id: string;
  type: 'token' | 'message' | 'error' | 'done';
  data: string | any;
  timestamp: number;
  conversationId?: string;
}

export interface ConversationContext {
  id: string;
  messages: Array<{ role: string; content: string }>;
  createdAt: number;
  lastActivity: number;
  tokenCount: number;
}

export class AIStreamingServer extends EventEmitter {
  private config: AIStreamConfig;
  private clients: Map<string, Response>;
  private conversations: Map<string, ConversationContext>;
  private rateLimitMap: Map<string, number[]>;
  private cache: Map<string, { data: any; expiry: number }>;
  private activeStreams: number = 0;

  constructor(config: AIStreamConfig = {}) {
    super();
    this.config = {
      ssePath: '/stream',
      wsPath: '/ws',
      rateLimit: 60, // 60 requests per minute
      cacheTTL: 300, // 5 minutes
      maxConcurrentStreams: 100,
      debug: false,
      ...config,
    };

    this.clients = new Map();
    this.conversations = new Map();
    this.rateLimitMap = new Map();
    this.cache = new Map();
  }

  /**
   * SSE Middleware - Stream AI responses
   */
  sseMiddleware(): (req: Request, res: Response, next: NextFunction) => void {
    return (req: Request, res: Response, next: NextFunction) => {
      if (req.path !== this.config.ssePath) {
        return next();
      }

      // Set SSE headers
      res.setHeader('Content-Type', 'text/event-stream');
      res.setHeader('Cache-Control', 'no-cache');
      res.setHeader('Connection', 'keep-alive');
      res.setHeader('X-Accel-Buffering', 'no'); // Disable nginx buffering

      const clientId = req.query.clientId as string || this.generateId();
      this.clients.set(clientId, res);
      this.activeStreams++;

      if (this.config.debug) {
        console.log(`[SSE] Client connected: ${clientId}, Active: ${this.activeStreams}`);
      }

      // Send connection event
      this.send(clientId, {
        id: this.generateId(),
        type: 'message',
        data: { event: 'connected', clientId },
        timestamp: Date.now(),
      });

      // Handle disconnect
      req.on('close', () => {
        this.clients.delete(clientId);
        this.activeStreams--;
        if (this.config.debug) {
          console.log(`[SSE] Client disconnected: ${clientId}, Active: ${this.activeStreams}`);
        }
      });

      // Handle errors
      req.on('error', (err) => {
        console.error(`[SSE] Error for client ${clientId}:`, err);
        this.clients.delete(clientId);
        this.activeStreams--;
      });
    };
  }

  /**
   * Send SSE message to client
   */
  send(clientId: string, message: StreamMessage): boolean {
    const client = this.clients.get(clientId);
    if (!client) {
      return false;
    }

    try {
      client.write(`data: ${JSON.stringify(message)}\n\n`);
      return true;
    } catch (error) {
      console.error(`[SSE] Send error for ${clientId}:`, error);
      this.clients.delete(clientId);
      return false;
    }
  }

  /**
   * Broadcast to all clients
   */
  broadcast(message: StreamMessage): void {
    for (const [clientId, client] of this.clients.entries()) {
      this.send(clientId, message);
    }
  }

  /**
   * Stream AI response (token by token)
   */
  async streamAIResponse(
    clientId: string,
    aiStream: AsyncIterable<string>,
    conversationId?: string
  ): Promise<void> {
    try {
      for await (const token of aiStream) {
        this.send(clientId, {
          id: this.generateId(),
          type: 'token',
          data: token,
          timestamp: Date.now(),
          conversationId,
        });
      }

      // Send done event
      this.send(clientId, {
        id: this.generateId(),
        type: 'done',
        data: { conversationId },
        timestamp: Date.now(),
      });
    } catch (error: any) {
      this.send(clientId, {
        id: this.generateId(),
        type: 'error',
        data: { message: error.message },
        timestamp: Date.now(),
        conversationId,
      });
    }
  }

  /**
   * Rate limiting middleware
   */
  rateLimitMiddleware(): (req: Request, res: Response, next: NextFunction) => void {
    return (req: Request, res: Response, next: NextFunction) => {
      const clientId = req.ip || req.query.clientId as string || 'unknown';
      const now = Date.now();
      const windowMs = 60000; // 1 minute window

      // Get client requests
      let clientRequests = this.rateLimitMap.get(clientId) || [];
      
      // Remove old requests
      clientRequests = clientRequests.filter(time => now - time < windowMs);

      // Check rate limit
      if (clientRequests.length >= this.config.rateLimit!) {
        if (this.config.debug) {
          console.log(`[Rate Limit] Exceeded for ${clientId}`);
        }
        
        return res.status(429).json({
          error: 'Rate limit exceeded',
          retryAfter: Math.ceil((clientRequests[0] + windowMs - now) / 1000),
        });
      }

      // Add current request
      clientRequests.push(now);
      this.rateLimitMap.set(clientId, clientRequests);

      next();
    };
  }

  /**
   * Cache AI response
   */
  cacheResponse(key: string, data: any, ttl?: number): void {
    const expiry = Date.now() + (ttl || this.config.cacheTTL!) * 1000;
    this.cache.set(key, { data, expiry });

    if (this.config.debug) {
      console.log(`[Cache] Stored: ${key}`);
    }
  }

  /**
   * Get cached response
   */
  getCachedResponse(key: string): any | null {
    const cached = this.cache.get(key);
    if (!cached) {
      return null;
    }

    // Check expiry
    if (Date.now() > cached.expiry) {
      this.cache.delete(key);
      return null;
    }

    return cached.data;
  }

  /**
   * Create conversation context
   */
  createConversation(id?: string): ConversationContext {
    const conversationId = id || this.generateId();
    const context: ConversationContext = {
      id: conversationId,
      messages: [],
      createdAt: Date.now(),
      lastActivity: Date.now(),
      tokenCount: 0,
    };

    this.conversations.set(conversationId, context);
    return context;
  }

  /**
   * Add message to conversation
   */
  addMessage(conversationId: string, role: string, content: string): void {
    const conversation = this.conversations.get(conversationId);
    if (!conversation) {
      throw new Error(`Conversation ${conversationId} not found`);
    }

    conversation.messages.push({ role, content });
    conversation.lastActivity = Date.now();
    conversation.tokenCount += this.estimateTokens(content);
  }

  /**
   * Get conversation context
   */
  getConversation(conversationId: string): ConversationContext | null {
    return this.conversations.get(conversationId) || null;
  }

  /**
   * Clear old conversations (cleanup)
   */
  cleanupConversations(maxAge: number = 3600000): void { // 1 hour default
    const now = Date.now();
    for (const [id, conversation] of this.conversations.entries()) {
      if (now - conversation.lastActivity > maxAge) {
        this.conversations.delete(id);
        if (this.config.debug) {
          console.log(`[Cleanup] Removed conversation: ${id}`);
        }
      }
    }
  }

  /**
   * Estimate token count (rough approximation)
   */
  private estimateTokens(text: string): number {
    // Rough estimate: 1 token ≈ 4 characters
    return Math.ceil(text.length / 4);
  }

  /**
   * Generate unique ID
   */
  private generateId(): string {
    return `${Date.now()}-${Math.random().toString(36).substr(2, 9)}`;
  }

  /**
   * Get stats
   */
  getStats(): any {
    return {
      activeStreams: this.activeStreams,
      connectedClients: this.clients.size,
      activeConversations: this.conversations.size,
      cacheSize: this.cache.size,
      rateLimitedClients: this.rateLimitMap.size,
    };
  }
}

export default AIStreamingServer;