// services/websocket.service.js
const WebSocket = require('ws');
const AIStreamingService = require('./ai-streaming.service');
const database = require('../config/database');

class WebSocketService {
  constructor(server) {
    this.wss = new WebSocket.Server({ 
      server,
      path: '/ws/ai-streaming',
      perMessageDeflate: false
    });
    
    this.aiStreamingService = new AIStreamingService();
    this.connections = new Map(); // Track active connections
    this.setupWebSocketHandlers();
    
    // Cleanup old sessions every 5 minutes
    setInterval(() => {
      this.aiStreamingService.cleanupOldSessions();
    }, 5 * 60 * 1000);
  }

  setupWebSocketHandlers() {
    this.wss.on('connection', (ws, req) => {
      console.log('New WebSocket connection established');
      
      const connectionId = `conn_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`;
      this.connections.set(connectionId, {
        ws,
        sessionId: null,
        repositoryId: null,
        startTime: Date.now(),
        lastActivity: Date.now()
      });

      // Send welcome message
      ws.send(JSON.stringify({
        type: 'connection_established',
        connection_id: connectionId,
        timestamp: new Date().toISOString()
      }));

      ws.on('message', async (message) => {
        try {
          const data = JSON.parse(message);
          await this.handleMessage(connectionId, data);
        } catch (error) {
          console.error('Error handling WebSocket message:', error);
          this.sendError(connectionId, 'Invalid message format');
        }
      });

      ws.on('close', () => {
        console.log(`WebSocket connection ${connectionId} closed`);
        this.cleanupConnection(connectionId);
      });

      ws.on('error', (error) => {
        console.error(`WebSocket error for connection ${connectionId}:`, error);
        this.cleanupConnection(connectionId);
      });
    });
  }

  async handleMessage(connectionId, data) {
    const connection = this.connections.get(connectionId);
    if (!connection) return;

    connection.lastActivity = Date.now();

    switch (data.type) {
      case 'start_streaming':
        await this.handleStartStreaming(connectionId, data);
        break;
      
      case 'get_status':
        await this.handleGetStatus(connectionId, data);
        break;
      
      case 'cancel_streaming':
        await this.handleCancelStreaming(connectionId, data);
        break;
      
      case 'start_bulk_analysis':
        await this.handleStartBulkAnalysis(connectionId, data);
        break;
      
      case 'get_bulk_status':
        await this.handleGetBulkStatus(connectionId, data);
        break;
      
      case 'cancel_bulk_analysis':
        await this.handleCancelBulkAnalysis(connectionId, data);
        break;
      
      case 'ping':
        this.sendMessage(connectionId, { type: 'pong', timestamp: new Date().toISOString() });
        break;
      
      default:
        this.sendError(connectionId, `Unknown message type: ${data.type}`);
    }
  }

