Embalming_Girl_Server/handler/ws.go

package handler

import (
	"Embalming_Girl_Server/db"
	"Embalming_Girl_Server/game"
	"Embalming_Girl_Server/room"
	"encoding/json"
	"log"
	"math/rand"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool { return true },
}

type Client struct {
	ID       string
	RoomID   string
	UserID   int64
	Username string
	Nickname string
	Token    string
	Conn     *websocket.Conn
	Send     chan []byte
	mu       sync.Mutex
}

type Hub struct {
	RoomMgr *room.Manager
	clients map[string]*Client
	engines map[string]*game.Engine
	mu      sync.RWMutex
}

func NewHub() *Hub {
	return &Hub{
		RoomMgr: room.NewManager(),
		clients: make(map[string]*Client),
		engines: make(map[string]*game.Engine),
	}
}

func (h *Hub) OnlineCount() int {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return len(h.clients)
}

func (h *Hub) InitTestRooms() {
	for _, n := range []int{3, 4, 5, 6} {
		rid := "TEST" + string(rune('0'+n))
		h.RoomMgr.CreateTestRoom(rid, n)
		log.Printf("Test room: %s (%d players, %d bots)", rid, n, n-1)
	}
}

func (h *Hub) ServeWS(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println("Upgrade:", err)
		return
	}
	clientID := generateClientID()
	client := &Client{ID: clientID, Conn: conn, Send: make(chan []byte, 256)}
	h.mu.Lock()
	h.clients[clientID] = client
	h.mu.Unlock()
	log.Printf("Connected: %s", clientID)
	go h.writePump(client)
	go h.readPump(client)
}

func (h *Hub) readPump(c *Client) {
	defer func() { h.removeClient(c); c.Conn.Close() }()
	for {
		msgType, msg, err := c.Conn.ReadMessage()
		if err != nil {
			break
		}
		if msgType == websocket.BinaryMessage {
			h.relayAudio(c, msg)
			continue
		}
		var m map[string]interface{}
		if json.Unmarshal(msg, &m) == nil {
			h.handleMessage(c, m)
		}
	}
}

func (h *Hub) writePump(c *Client) {
	defer c.Conn.Close()
	for data := range c.Send {
		c.mu.Lock()
		err := c.Conn.WriteMessage(websocket.TextMessage, data)
		c.mu.Unlock()
		if err != nil {
			break
		}
	}
}

func (h *Hub) removeClient(c *Client) {
	h.mu.Lock()
	delete(h.clients, c.ID)
	h.mu.Unlock()
	if c.RoomID != "" {
		h.RoomMgr.LeaveRoom(c.RoomID, c.ID)
		rm := h.RoomMgr.GetRoom(c.RoomID)
		if rm != nil {
			h.broadcastRoomState(rm)
		}
	}
	close(c.Send)
	log.Printf("Disconnected: %s", c.ID)
}

func (h *Hub) handleMessage(c *Client, msg map[string]interface{}) {
	t, _ := msg["type"].(string)
	p, _ := msg["payload"].(map[string]interface{})
	switch t {
	case "register":
		h.handleRegister(c, p)
	case "login":
		h.handleLogin(c, p)
	case "update_profile":
		h.handleUpdateProfile(c, p)
	case "get_profile":
		h.handleGetProfile(c, p)
	case "create_room":
		h.handleCreateRoom(c, p)
	case "join_room":
		h.handleJoinRoom(c, p)
	case "leave_room":
		h.handleLeaveRoom(c)
	case "set_ready":
		h.handleSetReady(c, p)
	case "start_game":
		h.handleStartGame(c)
	case "list_rooms":
		h.handleListRooms(c)
	case "select_card":
		h.handleSelectCard(c, p)
	case "confirm_action":
		h.handleConfirmAction(c, p)
	case "effect_response":
		h.handleEffectResponse(c, p)
	case "send_chat":
		h.handleChat(c, p)
	case "set_chaos":
		h.handleSetChaos(c, p)
	case "heartbeat":
		h.sendTo(c, "heartbeat_ack", nil)
	}
}

// ---- Room handlers ----

func (h *Hub) handleListRooms(c *Client) {
	h.sendTo(c, "room_list", map[string]interface{}{"rooms": h.RoomMgr.ListRooms()})
}

