import { WebSocket, WebSocketServer } from 'ws';
import { IncomingMessage } from 'http';
import { Server } from 'http';
import * as jwt from 'jsonwebtoken';

import type {
  IPlaygroundWebSocketMessage,
  TPlaygroundWebSocketMessageKind,
} from '@robota-sdk/agent-playground';
import type { ILogger, TUniversalValue } from '@robota-sdk/agent-core';
import { createLogger } from '@robota-sdk/agent-core';

const JWT_PART_COUNT = 3;
const CLEANUP_INTERVAL_MS = 30000;

function isUniversalObjectValue(value: TUniversalValue): value is Record<string, TUniversalValue> {
  return (
    typeof value === 'object' && value !== null && !Array.isArray(value) && !(value instanceof Date)
  );
}

function isPlaygroundAuthPayload(
  value: TUniversalValue | undefined,
): value is { userId: string; sessionId: string; token: string } {
  if (!value) return false;
  if (!isUniversalObjectValue(value)) return false;
  return (
    typeof value.userId === 'string' &&
    typeof value.sessionId === 'string' &&
    typeof value.token === 'string'
  );
}

export interface IPlaygroundClient {
  id: string;
  userId?: string;
  sessionId?: string;
  ws: WebSocket;
  lastActivity: Date;
  isAuthenticated: boolean;
}

/**
 * WebSocket Server for Playground real-time communication
 * Handles authentication, message routing, and connection management
 */
export class PlaygroundWebSocketServer {
  private wss: WebSocketServer;
  private clients = new Map<string, IPlaygroundClient>();
  private userSessions = new Map<string, Set<string>>(); // userId -> Set<clientId>
  private cleanupInterval: ReturnType<typeof setInterval>;
  private logger: ILogger;

  constructor(server: Server, logger?: ILogger) {
    this.logger = logger ?? createLogger('PlaygroundWebSocketServer');
    this.wss = new WebSocketServer({
      server,
      path: '/ws/playground',
    });

    this.wss.on('connection', this.handleConnection.bind(this));

    // Cleanup inactive connections every 30 seconds
    this.cleanupInterval = setInterval(
      this.cleanupInactiveConnections.bind(this),
      CLEANUP_INTERVAL_MS,
    );

    if (!process.env.JWT_SECRET) {
      this.logger.warn(
        'JWT_SECRET environment variable is not set. ' +
          'JWT validation will use format-only checking (development mode). ' +
          'Set JWT_SECRET for production use.',
      );
    }

    this.logger.info('PlaygroundWebSocketServer initialized on /ws/playground');
  }

  private handleConnection(ws: WebSocket, _req: IncomingMessage): void {
    const clientId = this.generateClientId();
    const client: IPlaygroundClient = {
      id: clientId,
      ws,
      lastActivity: new Date(),
      isAuthenticated: false,
    };

    this.clients.set(clientId, client);
    this.logger.info(`New WebSocket connection: ${clientId}`);

    // Set up message handling
    ws.on('message', (data: Buffer) => {
      try {
        const message: IPlaygroundWebSocketMessage = JSON.parse(data.toString());
        this.handleMessage(clientId, message);
      } catch (error) {
        this.logger.error(`Invalid message from ${clientId}:`, error as Error);
        this.sendError(clientId, 'Invalid message format');
      }
    });

    // Handle connection close
    ws.on('close', () => {
      this.handleDisconnection(clientId);
    });

    // Handle errors
    ws.on('error', (error) => {
      this.logger.error(`WebSocket error for ${clientId}:`, error);
      this.handleDisconnection(clientId);
    });

    // Send welcome message
    this.sendMessage(clientId, {
      type: 'auth',
      timestamp: new Date().toISOString(),
      data: {
        message: 'Connected to Playground WebSocket. Please authenticate.',
        clientId,
      },
    });
  }

  private handleMessage(clientId: string, message: IPlaygroundWebSocketMessage): void {
    const client = this.clients.get(clientId);
    if (!client) return;

    client.lastActivity = new Date();

    switch (message.type) {
      case 'auth':
        this.handleAuthentication(clientId, message);
        break;

      case 'ping':
        this.sendMessage(clientId, {
          type: 'pong',
          timestamp: new Date().toISOString(),
        });
        break;

      case 'playground_update':
        if (client.isAuthenticated) {
          if (!client.sessionId) {
            this.sendError(clientId, 'Session not established');
            return;
          }
          this.broadcastToSession(client.sessionId, message, clientId);
        } else {
          this.sendError(clientId, 'Authentication required');
        }
        break;

      default:
        this.sendError(clientId, `Unknown message type: ${message.type}`);
    }
  }