  async handleStartStreaming(connectionId, data) {
    try {
      const { repository_id, options = {} } = data;
      
      if (!repository_id) {
        this.sendError(connectionId, 'Repository ID is required');
        return;
      }

      // Validate repository exists
      const repoQuery = 'SELECT id, repository_name FROM all_repositories WHERE id = $1';
      const repoResult = await database.query(repoQuery, [repository_id]);
      
      if (repoResult.rows.length === 0) {
        this.sendError(connectionId, 'Repository not found');
        return;
      }

      // Parse options with auto-detection
      let file_types;
      let chunk_size;
      
      if (options.file_types === 'auto' || !options.file_types) {
        // Auto-detect file types from repository
        file_types = await this.aiStreamingService.getAvailableFileTypes(repository_id);
      } else {
        file_types = options.file_types;
      }
      
      if (options.chunk_size === 'auto' || !options.chunk_size) {
        // Auto-calculate optimal chunk size based on total files
        const totalFiles = await this.aiStreamingService.getRepositoryFilesCount(repository_id, {
          fileTypes: file_types,
          maxSize: options.max_size || 500000,
          includeBinary: options.include_binary || false,
          directoryFilter: options.directory_filter || '',
          excludePatterns: options.exclude_patterns || ['node_modules', 'dist', 'build', '.git', 'coverage', 'test', 'tests', '__tests__', 'spec', 'specs']
        });
        chunk_size = this.aiStreamingService.calculateOptimalChunkSize(totalFiles);
      } else {
        chunk_size = options.chunk_size;
      }
      
      const {
        max_size = 500000,
        include_binary = false,
        directory_filter = '',
        exclude_patterns = ['node_modules', 'dist', 'build', '.git', 'coverage', 'test', 'tests', '__tests__', 'spec', 'specs']
      } = options;

      // Get total files count
      const totalFiles = await this.aiStreamingService.getRepositoryFilesCount(repository_id, {
        fileTypes: file_types,
        maxSize: max_size,
        includeBinary: include_binary,
        directoryFilter: directory_filter,
        excludePatterns: exclude_patterns
      });

      if (totalFiles === 0) {
        this.sendError(connectionId, 'No files found matching the criteria');
        return;
      }

      // Calculate total chunks
      const totalChunks = Math.ceil(totalFiles / chunk_size);

      // Create streaming session
      const sessionId = this.aiStreamingService.createStreamingSession(repository_id, {
        fileTypes: file_types,
        maxSize: max_size,
        includeBinary: include_binary,
        directoryFilter: directory_filter,
        excludePatterns: exclude_patterns,
        chunkSize: chunk_size
      });

      // Update connection with session info
      const connection = this.connections.get(connectionId);
      connection.sessionId = sessionId;
      connection.repositoryId = repository_id;

      // Update session with total info
      this.aiStreamingService.updateStreamingSession(sessionId, {
        totalFiles,
        totalChunks,
        status: 'ready'
      });

      // Get repository info
      const repositoryInfo = await this.aiStreamingService.getRepositoryInfo(repository_id);

      // Send initial response
      this.sendMessage(connectionId, {
        type: 'streaming_started',
        session_id: sessionId,
        repository_info: {
          id: repositoryInfo.id,
          name: repositoryInfo.name,
          full_name: repositoryInfo.full_name,
          description: repositoryInfo.description,
          language: repositoryInfo.language,
          size: repositoryInfo.size,
          local_path: repositoryInfo.local_path
        },
        streaming_config: {
          total_files: totalFiles,
          total_chunks: totalChunks,
          chunk_size: chunk_size,
          file_types: file_types,
          max_size_bytes: max_size,
          include_binary: include_binary,
          directory_filter: directory_filter,
          exclude_patterns: exclude_patterns
        },
        status: 'ready'
      });

      // Start streaming in background
      this.startStreaming(connectionId, sessionId, repository_id, {
        fileTypes: file_types,
        maxSize: max_size,
        includeBinary: include_binary,
        directoryFilter: directory_filter,
        excludePatterns: exclude_patterns,
        chunkSize: chunk_size
      });

    } catch (error) {
      console.error('Error starting streaming:', error);
      this.sendError(connectionId, error.message);
    }
  }

  async startStreaming(connectionId, sessionId, repositoryId, options) {
    try {
      const session = this.aiStreamingService.getStreamingSession(sessionId);
      if (!session) return;

      const { totalFiles, totalChunks } = session;
      const { chunkSize } = options;
      
      let currentChunk = 0;
      let processedFiles = 0;

      // Update status to streaming
      this.aiStreamingService.updateStreamingSession(sessionId, {
        status: 'streaming'
      });

      this.sendMessage(connectionId, {
        type: 'streaming_progress',
        status: 'streaming',
        progress: {
          current_chunk: 0,
          total_chunks: totalChunks,
          processed_files: 0,
          total_files: totalFiles,
          percentage: 0
        }
      });

      while (currentChunk < totalChunks) {
        try {
          const offset = currentChunk * chunkSize;
          const files = await this.aiStreamingService.getFilesChunk(repositoryId, offset, chunkSize, options);

          if (files.length === 0) {
            break;
          }

          // Process chunk
          const chunkResult = await this.aiStreamingService.processFilesChunk(
            files, 
            currentChunk + 1, 
            totalChunks
          );

          processedFiles += chunkResult.files_processed;

          // Update session
          this.aiStreamingService.updateStreamingSession(sessionId, {
            currentChunk: currentChunk + 1,
            processedFiles,
            status: 'streaming'
          });

          // Send chunk data
          this.sendMessage(connectionId, {
            type: 'chunk_data',
            chunk_data: chunkResult,
            progress: {
              current_chunk: currentChunk + 1,
              total_chunks: totalChunks,
              processed_files: processedFiles,
              total_files: totalFiles,
              percentage: Math.round((processedFiles / totalFiles) * 100)
            }
          });

          currentChunk++;

          // Small delay to prevent overwhelming the client
          await new Promise(resolve => setTimeout(resolve, 100));

        } catch (error) {
          console.error(`Error processing chunk ${currentChunk + 1}:`, error);
          
          // Send error for this chunk
          this.sendMessage(connectionId, {
            type: 'chunk_error',
            chunk_number: currentChunk + 1,
            error: error.message
          });

          currentChunk++;
        }
      }

      // Send completion message
      this.sendMessage(connectionId, {
        type: 'streaming_complete',
        session_id: sessionId,
        total_files_processed: processedFiles,
        total_chunks_processed: currentChunk,
        processing_time_ms: Date.now() - session.startTime
      });

      // Clean up session
      this.aiStreamingService.removeStreamingSession(sessionId);

    } catch (error) {
      console.error('Error in streaming process:', error);
      this.sendError(connectionId, error.message);
    }
  }