func (h *Hub) handleRegister(c *Client, p map[string]interface{}) {
	username, _ := p["username"].(string)
	password, _ := p["password"].(string)
	nickname, _ := p["nickname"].(string)
	user, err := db.Register(username, password, nickname)
	if err != nil {
		h.sendTo(c, "register_result", map[string]interface{}{"success": false, "message": err.Error()})
		return
	}
	h.sendTo(c, "register_result", map[string]interface{}{
		"success": true, "message": "注册成功", "username": user.Username, "nickname": user.Nickname,
	})
}

func (h *Hub) handleLogin(c *Client, p map[string]interface{}) {
	username, _ := p["username"].(string)
	password, _ := p["password"].(string)
	token, _ := p["token"].(string)

	var user *db.User
	var err error
	if token != "" {
		user, err = db.ValidateToken(token)
		if err == nil {
			token = token
		}
	} else {
		token, user, err = db.Login(username, password)
	}
	if err != nil {
		h.sendTo(c, "login_result", map[string]interface{}{"success": false, "message": err.Error()})
		return
	}

	c.UserID = user.ID
	c.Username = user.Username
	c.Nickname = user.Nickname
	c.Token = token
	log.Printf("User logged in: %s (%s)", user.Username, c.ID)

	h.sendTo(c, "login_result", map[string]interface{}{
		"success": true, "token": token,
		"username": user.Username, "nickname": user.Nickname, "avatar": user.Avatar, "bio": user.Bio,
	})
}

func (h *Hub) handleUpdateProfile(c *Client, p map[string]interface{}) {
	if c.UserID == 0 {
		h.sendTo(c, "error", map[string]interface{}{"message": "未登录"})
		return
	}
	nick, _ := p["nickname"].(string)
	bio, _ := p["bio"].(string)
	avatar, _ := p["avatar"].(string)
	if nick == "" {
		nick = c.Nickname
	}
	db.UpdateProfile(c.UserID, nick, bio, avatar)
	c.Nickname = nick
	h.sendTo(c, "profile_updated", map[string]interface{}{"success": true, "nickname": nick, "bio": bio, "avatar": avatar})
}

func (h *Hub) handleGetProfile(c *Client, p map[string]interface{}) {
	username, _ := p["username"].(string)
	user, err := db.GetProfile(username)
	if err != nil {
		h.sendTo(c, "profile_info", map[string]interface{}{"success": false})
		return
	}
	h.sendTo(c, "profile_info", map[string]interface{}{
		"success": true, "username": user.Username, "nickname": user.Nickname, "bio": user.Bio, "avatar": user.Avatar,
	})
}

func (h *Hub) handleCreateRoom(c *Client, p map[string]interface{}) {
	nick, _ := p["nickname"].(string)
	mp := 4
	if v, ok := p["max_players"].(float64); ok {
		mp = int(v)
	}
	if nick == "" {
		nick = "Player"
	}
	rm := h.RoomMgr.CreateRoom(c.ID, nick, mp)
	c.RoomID = rm.ID
	rm.Players[c.ID].Send = c.Send
	h.sendTo(c, "room_created", map[string]interface{}{"room_id": rm.ID})
	h.broadcastRoomState(rm)
}

func (h *Hub) handleJoinRoom(c *Client, p map[string]interface{}) {
	rid, _ := p["room_id"].(string)
	nick, _ := p["nickname"].(string)
	if nick == "" {
		nick = "Player"
	}
	rm, ok := h.RoomMgr.JoinRoom(rid, c.ID, nick)
	if !ok {
		h.sendTo(c, "error", map[string]interface{}{"message": "Cannot join room"})
		return
	}
	c.RoomID = rm.ID
	rm.Players[c.ID].Send = c.Send
	h.sendTo(c, "room_joined", map[string]interface{}{"room_id": rm.ID})
	h.broadcastRoomState(rm)
}

func (h *Hub) handleLeaveRoom(c *Client) {
	if c.RoomID == "" {
		return
	}
	rid := c.RoomID
	h.RoomMgr.LeaveRoom(rid, c.ID)
	c.RoomID = ""
	h.sendTo(c, "room_left", nil)
	if rm := h.RoomMgr.GetRoom(rid); rm != nil {
		h.broadcastRoomState(rm)
	}
}

