import { Server as SocketIOServer, Socket } from "socket.io"

const courseSocketMap = new Map<string, Set<string>>() // slug -> Set<socket.id>

const notificationSocketMap = new Map<string, Set<string>>() // notificationId -> Set<socket.id>
const socketStore = new Map<string, Socket>() // socket.id -> socket

let io: SocketIOServer | null = null

export function initSocketIO(server: any) {
  io = new SocketIOServer(server, {
    cors: {
      origin: "*",
      methods: ["GET", "POST"],
    },
    path: "/sockete",
  })

  io.on("connection", socket => {
    console.log("🧠 Socket connected:", socket.id)
    socketStore.set(socket.id, socket)

    socket.on("register", (data: { courseSlug: string }) => {
      const { courseSlug } = data
      if (!courseSlug) return

      if (!courseSocketMap.has(courseSlug)) {
        courseSocketMap.set(courseSlug, new Set())
      }

      courseSocketMap.get(courseSlug)?.add(socket.id)
      console.log(`📦 Socket ${socket.id} registered to course: ${courseSlug}`)
    })

    socket.on("registerNotification", (data: { notificationId: string }) => {
      const { notificationId } = data
      if (!notificationId) return

      if (!notificationSocketMap.has(notificationId)) {
        notificationSocketMap.set(notificationId, new Set())
      }

      notificationSocketMap.get(notificationId)?.add(socket.id)
      console.log(
        `📧 Socket ${socket.id} registered to notification: ${notificationId}`
      )
    })

    socket.on("disconnect", () => {
      console.log("🔥 Socket disconnected:", socket.id)
      socketStore.delete(socket.id)

      for (const set of courseSocketMap.values()) {
        set.delete(socket.id)
      }
    })
  })

  return io
}

export function emitToCourse(courseSlug: string, event: string, payload: any) {
  const socketIds = courseSocketMap.get(courseSlug)
  if (!socketIds || socketIds.size === 0) return

  for (const id of socketIds) {
    const socket = socketStore.get(id)
    if (socket) socket.emit(event, payload)
  }
}

export function emitToNotification(
  notificationId: string,
  payload: any,
  retry = 0
) {
  const socketIds = notificationSocketMap.get(notificationId)
  if (!socketIds || socketIds.size === 0) {
    if (retry > 3) {
      console.log("❌ Notification", notificationId, "not found")
      return
    }

    setTimeout(() => {
      emitToNotification(notificationId, payload, retry + 1)
    }, 3000)
    return
  }

  for (const id of socketIds) {
    const socket = socketStore.get(id)
    if (socket) socket.emit(notificationId, payload)
  }
}

export function getSocketIO() {
  if (!io) throw new Error("Socket.IO not initialized")
  return io
}