  private handleAuthentication(clientId: string, message: IPlaygroundWebSocketMessage): void {
    const client = this.clients.get(clientId);
    if (!client) return;

    if (!isPlaygroundAuthPayload(message.data)) {
      this.sendError(clientId, 'Invalid auth payload (missing userId/sessionId/token)');
      return;
    }
    const { userId, sessionId, token } = message.data;

    if (!userId || !sessionId) {
      this.sendError(clientId, 'Missing userId or sessionId');
      return;
    }

    if (!token) {
      this.sendError(clientId, 'Missing authentication token');
      client.ws.close();
      return;
    }

    const jwtSecret = process.env.JWT_SECRET;
    if (jwtSecret) {
      try {
        jwt.verify(token, jwtSecret);
      } catch {
        this.sendError(clientId, 'Invalid or expired authentication token');
        client.ws.close();
        return;
      }
    } else {
      // Development fallback: verify JWT format (3 dot-separated parts)
      const parts = token.split('.');
      if (parts.length !== JWT_PART_COUNT) {
        this.sendError(clientId, 'Invalid or expired authentication token');
        client.ws.close();
        return;
      }
    }

    // Update client info
    client.userId = userId;
    client.sessionId = sessionId;
    client.isAuthenticated = true;

    // Track user sessions
    let sessionSet = this.userSessions.get(userId);
    if (!sessionSet) {
      sessionSet = new Set();
      this.userSessions.set(userId, sessionSet);
    }
    sessionSet.add(clientId);

    // Send authentication success
    this.sendMessage(clientId, {
      type: 'auth',
      timestamp: new Date().toISOString(),
      data: {
        success: true,
        message: 'Authentication successful',
        userId,
        sessionId,
        clientId,
      },
    });

    this.logger.info(`Client ${clientId} authenticated as user ${userId}, session ${sessionId}`);
  }

  private handleDisconnection(clientId: string): void {
    const client = this.clients.get(clientId);
    if (client) {
      // Remove from user sessions
      if (client.userId) {
        const userSessions = this.userSessions.get(client.userId);
        if (userSessions) {
          userSessions.delete(clientId);
          if (userSessions.size === 0) {
            this.userSessions.delete(client.userId);
          }
        }
      }

      this.clients.delete(clientId);
      this.logger.info(`Client ${clientId} disconnected`);
    }
  }

  private sendMessage(clientId: string, message: IPlaygroundWebSocketMessage): void {
    const client = this.clients.get(clientId);
    if (client && client.ws.readyState === WebSocket.OPEN) {
      try {
        client.ws.send(JSON.stringify(message));
      } catch (error) {
        this.logger.error(`Failed to send message to ${clientId}:`, error as Error);
        this.handleDisconnection(clientId);
      }
    }
  }

  private sendError(clientId: string, error: string): void {
    this.sendMessage(clientId, {
      type: 'error',
      timestamp: new Date().toISOString(),
      data: {
        success: false,
        error,
      },
    });
  }

  private broadcastToSession(
    sessionId: string,
    message: IPlaygroundWebSocketMessage,
    excludeClientId?: string,
  ): void {
    let broadcastCount = 0;

    for (const [clientId, client] of this.clients) {
      if (
        client.sessionId === sessionId &&
        client.isAuthenticated &&
        clientId !== excludeClientId &&
        client.ws.readyState === WebSocket.OPEN
      ) {
        this.sendMessage(clientId, message);
        broadcastCount++;
      }
    }

    this.logger.debug(`Broadcasted to ${broadcastCount} clients in session ${sessionId}`);
  }

  private cleanupInactiveConnections(): void {
    const now = new Date();
    const timeout = 5 * 60 * 1000; // 5 minutes

    for (const [clientId, client] of this.clients) {
      if (now.getTime() - client.lastActivity.getTime() > timeout) {
        this.logger.debug(`Cleaning up inactive client: ${clientId}`);
        client.ws.close();
        this.handleDisconnection(clientId);
      }
    }
  }

  private generateClientId(): string {
    return `client_${Date.now()}_${Math.random().toString(36).slice(2, 11)}`;
  }

  // Public API for external integration

  /**
   * Broadcast a message to all clients of a specific user
   */
  public broadcastToUser(
    userId: string,
    message: Omit<IPlaygroundWebSocketMessage, 'timestamp'>,
  ): void {
    const userSessions = this.userSessions.get(userId);
    if (!userSessions) return;

    const messageWithTimestamp: IPlaygroundWebSocketMessage = {
      ...message,
      type: message.type as TPlaygroundWebSocketMessageKind,
      timestamp: new Date().toISOString(),
    };

    for (const clientId of userSessions) {
      this.sendMessage(clientId, messageWithTimestamp);
    }
  }

  /**
   * Get connection statistics
   */
  public getStats(): {
    totalConnections: number;
    authenticatedConnections: number;
    uniqueUsers: number;
    uniqueSessions: number;
  } {
    const authenticatedClients = Array.from(this.clients.values()).filter((c) => c.isAuthenticated);
    const uniqueSessions = new Set(authenticatedClients.map((c) => c.sessionId).filter(Boolean));

    return {
      totalConnections: this.clients.size,
      authenticatedConnections: authenticatedClients.length,
      uniqueUsers: this.userSessions.size,
      uniqueSessions: uniqueSessions.size,
    };
  }

  /**
   * Close all connections and cleanup
   */
  public close(): void {
    clearInterval(this.cleanupInterval);
    for (const client of this.clients.values()) {
      client.ws.close();
    }
    this.clients.clear();
    this.userSessions.clear();
    this.wss.close();
    this.logger.info('PlaygroundWebSocketServer closed');
  }
}