func (h *Hub) handleSetChaos(c *Client, p map[string]interface{}) {
	if c.RoomID == "" {
		return
	}
	rm := h.RoomMgr.GetRoom(c.RoomID)
	if rm == nil {
		return
	}
	enabled, _ := p["enabled"].(bool)
	rm.ChaosMode = enabled
	h.broadcastRoomState(rm)
}

func (h *Hub) handleChat(c *Client, p map[string]interface{}) {
	if c.RoomID == "" {
		return
	}
	rm := h.RoomMgr.GetRoom(c.RoomID)
	if rm == nil {
		return
	}
	text, _ := p["text"].(string)
	if text == "" {
		return
	}
	nick := ""
	if pl, ok := rm.Players[c.ID]; ok {
		nick = pl.Nickname
	}
	msg := map[string]interface{}{
		"sender_id": c.ID, "nickname": nick, "text": text,
	}
	for _, pid := range rm.PlayerOrder {
		pl := rm.Players[pid]
		if pl.IsBot || pl.Send == nil {
			continue
		}
		h.sendToChannel(pl.Send, "chat", msg)
	}
}

func (h *Hub) handleSetReady(c *Client, p map[string]interface{}) {
	ready, _ := p["ready"].(bool)
	if c.RoomID == "" {
		return
	}
	h.RoomMgr.SetReady(c.RoomID, c.ID, ready)
	if rm := h.RoomMgr.GetRoom(c.RoomID); rm != nil {
		h.broadcastRoomState(rm)
	}
}

// ---- Game start ----

func (h *Hub) handleStartGame(c *Client) {
	if c.RoomID == "" {
		return
	}
	rm := h.RoomMgr.GetRoom(c.RoomID)
	if rm == nil || !rm.Players[c.ID].IsHost {
		h.sendTo(c, "error", map[string]interface{}{"message": "Only host can start"})
		return
	}
	if len(rm.Players) < 3 {
		h.sendTo(c, "error", map[string]interface{}{"message": "Need 3+ players"})
		return
	}

	var ids, nicks []string
	var bots []bool
	for _, pid := range rm.PlayerOrder {
		p := rm.Players[pid]
		ids = append(ids, p.ID)
		nicks = append(nicks, p.Nickname)
		bots = append(bots, p.IsBot)
	}

	eng := game.NewEngine(ids, nicks, bots, len(ids))
	eng.ChaosMode = rm.ChaosMode
	eng.Deal()

	h.mu.Lock()
	h.engines[rm.ID] = eng
	h.mu.Unlock()

	rm.State = "playing"
	rm.HarmonyTarget = eng.HarmonyTarget

	for _, pid := range rm.PlayerOrder {
		p := rm.Players[pid]
		if p.IsBot || p.Send == nil {
			continue
		}
		h.sendToChannel(p.Send, "game_start", map[string]interface{}{
			"your_player_id": pid,
			"harmony_target": eng.HarmonyTarget,
			"player_count":   len(ids),
		})
	}

	log.Printf("Game started in %s, %d players, target=%d, first=%s",
		rm.ID, len(ids), eng.HarmonyTarget, eng.CurrentPlayer().ID)

	h.broadcastSnapshots(rm.ID)
	h.processBotTurns(rm.ID)
}

// ---- Game action handlers ----

func (h *Hub) handleSelectCard(c *Client, p map[string]interface{}) {
	h.mu.RLock()
	eng := h.engines[c.RoomID]
	h.mu.RUnlock()
	if eng == nil {
		return
	}

	uid, _ := p["card_uid"].(string)
	if err := eng.SelectCard(c.ID, uid); err != nil {
		h.sendTo(c, "error", map[string]interface{}{"message": err.Error()})
		return
	}
	h.broadcastSnapshots(c.RoomID)
}

func (h *Hub) handleConfirmAction(c *Client, p map[string]interface{}) {
	h.mu.RLock()
	eng := h.engines[c.RoomID]
	h.mu.RUnlock()
	if eng == nil {
		return
	}

	action, _ := p["action"].(string)
	target, _ := p["target_player_id"].(string)

	if err := eng.ConfirmAction(c.ID, action, target); err != nil {
		h.sendTo(c, "error", map[string]interface{}{"message": err.Error()})
		return
	}

	h.broadcastSnapshots(c.RoomID)
	h.flushGameLogs(c.RoomID)

	if eng.IsEffectResolving() {
		h.processBotEffectResponses(c.RoomID)
		return
	}

	if eng.Phase == "all_exited" {
		h.handleAllExited(c.RoomID)
		return
	}

	h.processBotTurns(c.RoomID)
}

