// Package websocket 提供通用的 WebSocket Hub，支持按 session_id 分房间广播
package websocket

import (
	"encoding/json"
	"log"
	"sync"

	"github.com/gorilla/websocket"
)

// Hub 管理所有 WebSocket 连接和房间
// 每个 session_id 对应一个房间，房间内的连接互相广播消息
type Hub struct {
	mu         sync.RWMutex
	rooms      map[string]map[*Conn]struct{} // session_id → 连接集合
	register   chan *Conn
	unregister chan *Conn
}

// Conn 封装一条 WebSocket 连接
type Conn struct {
	Hub        *Hub
	SessionID  string
	Role       string // "A" 或 "B"，可选
	RoleMu     sync.Mutex
	WS         *websocket.Conn
	Send       chan []byte
	closeOnce  sync.Once
}

// NewHub 创建一个 Hub 并启动后台处理协程
func NewHub() *Hub {
	h := &Hub{
		rooms:      make(map[string]map[*Conn]struct{}),
		register:   make(chan *Conn),
		unregister: make(chan *Conn),
	}
	go h.run()
	return h
}

func (h *Hub) run() {
	for {
		select {
		case conn := <-h.register:
			h.mu.Lock()
			if h.rooms[conn.SessionID] == nil {
				h.rooms[conn.SessionID] = make(map[*Conn]struct{})
			}
			h.rooms[conn.SessionID][conn] = struct{}{}
			h.mu.Unlock()

		case conn := <-h.unregister:
			h.mu.Lock()
			if room, ok := h.rooms[conn.SessionID]; ok {
				delete(room, conn)
				if len(room) == 0 {
					delete(h.rooms, conn.SessionID)
				}
			}
			// 用 sync.Once 保证 Send channel 只被 close 一次
			conn.closeOnce.Do(func() { close(conn.Send) })
			h.mu.Unlock()
		}
	}
}

// Register 将连接加入指定 session 的房间
func (h *Hub) Register(conn *Conn) {
	h.register <- conn
}

// Unregister 将连接从房间移除（幂等，可重复调用）
func (h *Hub) Unregister(conn *Conn) {
	h.unregister <- conn
}

// Broadcast 向指定 session 房间内的所有连接发送消息
// msg 会被序列化为 JSON
func (h *Hub) Broadcast(sessionID string, msg interface{}) {
	data, err := json.Marshal(msg)
	if err != nil {
		log.Printf("[ws] marshal broadcast msg error: %v", err)
		return
	}
	h.mu.RLock()
	defer h.mu.RUnlock()
	if room, ok := h.rooms[sessionID]; ok {
		for conn := range room {
			select {
			case conn.Send <- data:
			default:
				// send channel 满了，丢弃，避免阻塞
				log.Printf("[ws] send channel full for session %s, dropping", sessionID)
			}
		}
	}
}

// RoomSize 返回指定 session 房间内的连接数（调试用）
func (h *Hub) RoomSize(sessionID string) int {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return len(h.rooms[sessionID])
}

// WritePump 从 Send channel 读消息并写入 WebSocket
// 在独立 goroutine 中运行，直到连接关闭
func WritePump(c *Conn) {
	defer func() {
		c.Hub.Unregister(c)
		c.WS.Close()
	}()
	for msg := range c.Send {
		if err := c.WS.WriteMessage(websocket.TextMessage, msg); err != nil {
			return
		}
	}
}

// ReadPump 从 WebSocket 读消息并丢弃
// 生产环境如果需要客户端主动发消息可以在这里处理
// 当前只负责检测断连
func ReadPump(c *Conn) {
	defer func() {
		c.Hub.Unregister(c)
		c.WS.Close()
	}()
	c.WS.SetPongHandler(func(string) error { return nil })
	for {
		if _, _, err := c.WS.ReadMessage(); err != nil {
			return
		}
	}
}