import WebSocket from 'ws';
import { Server as HTTPServer } from 'http';
import jwt from 'jsonwebtoken';
import { BrowserPool } from '../core/BrowserPool';
import { BrowserSession } from '../core/BrowserSession';
import { createLogger } from '../utils/logger';
import { config } from '../config';
import { rewriteResourceUrls } from '../utils/urlRewriter';
import {
  AuthPayload,
  DOMUpdate,
  NavigationRequest,
  InteractionEvent,
  PreviewInteractionEvent,
} from '../types';
import {
  wsConnectionsCounter,
  activeWsConnectionsGauge,
  navigationDurationHistogram,
  navigationsCounter,
} from '../utils/metrics';

const logger = createLogger('StreamingServer');

interface AuthenticatedWebSocket extends WebSocket {
  user?: AuthPayload;
  token?: string;
  sessionId?: string;
  isAlive?: boolean;
}

/**
 * WebSocket streaming server for real-time DOM updates
 */
export class StreamingServer {
  private wss: WebSocket.Server;
  private browserPool: BrowserPool;
  private clients: Map<string, Set<AuthenticatedWebSocket>> = new Map();
  private heartbeatInterval: NodeJS.Timeout | null = null;

  constructor(server: HTTPServer, browserPool: BrowserPool) {
    this.browserPool = browserPool;

    // Create WebSocket server
    this.wss = new WebSocket.Server({
      server,
      path: '/ws',
    });

    this.setupWebSocketServer();
    this.startHeartbeat();

    logger.info('WebSocket streaming server initialized');
  }

  /**
   * Setup WebSocket server handlers
   */
  private setupWebSocketServer(): void {
    this.wss.on('connection', (ws: AuthenticatedWebSocket, req) => {
      logger.info('New WebSocket connection', { ip: req.socket.remoteAddress });

      // Mark as alive for heartbeat
      ws.isAlive = true;

      ws.on('pong', () => {
        ws.isAlive = true;
      });

      // Handle authentication
      ws.on('message', async (message: string) => {
        try {
          const data = JSON.parse(message);

          // Handle authentication message
          if (data.type === 'auth') {
            await this.handleAuthentication(ws, data.token);
            return;
          }

          // Require authentication for other messages
          if (!ws.user) {
            ws.send(JSON.stringify({
              type: 'error',
              message: 'Authentication required',
            }));
            return;
          }

          // Handle different message types
          switch (data.type) {
            case 'subscribe':
              await this.handleSubscribe(ws, data.sessionId);
              break;

            case 'navigate':
              await this.handleNavigate(ws, data);
              break;

            case 'interact':
              await this.handleInteract(ws, data);
              break;

            case 'preview_interact':
              await this.handlePreviewInteract(ws, data);
              break;

            case 'unsubscribe':
              this.handleUnsubscribe(ws);
              break;

            default:
              ws.send(JSON.stringify({
                type: 'error',
                message: 'Unknown message type',
              }));
          }
        } catch (error) {
          logger.error('Error handling WebSocket message', {
            error: error instanceof Error ? error.message : String(error),
          });

          ws.send(JSON.stringify({
            type: 'error',
            message: error instanceof Error ? error.message : 'Internal error',
          }));
        }
      });

      ws.on('close', () => {
        this.handleDisconnect(ws);
      });

      ws.on('error', (error) => {
        logger.error('WebSocket error', { error: error.message });
        this.handleDisconnect(ws);
      });

      // Send welcome message
      ws.send(JSON.stringify({
        type: 'connected',
        message: 'WebSocket connection established. Please authenticate.',
      }));
    });

    this.wss.on('error', (error) => {
      logger.error('WebSocket server error', { error: error.message });
    });
  }

  /**
   * Handle authentication
   */
  private async handleAuthentication(
    ws: AuthenticatedWebSocket,
    token: string
  ): Promise<void> {
    try {
      const decoded = jwt.verify(token, config.jwt.secret) as AuthPayload;
      ws.user = decoded;
      ws.token = token; // Store token for URL rewriting

      wsConnectionsCounter.inc({ status: 'authenticated' });
      activeWsConnectionsGauge.inc();

      logger.info('WebSocket authenticated', { userId: decoded.userId });

      ws.send(JSON.stringify({
        type: 'authenticated',
        user: {
          userId: decoded.userId,
          email: decoded.email,
        },
      }));
    } catch (error) {
      wsConnectionsCounter.inc({ status: 'auth_failed' });

      logger.warn('WebSocket authentication failed', {
        error: error instanceof Error ? error.message : String(error),
      });

      ws.send(JSON.stringify({
        type: 'error',
        message: 'Authentication failed',
      }));

      ws.close();
    }
  }