func (h *Hub) handleEffectResponse(c *Client, p map[string]interface{}) {
	h.mu.RLock()
	eng := h.engines[c.RoomID]
	h.mu.RUnlock()
	if eng == nil || !eng.IsEffectResolving() {
		return
	}

	if err := eng.HandleEffectResponse(c.ID, p); err != nil {
		h.sendTo(c, "error", map[string]interface{}{"message": err.Error()})
		return
	}

	h.broadcastSnapshots(c.RoomID)

	if eng.IsEffectResolving() {
		h.processBotEffectResponses(c.RoomID)
		return
	}

	if eng.Phase == "all_exited" {
		h.handleAllExited(c.RoomID)
		return
	}

	h.processBotTurns(c.RoomID)
}

func (h *Hub) processBotEffectResponses(roomID string) {
	h.mu.RLock()
	eng := h.engines[roomID]
	h.mu.RUnlock()
	if eng == nil {
		return
	}

	for eng.IsEffectResolving() {
		ef := eng.Effect
		if ef == nil {
			break
		}

		if ef.ChoiceType == "all_select_card" {
			anyBotLeft := false
			for _, p := range eng.Players {
				if p.IsBot && eng.EffectNeedsResponseFrom(p.ID) {
					time.Sleep(200 * time.Millisecond)
					resp := eng.BotAutoEffectResponse(p.ID)
					if resp != nil {
						eng.HandleEffectResponse(p.ID, resp)
					}
					anyBotLeft = true
				}
			}
			if !anyBotLeft {
				break
			}
			h.broadcastSnapshots(roomID)
			continue
		}

		waitFor := ef.WaitingFor
		bot := eng.GetPlayer(waitFor)
		if bot == nil || !bot.IsBot {
			break
		}

		time.Sleep(200 * time.Millisecond)
		resp := eng.BotAutoEffectResponse(waitFor)
		if resp != nil {
			eng.HandleEffectResponse(waitFor, resp)
		}
		h.broadcastSnapshots(roomID)
	}
}

func (h *Hub) processBotTurns(roomID string) {
	h.mu.RLock()
	eng := h.engines[roomID]
	h.mu.RUnlock()
	if eng == nil {
		return
	}

	for eng.Phase == "playing" && !eng.IsEffectResolving() {
		cp := eng.CurrentPlayer()
		if cp == nil || !cp.IsBot {
			break
		}

		time.Sleep(300 * time.Millisecond)

		action, target := eng.BotAutoPlay(cp.ID)
		if action == "" {
			break
		}

		eng.ConfirmAction(cp.ID, action, target)
		h.broadcastSnapshots(roomID)
		h.flushGameLogs(roomID)

		if eng.IsEffectResolving() {
			h.processBotEffectResponses(roomID)
			if eng.IsEffectResolving() {
				break
			}
			h.broadcastSnapshots(roomID)
		}

		if eng.Phase == "all_exited" {
			h.handleAllExited(roomID)
			return
		}
	}
}

func (h *Hub) handleAllExited(roomID string) {
	rm := h.RoomMgr.GetRoom(roomID)
	if rm == nil {
		return
	}
	h.mu.RLock()
	eng := h.engines[roomID]
	h.mu.RUnlock()
	if eng == nil {
		return
	}

	result := eng.RunSettlement()
	log.Printf("Settlement in room %s: %s", roomID, result["victory"].(map[string]interface{})["end_type"])

	for _, pid := range rm.PlayerOrder {
		p := rm.Players[pid]
		if p.IsBot || p.Send == nil {
			continue
		}
		h.sendToChannel(p.Send, "settlement_result", result)
	}

	endType, _ := result["victory"].(map[string]interface{})["end_type"].(string)
	winners, _ := result["victory"].(map[string]interface{})["winners"].([]map[string]interface{})
	var winnerNames []string
	for _, w := range winners {
		winnerNames = append(winnerNames, w["nickname"].(string))
	}

	endMsg := map[string]interface{}{"end_type": endType}
	if endType == "all_dead" {
		endMsg["message"] = "全灭结局 — 没有任何人获胜"
	} else {
		msg := ""
		for i, n := range winnerNames {
			if i > 0 {
				msg += ", "
			}
			msg += n
		}
		endMsg["message"] = "获胜者: " + msg
		endMsg["winners"] = winners
	}

	for _, pid := range rm.PlayerOrder {
		p := rm.Players[pid]
		if p.IsBot || p.Send == nil {
			continue
		}
		h.sendToChannel(p.Send, "game_end", endMsg)
	}

	rm.State = "waiting"
	for _, p := range rm.Players {
		p.IsReady = false
		if p.IsBot {
			p.IsReady = true
		}
	}
	h.mu.Lock()
	delete(h.engines, roomID)
	h.mu.Unlock()
	h.broadcastRoomState(rm)
}

