/**
 * WebSocket Server - Handles WebSocket JSON-RPC requests and subscriptions
 */

import WebSocket, { WebSocketServer } from 'ws';
import { RpcHandler, JsonRpcRequest, JsonRpcResponse } from '../handlers/rpc-handler';
import { BesuClient } from '../clients/besu-client';
import { v4 as uuidv4 } from 'uuid';

interface ClientConnection {
  ws: WebSocket;
  upstreamWs?: WebSocket;
  subscriptionMap: Map<string, string>; // clientSubId -> upstreamSubId
  requestId: string;
}

export class WsServer {
  private wss: WebSocketServer;
  private port: number;
  private rpcHandler: RpcHandler;
  private besuClient: BesuClient;
  private connections: Map<string, ClientConnection> = new Map();
  private pingInterval?: NodeJS.Timeout;

  constructor(port: number, rpcHandler: RpcHandler, besuClient: BesuClient) {
    this.port = port;
    this.rpcHandler = rpcHandler;
    this.besuClient = besuClient;
    this.wss = new WebSocketServer({ port });
    this.setupServer();
  }

  private setupServer() {
    this.wss.on('connection', (ws: WebSocket) => {
      const connectionId = uuidv4();
      const connection: ClientConnection = {
        ws,
        subscriptionMap: new Map(),
        requestId: connectionId,
      };

      this.connections.set(connectionId, connection);

      console.log(`[${connectionId}] WebSocket client connected`);

      // Set up upstream WebSocket connection for subscriptions
      this.besuClient.createWebSocketConnection(connectionId)
        .then((upstreamWs) => {
          connection.upstreamWs = upstreamWs;

          // Forward upstream messages to client
          upstreamWs.on('message', (data: WebSocket.Data) => {
            try {
              const message = JSON.parse(data.toString());
              
              // If it's a subscription notification, map the subscription ID
              if (message.method === 'eth_subscription' && message.params) {
                const upstreamSubId = message.params.subscription;
                // Find the client subscription ID
                for (const [clientSubId, mappedUpstreamSubId] of connection.subscriptionMap.entries()) {
                  if (mappedUpstreamSubId === upstreamSubId) {
                    message.params.subscription = clientSubId;
                    break;
                  }
                }
              }

              ws.send(JSON.stringify(message));
            } catch (error) {
              console.error(`[${connectionId}] Error forwarding upstream message:`, error);
            }
          });

          upstreamWs.on('error', (error: Error) => {
            console.error(`[${connectionId}] Upstream WebSocket error:`, error);
          });

          upstreamWs.on('close', () => {
            console.log(`[${connectionId}] Upstream WebSocket closed`);
            connection.upstreamWs = undefined;
          });
        })
        .catch((error) => {
          console.error(`[${connectionId}] Failed to create upstream WebSocket:`, error);
        });

      // Handle client messages
      ws.on('message', async (data: WebSocket.Data) => {
        try {
          const message = JSON.parse(data.toString());
          await this.handleMessage(connectionId, message);
        } catch (error: unknown) {
          const err = error as Error;
          console.error(`[${connectionId}] Error handling message:`, err);
            const errorResponse: JsonRpcResponse = {
              jsonrpc: '2.0',
              error: {
                code: -32700,
                message: 'Parse error',
                data: err.message,
              },
              id: null,
            };
          ws.send(JSON.stringify(errorResponse));
        }
      });

      ws.on('close', () => {
        console.log(`[${connectionId}] WebSocket client disconnected`);
        this.cleanupConnection(connectionId);
      });

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

      // Send ping to keep connection alive
      const pingInterval = setInterval(() => {
        if (ws.readyState === WebSocket.OPEN) {
          ws.ping();
        }
      }, 30000); // 30 seconds

      ws.on('close', () => {
        clearInterval(pingInterval);
      });
    });

    // Start ping interval for all connections
    this.pingInterval = setInterval(() => {
      for (const [connectionId, connection] of this.connections.entries()) {
        if (connection.ws.readyState === WebSocket.OPEN) {
          connection.ws.ping();
        }
      }
    }, 30000);
  }