  /**
   * Handle session subscription
   */
  private async handleSubscribe(
    ws: AuthenticatedWebSocket,
    sessionId: string
  ): Promise<void> {
    try {
      // Get or create session
      const session = await this.browserPool.getSession({
        sessionId,
        userId: ws.user!.userId,
        deviceMode: 'desktop', // Will reuse existing if available
      });

      // Store client connection
      if (!this.clients.has(sessionId)) {
        this.clients.set(sessionId, new Set());
      }
      this.clients.get(sessionId)!.add(ws);
      ws.sessionId = sessionId;

      // Subscribe to session updates
      if (session instanceof BrowserSession) {
        session.subscribe((update: DOMUpdate) => {
          this.broadcastUpdate(sessionId, update);
        });
      }

      logger.info('Client subscribed to session', {
        sessionId,
        userId: ws.user!.userId,
      });

      ws.send(JSON.stringify({
        type: 'subscribed',
        sessionId,
        state: session.state,
      }));

      // Send initial HTML snapshot to the newly subscribed client
      if (session instanceof BrowserSession && ws.token) {
        try {
          logger.info('Fetching HTML for initial snapshot', { sessionId });
          const html = await session.getHTML();
          const currentUrl = session.state.currentUrl;
          logger.info('About to rewrite HTML URLs', { sessionId, htmlLength: html.length, currentUrl });
          const rewrittenHtml = rewriteResourceUrls(html, currentUrl, sessionId, ws.token);
          logger.info('HTML rewritten, sending to client', { sessionId, rewrittenLength: rewrittenHtml.length });

          ws.send(JSON.stringify({
            type: 'update',
            update: {
              type: 'html_snapshot',
              timestamp: Date.now(),
              data: { html: rewrittenHtml },
            },
          }));
        } catch (error) {
          logger.error('Failed to send initial HTML snapshot', {
            sessionId,
            error: error instanceof Error ? error.message : String(error),
          });
        }
      }
    } catch (error) {
      logger.error('Subscription failed', {
        sessionId,
        error: error instanceof Error ? error.message : String(error),
      });

      ws.send(JSON.stringify({
        type: 'error',
        message: 'Failed to subscribe to session',
      }));
    }
  }

  /**
   * Handle navigation request
   */
  private async handleNavigate(
    ws: AuthenticatedWebSocket,
    data: NavigationRequest
  ): Promise<void> {
    try {
      if (!ws.sessionId) {
        ws.send(JSON.stringify({
          type: 'error',
          message: 'Not subscribed to any session',
        }));
        return;
      }

      const session = await this.browserPool.getSession({
        sessionId: ws.sessionId,
        userId: ws.user!.userId,
        deviceMode: 'desktop',
      });

      const start = Date.now();

      try {
        await session.navigate(data.url);
        const duration = (Date.now() - start) / 1000;
        navigationDurationHistogram.observe({ status: 'success' }, duration);
        navigationsCounter.inc({ status: 'success' });

        ws.send(JSON.stringify({
          type: 'navigation_complete',
          url: data.url,
          currentUrl: session.state.currentUrl,
        }));
      } catch (error) {
        const duration = (Date.now() - start) / 1000;
        navigationDurationHistogram.observe({ status: 'error' }, duration);
        navigationsCounter.inc({ status: 'error' });
        throw error;
      }
    } catch (error) {
      logger.error('Navigation failed', {
        sessionId: ws.sessionId,
        error: error instanceof Error ? error.message : String(error),
      });

      ws.send(JSON.stringify({
        type: 'error',
        message: 'Navigation failed',
      }));
    }
  }

  /**
   * Handle interaction event
   */
  private async handleInteract(
    ws: AuthenticatedWebSocket,
    data: InteractionEvent
  ): Promise<void> {
    try {
      if (!ws.sessionId) {
        ws.send(JSON.stringify({
          type: 'error',
          message: 'Not subscribed to any session',
        }));
        return;
      }

      const session = await this.browserPool.getSession({
        sessionId: ws.sessionId,
        userId: ws.user!.userId,
        deviceMode: 'desktop',
      });

      await session.interact({ ...data, sessionId: ws.sessionId });

      ws.send(JSON.stringify({
        type: 'interaction_complete',
      }));
    } catch (error) {
      logger.error('Interaction failed', {
        sessionId: ws.sessionId,
        error: error instanceof Error ? error.message : String(error),
      });

      ws.send(JSON.stringify({
        type: 'error',
        message: 'Interaction failed',
      }));
    }
  }