  async handleGetStatus(connectionId, data) {
    const connection = this.connections.get(connectionId);
    if (!connection || !connection.sessionId) {
      this.sendError(connectionId, 'No active streaming session');
      return;
    }

    const session = this.aiStreamingService.getStreamingSession(connection.sessionId);
    if (!session) {
      this.sendError(connectionId, 'Streaming session not found');
      return;
    }

    this.sendMessage(connectionId, {
      type: 'status_update',
      session: {
        session_id: connection.sessionId,
        repository_id: session.repositoryId,
        status: session.status,
        current_chunk: session.currentChunk,
        total_chunks: session.totalChunks,
        total_files: session.totalFiles,
        processed_files: session.processedFiles,
        progress_percentage: session.totalFiles > 0 ? 
          Math.round((session.processedFiles / session.totalFiles) * 100) : 0,
        start_time: new Date(session.startTime).toISOString(),
        last_activity: new Date(session.lastActivity).toISOString()
      }
    });
  }

  async handleCancelStreaming(connectionId, data) {
    const connection = this.connections.get(connectionId);
    if (!connection || !connection.sessionId) {
      this.sendError(connectionId, 'No active streaming session');
      return;
    }

    this.aiStreamingService.removeStreamingSession(connection.sessionId);
    connection.sessionId = null;
    connection.repositoryId = null;

    this.sendMessage(connectionId, {
      type: 'streaming_cancelled',
      message: 'Streaming session cancelled successfully'
    });
  }

  sendMessage(connectionId, message) {
    const connection = this.connections.get(connectionId);
    if (connection && connection.ws.readyState === WebSocket.OPEN) {
      try {
        connection.ws.send(JSON.stringify(message));
      } catch (error) {
        console.error('Error sending WebSocket message:', error);
      }
    }
  }

  sendError(connectionId, errorMessage) {
    this.sendMessage(connectionId, {
      type: 'error',
      error: errorMessage,
      timestamp: new Date().toISOString()
    });
  }

  cleanupConnection(connectionId) {
    const connection = this.connections.get(connectionId);
    if (connection && connection.sessionId) {
      this.aiStreamingService.removeStreamingSession(connection.sessionId);
    }
    this.connections.delete(connectionId);
  }

  // ==================== BULK ANALYSIS WEBSOCKET HANDLERS ====================

  async handleStartBulkAnalysis(connectionId, data) {
    try {
      const { repository_id, commit_ids = [], options = {} } = data;
      
      if (!repository_id) {
        this.sendError(connectionId, 'Repository ID is required');
        return;
      }

      if (!Array.isArray(commit_ids) || commit_ids.length === 0) {
        this.sendError(connectionId, 'commit_ids must be a non-empty array');
        return;
      }

      // Limit the number of commits to prevent overload
      const maxCommits = 50;
      if (commit_ids.length > maxCommits) {
        this.sendError(connectionId, `Maximum ${maxCommits} commits allowed per request`);
        return;
      }

      // Validate repository exists
      const repoQuery = 'SELECT id, repository_name, owner_name FROM all_repositories WHERE id = $1';
      const repoResult = await database.query(repoQuery, [repository_id]);
      
      if (repoResult.rows.length === 0) {
        this.sendError(connectionId, 'Repository not found');
        return;
      }

      // Create bulk analysis session
      const sessionId = `bulk_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`;
      
      // Update connection with session info
      const connection = this.connections.get(connectionId);
      connection.sessionId = sessionId;
      connection.repositoryId = repository_id;

      // Send initial response
      this.sendMessage(connectionId, {
        type: 'bulk_analysis_started',
        session_id: sessionId,
        repository_id: repository_id,
        total_commits: commit_ids.length,
        status: 'processing'
      });

      // Start bulk analysis in background
      this.startBulkAnalysis(connectionId, sessionId, repository_id, commit_ids, options);

    } catch (error) {
      console.error('Error starting bulk analysis:', error);
      this.sendError(connectionId, error.message);
    }
  }