// ---- Audio relay ----

func (h *Hub) flushGameLogs(roomID string) {
	h.mu.RLock()
	eng := h.engines[roomID]
	h.mu.RUnlock()
	if eng == nil {
		return
	}
	rm := h.RoomMgr.GetRoom(roomID)
	if rm == nil {
		return
	}
	for _, msg := range eng.FlushLogs() {
		for _, pid := range rm.PlayerOrder {
			p := rm.Players[pid]
			if p.IsBot || p.Send == nil {
				continue
			}
			h.sendToChannel(p.Send, "game_log", map[string]interface{}{"text": msg})
		}
	}
}

func (h *Hub) relayAudio(sender *Client, pcm []byte) {
	if sender.RoomID == "" {
		return
	}
	rm := h.RoomMgr.GetRoom(sender.RoomID)
	if rm == nil {
		return
	}

	idBytes := []byte(sender.ID)
	header := make([]byte, 2+len(idBytes))
	header[0] = byte(len(idBytes) >> 8)
	header[1] = byte(len(idBytes))
	copy(header[2:], idBytes)
	frame := append(header, pcm...)

	h.mu.RLock()
	defer h.mu.RUnlock()
	for _, pid := range rm.PlayerOrder {
		p := rm.Players[pid]
		if p.IsBot || pid == sender.ID {
			continue
		}
		if cl, ok := h.clients[pid]; ok {
			cl.mu.Lock()
			cl.Conn.WriteMessage(websocket.BinaryMessage, frame)
			cl.mu.Unlock()
		}
	}
}

// ---- Broadcast helpers ----

func (h *Hub) broadcastSnapshots(roomID string) {
	rm := h.RoomMgr.GetRoom(roomID)
	if rm == nil {
		return
	}
	h.mu.RLock()
	eng := h.engines[roomID]
	h.mu.RUnlock()
	if eng == nil {
		return
	}

	for _, pid := range rm.PlayerOrder {
		p := rm.Players[pid]
		if p.IsBot || p.Send == nil {
			continue
		}
		snap := eng.Snapshot(pid)
		data := buildMessage("state_snapshot", snap)
		select {
		case p.Send <- data:
		default:
		}
	}
}

func (h *Hub) broadcastRoomState(rm *room.Room) {
	players := rm.GetPlayerInfos()
	for _, pid := range rm.PlayerOrder {
		p, ok := rm.Players[pid]
		if !ok || p.IsBot || p.Send == nil {
			continue
		}
		st := map[string]interface{}{
			"room_id": rm.ID, "state": rm.State,
			"harmony_target": rm.HarmonyTarget,
			"players":        players, "your_player_id": pid,
		}
		select {
		case p.Send <- buildMessage("room_state", st):
		default:
		}
	}
}

func (h *Hub) sendTo(c *Client, t string, p map[string]interface{}) {
	select {
	case c.Send <- buildMessage(t, p):
	default:
	}
}

func (h *Hub) sendToChannel(ch chan []byte, t string, p map[string]interface{}) {
	select {
	case ch <- buildMessage(t, p):
	default:
	}
}

func buildMessage(t string, p map[string]interface{}) []byte {
	m := map[string]interface{}{"type": t}
	if p != nil {
		m["payload"] = p
	}
	d, _ := json.Marshal(m)
	return d
}

func generateClientID() string {
	const cs = "abcdefghijklmnopqrstuvwxyz0123456789"
	b := make([]byte, 12)
	for i := range b {
		b[i] = cs[rand.Intn(len(cs))]
	}
	return string(b)
}