  /**
   * Handle preview interaction event
   */
  private async handlePreviewInteract(
    ws: AuthenticatedWebSocket,
    data: any
  ): Promise<void> {
    try {
      if (!ws.sessionId) {
        ws.send(JSON.stringify({
          type: 'error',
          message: 'Not subscribed to any session',
        }));
        return;
      }

      const session = await this.browserPool.getSession({
        sessionId: ws.sessionId,
        userId: ws.user!.userId,
        deviceMode: 'desktop',
      });

      // Extract eventType and reconstruct the event
      const { eventType, ...eventData } = data;
      const previewEvent: PreviewInteractionEvent = {
        type: eventType,
        sessionId: ws.sessionId,
        ...eventData,
      };

      logger.info('Handling preview interaction', {
        sessionId: ws.sessionId,
        eventType,
        eventData,
      });

      // Cast to BrowserSession to access handlePreviewInteraction
      if (session instanceof BrowserSession) {
        await session.handlePreviewInteraction(previewEvent);
      }

      ws.send(JSON.stringify({
        type: 'preview_interaction_complete',
      }));
    } catch (error) {
      logger.error('Preview interaction failed', {
        sessionId: ws.sessionId,
        error: error instanceof Error ? error.message : String(error),
      });

      ws.send(JSON.stringify({
        type: 'error',
        message: 'Preview interaction failed',
      }));
    }
  }

  /**
   * Handle unsubscribe
   */
  private handleUnsubscribe(ws: AuthenticatedWebSocket): void {
    if (ws.sessionId) {
      const clients = this.clients.get(ws.sessionId);
      if (clients) {
        clients.delete(ws);
        if (clients.size === 0) {
          this.clients.delete(ws.sessionId);
        }
      }

      logger.info('Client unsubscribed', { sessionId: ws.sessionId });

      ws.sessionId = undefined;

      ws.send(JSON.stringify({
        type: 'unsubscribed',
      }));
    }
  }

  /**
   * Handle client disconnect
   */
  private handleDisconnect(ws: AuthenticatedWebSocket): void {
    this.handleUnsubscribe(ws);

    if (ws.user) {
      activeWsConnectionsGauge.dec();
      wsConnectionsCounter.inc({ status: 'disconnected' });

      logger.info('WebSocket disconnected', { userId: ws.user.userId });
    }
  }

  /**
   * Broadcast update to all clients subscribed to a session
   */
  private async broadcastUpdate(sessionId: string, update: DOMUpdate): Promise<void> {
    const clients = this.clients.get(sessionId);

    if (clients && clients.size > 0) {
      // If it's an HTML snapshot, we need to rewrite URLs for each client (since each has their own token)
      if (update.type === 'html_snapshot' && update.data?.html) {
        // Get session to retrieve current URL
        try {
          logger.info('Broadcasting HTML snapshot update', { sessionId, clientCount: clients.size });
          // Get any client to fetch the session (they all share the same session)
          const anyClient = Array.from(clients)[0];
          if (!anyClient.user) return;

          const session = await this.browserPool.getSession({
            sessionId,
            userId: anyClient.user.userId,
            deviceMode: 'desktop',
          });

          const currentUrl = session.state.currentUrl;

          // Send rewritten HTML to each client
          clients.forEach((client) => {
            if (client.readyState === WebSocket.OPEN && client.token) {
              logger.info('Rewriting HTML for client', { sessionId, currentUrl });
              const rewrittenHtml = rewriteResourceUrls(
                update.data.html,
                currentUrl,
                sessionId,
                client.token
              );

              const message = JSON.stringify({
                type: 'update',
                update: {
                  ...update,
                  data: { html: rewrittenHtml },
                },
              });

              client.send(message);
            }
          });
        } catch (error) {
          logger.error('Failed to rewrite HTML snapshot', {
            sessionId,
            error: error instanceof Error ? error.message : String(error),
          });
        }
      } else {
        // For non-HTML updates, broadcast as-is
        const message = JSON.stringify({
          type: 'update',
          update,
        });

        clients.forEach((client) => {
          if (client.readyState === WebSocket.OPEN) {
            client.send(message);
          }
        });
      }
    }
  }

  /**
   * Start heartbeat to detect dead connections
   */
  private startHeartbeat(): void {
    this.heartbeatInterval = setInterval(() => {
      this.wss.clients.forEach((ws: WebSocket) => {
        const client = ws as AuthenticatedWebSocket;

        if (client.isAlive === false) {
          logger.info('Terminating inactive WebSocket connection');
          return client.terminate();
        }

        client.isAlive = false;
        client.ping();
      });
    }, 30000); // 30 seconds

    logger.info('WebSocket heartbeat started');
  }

  /**
   * Shutdown the WebSocket server
   */
  async shutdown(): Promise<void> {
    logger.info('Shutting down WebSocket server');

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

    // Close all connections
    this.wss.clients.forEach((ws) => {
      ws.close(1000, 'Server shutting down');
    });

    return new Promise((resolve) => {
      this.wss.close(() => {
        logger.info('WebSocket server closed');
        resolve();
      });
    });
  }
}