src/server/websocket.ts

import type { ServerWebSocket } from "bun";

interface Session {
  sessionId: string;
  sender: string;
  createdAt: number;
  ws?: ServerWebSocket;
}

interface ComicSession {
  sessionId: string;
  sender: string;
  password: string;
  data: any;
  createdAt: number;
  expiresAt: number;
  clientWs?: ServerWebSocket;
}

const sessions = new Map<string, Session>();
const comicSessions = new Map<string, ComicSession>();
const botConnections = new Set<ServerWebSocket>();
const clientConnections = new Map<string, ServerWebSocket>();

export function handleWebSocketConnection(ws: ServerWebSocket) {
  // @ts-ignore
  ws.data = {
    handleMessage: (message: any) => handleMessage(ws, message),
    cleanup: () => cleanup(ws),
  };

  ws.send(JSON.stringify({
    type: 'connected',
    message: 'Connected to server'
  }));
}

function handleMessage(ws: ServerWebSocket, message: any) {
  const { type, sessionId } = message;

  switch (type) {
    case 'bot_connect':
      botConnections.add(ws);
      console.log('[WebSocket] Bot connected, total bots:', botConnections.size);
      break;

    case 'create_session':
      sessions.set(sessionId, {
        sessionId,
        sender: message.sender,
        createdAt: Date.now(),
      });
      console.log('[WebSocket] Session created:', sessionId);
      break;
      
    case 'force_cleanup_session':
      handleForceCleanupSession(message);
      break;

    case 'user_connect':
      const session = sessions.get(sessionId);
      if (session) {
        session.ws = ws;
        ws.send(JSON.stringify({
          type: 'session_valid',
          sessionId
        }));
      } else {
        ws.send(JSON.stringify({
          type: 'session_invalid',
          message: 'Session not found or expired'
        }));
      }
      break;

    case 'registration_submit':
      handleRegistrationSubmit(sessionId, message.data);
      break;

    case 'comic_stream':
    case 'chapter_stream':
      handleComicStream(message);
      break;

    case 'chapter_fetched':
      handleChapterFetched(message);
      break;

    case 'request_comic_data':
      handleComicDataRequest(ws, message);
      break;

    case 'validate_comic_password':
      handlePasswordValidation(ws, message);
      break;

    case 'request_chapter_data':
      handleChapterDataRequest(ws, message);
      break;

    default:
      console.log('[WebSocket] Unknown message type:', type);
  }
}

function handleForceCleanupSession(message: any) {
  const { sessionId, sender, reason } = message;  
  const comicSession = comicSessions.get(sessionId);
  
  if (comicSession) {
    comicSessions.delete(sessionId);

    const clientWs = clientConnections.get(sessionId);
    if (clientWs) {
      try {
        clientWs.send(JSON.stringify({
          type: 'session_force_closed',
          reason: reason || 'Session removed by admin',
          message: 'Your session has been terminated'
        }));
      } catch (err) {
        console.error('[WebSocket] Error notifying client:', err);
      }
      clientConnections.delete(sessionId);
    }
    
    console.log(`[WebSocket] Force cleanup session: ${sessionId}, reason: ${reason}`);

    botConnections.forEach(botWs => {
      try {
        botWs.send(JSON.stringify({
          type: 'cleanup_confirmed',
          sessionId,
          sender
        }));
      } catch (error) {
        console.error('[WebSocket] Error sending cleanup confirmation:', error);
      }
    });
  } else {
    console.log(`[WebSocket] Session not found for cleanup: ${sessionId}`);
  }
}

function handleComicStream(message: any) {
  const { sessionId, sender, password, data, type } = message;
  const expiresAt = Date.now() + (24 * 60 * 60 * 1000);
  
  comicSessions.set(sessionId, {
    sessionId,
    sender,
    password,
    data,
    createdAt: Date.now(),
    expiresAt
  });

  console.log(`[WebSocket] Comic stream created: ${sessionId}, type: ${type}`);
}

function handleComicDataRequest(ws: ServerWebSocket, message: any) {
  const { sessionId, password } = message;
  
  const comicSession = comicSessions.get(sessionId);
  
  if (!comicSession) {
    ws.send(JSON.stringify({
      type: 'comic_data_response',
      success: false,
      error: 'Session not found or expired'
    }));
    return;
  }

  if (comicSession.password !== password) {
    ws.send(JSON.stringify({
      type: 'comic_data_response',
      success: false,
      error: 'Invalid password'
    }));
    return;
  }

  if (Date.now() > comicSession.expiresAt) {
    comicSessions.delete(sessionId);
    ws.send(JSON.stringify({
      type: 'comic_data_response',
      success: false,
      error: 'Session expired'
    }));
    return;
  }

  comicSession.clientWs = ws;
  clientConnections.set(sessionId, ws);

  ws.send(JSON.stringify({
    type: 'comic_data_response',
    success: true,
    data: comicSession.data,
    expiresAt: comicSession.expiresAt
  }));
}

