package websocket

import (
	"encoding/json"
	"sync"
	"time"

	"go.uber.org/zap"
)

type Message struct {
	Type    string      `json:"type"`
	Payload interface{} `json:"payload"`
}

type Client struct {
	hub      *Hub
	userID   uint
	tenantID string
	send     chan []byte
	done     chan struct{}
}

type Hub struct {
	mu       sync.RWMutex
	clients  map[uint]map[*Client]bool
	log      *zap.Logger
	register chan *Client
	unreg    chan *Client
}

func NewHub(log *zap.Logger) *Hub {
	h := &Hub{
		clients:  make(map[uint]map[*Client]bool),
		log:      log,
		register: make(chan *Client, 64),
		unreg:    make(chan *Client, 64),
	}
	go h.run()
	return h
}

func (h *Hub) run() {
	for {
		select {
		case c := <-h.register:
			h.mu.Lock()
			if h.clients[c.userID] == nil {
				h.clients[c.userID] = make(map[*Client]bool)
			}
			h.clients[c.userID][c] = true
			h.mu.Unlock()
			h.log.Debug("websocket client connected", zap.Uint("user_id", c.userID))

		case c := <-h.unreg:
			h.mu.Lock()
			if m := h.clients[c.userID]; m != nil {
				delete(m, c)
				if len(m) == 0 {
					delete(h.clients, c.userID)
				}
			}
			h.mu.Unlock()
			close(c.send)
			h.log.Debug("websocket client disconnected", zap.Uint("user_id", c.userID))
		}
	}
}

func (h *Hub) Register(c *Client) {
	h.register <- c
}

func (h *Hub) Unregister(c *Client) {
	h.unreg <- c
}

func (h *Hub) SendToUser(userID uint, msg Message) {
	data, err := json.Marshal(msg)
	if err != nil {
		h.log.Error("websocket marshal", zap.Error(err))
		return
	}
	h.mu.RLock()
	defer h.mu.RUnlock()
	for c := range h.clients[userID] {
		select {
		case c.send <- data:
		default:
			go func(c *Client) {
				h.unreg <- c
			}(c)
		}
	}
}

func (h *Hub) Broadcast(msg Message) {
	data, err := json.Marshal(msg)
	if err != nil {
		h.log.Error("websocket marshal", zap.Error(err))
		return
	}
	h.mu.RLock()
	defer h.mu.RUnlock()
	for _, clients := range h.clients {
		for c := range clients {
			select {
			case c.send <- data:
			default:
				go func(c *Client) {
					h.unreg <- c
				}(c)
			}
		}
	}
}

func (h *Hub) ConnectedCount() int {
	h.mu.RLock()
	defer h.mu.RUnlock()
	n := 0
	for _, clients := range h.clients {
		n += len(clients)
	}
	return n
}

func NewClient(hub *Hub, userID uint, tenantID string) *Client {
	return &Client{
		hub:      hub,
		userID:   userID,
		tenantID: tenantID,
		send:     make(chan []byte, 256),
		done:     make(chan struct{}),
	}
}

func (c *Client) Done() <-chan struct{} {
	return c.done
}

const (
	writeWait      = 10 * time.Second
	pongWait       = 60 * time.Second
	pingPeriod     = (pongWait * 9) / 10
	maxMessageSize = 4096
)