  async startBulkAnalysis(connectionId, sessionId, repositoryId, commitIds, options) {
    try {
      console.log(`🔍 Starting bulk analysis for ${commitIds.length} commits`);

      // Send progress update
      this.sendMessage(connectionId, {
        type: 'bulk_analysis_progress',
        session_id: sessionId,
        status: 'fetching_commits',
        progress: {
          current_step: 1,
          total_steps: 4,
          step_name: 'Fetching commit details',
          percentage: 0
        }
      });

      // Step 1: Get bulk commit details
      const commitResults = await this.aiStreamingService.getBulkCommitDetails(commitIds);
      
      this.sendMessage(connectionId, {
        type: 'bulk_analysis_progress',
        session_id: sessionId,
        status: 'reading_files',
        progress: {
          current_step: 2,
          total_steps: 4,
          step_name: 'Reading diff files',
          percentage: 25,
          commits_processed: commitResults.length
        }
      });

      // Step 2: Read diff files in batch
      const enrichedResults = await this.aiStreamingService.batchReadDiffFiles(commitResults);
      
      this.sendMessage(connectionId, {
        type: 'bulk_analysis_progress',
        session_id: sessionId,
        status: 'generating_summary',
        progress: {
          current_step: 3,
          total_steps: 4,
          step_name: 'Generating analysis summary',
          percentage: 50,
          files_read: enrichedResults.filter(r => r.diffContent?.content).length
        }
      });

      // Step 3: Get analysis summary
      const summary = await this.aiStreamingService.getBulkAnalysisSummary(enrichedResults);
      
      this.sendMessage(connectionId, {
        type: 'bulk_analysis_progress',
        session_id: sessionId,
        status: 'preparing_ai_inputs',
        progress: {
          current_step: 4,
          total_steps: 4,
          step_name: 'Preparing AI analysis inputs',
          percentage: 75
        }
      });

      // Step 4: Process for AI analysis
      const aiInputs = await this.aiStreamingService.processBulkCommitsForAI(enrichedResults);

      // Send completion message
      this.sendMessage(connectionId, {
        type: 'bulk_analysis_complete',
        session_id: sessionId,
        repository_id: repositoryId,
        summary: summary,
        commits: enrichedResults,
        ai_ready_commits: aiInputs.length,
        ai_inputs: options.include_content !== false ? aiInputs : [],
        processing_time_ms: Date.now() - this.connections.get(connectionId).startTime
      });

    } catch (error) {
      console.error('Error in bulk analysis process:', error);
      this.sendError(connectionId, error.message);
    }
  }

  async handleGetBulkStatus(connectionId, data) {
    const connection = this.connections.get(connectionId);
    if (!connection || !connection.sessionId) {
      this.sendError(connectionId, 'No active bulk analysis session');
      return;
    }

    this.sendMessage(connectionId, {
      type: 'bulk_analysis_status',
      session_id: connection.sessionId,
      repository_id: connection.repositoryId,
      status: 'active',
      start_time: new Date(connection.startTime).toISOString(),
      last_activity: new Date(connection.lastActivity).toISOString()
    });
  }

  async handleCancelBulkAnalysis(connectionId, data) {
    const connection = this.connections.get(connectionId);
    if (!connection || !connection.sessionId) {
      this.sendError(connectionId, 'No active bulk analysis session');
      return;
    }

    // Clean up session
    connection.sessionId = null;
    connection.repositoryId = null;

    this.sendMessage(connectionId, {
      type: 'bulk_analysis_cancelled',
      message: 'Bulk analysis session cancelled successfully'
    });
  }

  getActiveConnections() {
    return Array.from(this.connections.entries()).map(([connectionId, connection]) => ({
      connection_id: connectionId,
      session_id: connection.sessionId,
      repository_id: connection.repositoryId,
      start_time: new Date(connection.startTime).toISOString(),
      last_activity: new Date(connection.lastActivity).toISOString()
    }));
  }
}

module.exports = WebSocketService;