  private async handleMessage(connectionId: string, message: any) {
    const connection = this.connections.get(connectionId);
    if (!connection) {
      return;
    }

    // Handle batch requests
    if (Array.isArray(message)) {
      const requests: JsonRpcRequest[] = message;
      const responses = await this.rpcHandler.handleBatch(requests);
      connection.ws.send(JSON.stringify(responses));
      return;
    }

    // Handle subscription methods - pass through to upstream
    if (message.method === 'eth_subscribe' || message.method === 'eth_unsubscribe') {
      await this.handleSubscription(connectionId, message);
      return;
    }

    // Handle regular JSON-RPC requests
    const request: JsonRpcRequest = message;
    const response = await this.rpcHandler.handleRequest(request);
    
    if (response.error) {
      console.log(`[${connectionId}] Error: ${response.error.message} (code: ${response.error.code})`);
    } else {
      console.log(`[${connectionId}] Success: ${request.method}`);
    }

    connection.ws.send(JSON.stringify(response));
  }

  private async handleSubscription(connectionId: string, message: any) {
    const connection = this.connections.get(connectionId);
    if (!connection || !connection.upstreamWs) {
      const errorResponse: JsonRpcResponse = {
        jsonrpc: '2.0',
        error: {
          code: -32000,
          message: 'Server error',
          data: 'Upstream WebSocket connection not available',
        },
        id: message.id,
      };
      connection?.ws.send(JSON.stringify(errorResponse));
      return;
    }

    if (message.method === 'eth_subscribe') {
      // Generate a client subscription ID
      const clientSubId = uuidv4();
      
      // Forward to upstream with original request ID
      const upstreamRequest = {
        ...message,
        id: message.id,
      };

      connection.upstreamWs.send(JSON.stringify(upstreamRequest));

      // Store mapping (we'll update it when we get the response)
      // For now, we'll assume the upstream returns the subscription ID in the response
      // This is a simplified implementation - in production, you'd need to track the request/response correlation
      const upstreamResponseHandler = (data: WebSocket.Data) => {
        try {
          const response = JSON.parse(data.toString());
          if (response.id === message.id && response.result) {
            const upstreamSubId = response.result;
            connection.subscriptionMap.set(clientSubId, upstreamSubId);
            
            // Send response to client with client subscription ID
            connection.ws.send(JSON.stringify({
              ...response,
              result: clientSubId,
            }));
          }
        } catch (error) {
          console.error(`[${connectionId}] Error handling subscription response:`, error);
        }
      };

      // Note: This is a simplified implementation. In production, you'd need to properly correlate requests/responses
      connection.upstreamWs.once('message', upstreamResponseHandler);
      
    } else if (message.method === 'eth_unsubscribe') {
      // Map client subscription ID to upstream subscription ID
      const clientSubId = message.params[0];
      const upstreamSubId = connection.subscriptionMap.get(clientSubId);
      
      if (upstreamSubId) {
        const upstreamRequest = {
          ...message,
          params: [upstreamSubId],
          id: message.id,
        };
        connection.upstreamWs.send(JSON.stringify(upstreamRequest));
        connection.subscriptionMap.delete(clientSubId);
      } else {
        const errorResponse: JsonRpcResponse = {
          jsonrpc: '2.0',
          error: {
            code: -32602,
            message: 'Invalid params',
            data: `Subscription ${clientSubId} not found`,
          },
          id: message.id,
        };
        connection.ws.send(JSON.stringify(errorResponse));
      }
    }
  }

  private cleanupConnection(connectionId: string) {
    const connection = this.connections.get(connectionId);
    if (connection) {
      if (connection.upstreamWs) {
        connection.upstreamWs.close();
      }
      this.besuClient.closeWebSocketConnection(connectionId);
      this.connections.delete(connectionId);
    }
  }

  async start(): Promise<void> {
    return new Promise((resolve, reject) => {
      this.wss.on('listening', () => {
        console.log(`WebSocket server started on port ${this.port}`);
        resolve();
      });

      this.wss.on('error', (error: any) => {
        if (error.code === 'EADDRINUSE') {
          reject(new Error(`Port ${this.port} is already in use`));
        } else {
          reject(error);
        }
      });
    });
  }

  async stop(): Promise<void> {
    return new Promise((resolve) => {
      // Close all connections
      for (const [connectionId, connection] of this.connections.entries()) {
        this.cleanupConnection(connectionId);
      }

      if (this.pingInterval) {
        clearInterval(this.pingInterval);
      }

      this.wss.close(() => {
        console.log('WebSocket server stopped');
        resolve();
      });
    });
  }
}