function handlePasswordValidation(ws: ServerWebSocket, message: any) {
  const { sessionId } = message;
  
  const comicSession = comicSessions.get(sessionId);
  
  if (!comicSession) {
    ws.send(JSON.stringify({
      type: 'password_validation_response',
      success: false,
      error: 'Session not found or expired'
    }));
    return;
  }

  if (Date.now() > comicSession.expiresAt) {
    comicSessions.delete(sessionId);
    ws.send(JSON.stringify({
      type: 'password_validation_response',
      success: false,
      error: 'Session expired'
    }));
    return;
  }

  ws.send(JSON.stringify({
    type: 'password_validation_response',
    success: true,
    requiresPassword: true
  }));
}

function handleChapterDataRequest(ws: ServerWebSocket, message: any) {
  const { sessionId, password, chapterSlug, allChapters } = message;
  
  const comicSession = comicSessions.get(sessionId);
  
  if (!comicSession) {
    ws.send(JSON.stringify({
      type: 'chapter_data_response',
      success: false,
      error: 'Session not found or expired'
    }));
    return;
  }

  if (comicSession.password !== password) {
    ws.send(JSON.stringify({
      type: 'chapter_data_response',
      success: false,
      error: 'Invalid password'
    }));
    return;
  }

  if (Date.now() > comicSession.expiresAt) {
    comicSessions.delete(sessionId);
    ws.send(JSON.stringify({
      type: 'chapter_data_response',
      success: false,
      error: 'Session expired'
    }));
    return;
  }

  comicSession.clientWs = ws;
  clientConnections.set(sessionId, ws);

  botConnections.forEach(botWs => {
    try {
      botWs.send(JSON.stringify({
        type: 'fetch_chapter_request',
        sessionId,
        chapterSlug,
        sender: comicSession.sender,
        allChapters: allChapters || comicSession.data?.chapters || []
      }));
    } catch (error) {
      console.error('[WebSocket] Error sending to bot:', error);
    }
  });

  console.log(`[WebSocket] Chapter fetch requested: ${chapterSlug} for session: ${sessionId}`);
}

function handleChapterFetched(message: any) {
  const { sessionId, chapterData, success, error } = message;
  const clientWs = clientConnections.get(sessionId);
  
  if (!clientWs) {
    console.log(`[WebSocket] No client connection found for session: ${sessionId}`);
    return;
  }

  try {
    clientWs.send(JSON.stringify({
      type: 'chapter_data_response',
      success,
      data: chapterData,
      error: error || null
    }));
    console.log(`[WebSocket] Chapter data sent to client for session: ${sessionId}`);
  } catch (err) {
    console.error('[WebSocket] Error sending chapter data to client:', err);
  }
}

function handleRegistrationSubmit(sessionId: string, data: any) {
  const session = sessions.get(sessionId);
  if (!session) {
    return;
  }

  const { username, age, password, confirmPassword } = data;
  let success = true;
  let errorMessage = '';

  if (!username || username.trim().length < 3) {
    success = false;
    errorMessage = 'Username must be at least 3 characters long';
  } else if (!age || age < 13 || age > 100) {
    success = false;
    errorMessage = 'Age must be between 13 and 100';
  } else if (!password || password.length < 6 || password.length > 20) {
    success = false;
    errorMessage = 'Password must be 6-20 characters long';
  } else if (password !== confirmPassword) {
    success = false;
    errorMessage = 'Passwords do not match';
  }

  if (session.ws) {
    session.ws.send(JSON.stringify({
      type: 'registration_result',
      success,
      message: success ? 'Registration successful! Please check your WhatsApp.' : errorMessage
    }));
  }

  const registrationData = {
    type: 'registration_complete',
    sessionId,
    data: {
      success,
      message: errorMessage,
      userData: success ? {
        username: username.trim(),
        age: parseInt(age),
        password
      } : undefined
    }
  };

  botConnections.forEach(botWs => {
    try {
      botWs.send(JSON.stringify(registrationData));
    } catch (error) {
      console.error('[WebSocket] Error sending to bot:', error);
    }
  });

  setTimeout(() => {
    sessions.delete(sessionId);
  }, 5000);
}

function cleanup(ws: ServerWebSocket) {
  botConnections.delete(ws);

  for (const [_, session] of sessions.entries()) {
    if (session.ws === ws) {
      session.ws = undefined;
    }
  }

  for (const [sessionId, clientWs] of clientConnections.entries()) {
    if (clientWs === ws) {
      clientConnections.delete(sessionId);
    }
  }

  for (const [_, comicSession] of comicSessions.entries()) {
    if (comicSession.clientWs === ws) {
      comicSession.clientWs = undefined;
    }
  }
}

setInterval(() => {
  const now = Date.now();

  for (const [sessionId, session] of sessions.entries()) {
    if (now - session.createdAt > 600000) {
      sessions.delete(sessionId);
      console.log('[WebSocket] Cleaned up expired session:', sessionId);
    }
  }

  for (const [sessionId, comicSession] of comicSessions.entries()) {
    if (now > comicSession.expiresAt) {
      botConnections.forEach(botWs => {
        try {
          botWs.send(JSON.stringify({
            type: 'session_cleanup',
            sessionId,
            sender: comicSession.sender
          }));
        } catch (error) {
          console.error('[WebSocket] Error notifying bot:', error);
        }
      });
      
      comicSessions.delete(sessionId);
      clientConnections.delete(sessionId);
      console.log('[WebSocket] Cleaned up expired comic session:', sessionId);
    }
  }
}, 60 * 60 * 1000);