package main

import (
	"errors"
	"net"
	"sync"
)

type Hub struct {
	mu       sync.RWMutex
	Channels map[string]map[net.Conn]struct{}
}

func handleSubscribe(conn net.Conn, args []string) {
	if len(args) < 2 {
		conn.Write(serializeRESP(errors.New("wrong number of arguments for 'SUBSCRIBE'")))
		return
	}
	channels := args[1:]
	hub.mu.Lock()
	defer hub.mu.Unlock()
	for _, channel := range channels {
		if _, exists := hub.Channels[channel]; !exists {
			hub.Channels[channel] = make(map[net.Conn]struct{})
		}
		hub.Channels[channel][conn] = struct{}{}

		subCount := 0
		for _, subscribers := range hub.Channels {
			if _, ok := subscribers[conn]; ok {
				subCount++
			}
		}

		conn.Write(serializeRESP([]any{"subscribe", channel, subCount}))
	}
}

func handleUnsubscribe(conn net.Conn, args []string) {
	hub.mu.Lock()
	defer hub.mu.Unlock()

	var channels []string
	if len(args) < 2 {
		for channel, subscribers := range hub.Channels {
			if _, ok := subscribers[conn]; ok {
				channels = append(channels, channel)
			}
		}
	} else {
		channels = args[1:]
	}

	for _, channel := range channels {
		if subscribers, exists := hub.Channels[channel]; exists {
			delete(subscribers, conn)
			if len(subscribers) == 0 {
				delete(hub.Channels, channel)
			}
		}

		subCount := 0
		for _, subscribers := range hub.Channels {
			if _, ok := subscribers[conn]; ok {
				subCount++
			}
		}

		conn.Write(serializeRESP([]any{"unsubscribe", channel, subCount}))
	}
}

func handlePublish(conn net.Conn, args []string) {
	if len(args) != 3 {
		conn.Write(serializeRESP(errors.New("wrong number of arguments for 'PUBLISH'")))
		return
	}
	conn.Write(serializeRESP(hub.Publish(args[1], args[2])))
}

func handleDisconnect(conn net.Conn) {
	hub.mu.Lock()
	defer hub.mu.Unlock()
	for channel, subscribers := range hub.Channels {
		delete(subscribers, conn)
		if len(subscribers) == 0 {
			delete(hub.Channels, channel)
		}
	}
}

func (hub *Hub) Publish(channel, message string) int {
	hub.mu.RLock()
	subscribers, exists := hub.Channels[channel]
	if !exists || len(subscribers) == 0 {
		hub.mu.RUnlock()
		return 0
	}

	conns := make([]net.Conn, 0, len(subscribers))
	for conn := range subscribers {
		conns = append(conns, conn)
	}
	hub.mu.RUnlock()

	payload := serializeRESP([]string{"message", channel, message})
	count := 0
	for _, conn := range conns {
		_, err := conn.Write(payload)
		if err == nil {
			count++
		}
	}
	return count
}