ruby-sea-battle/game-server/internal/game/manager.go

// Package game keeps the in-memory state of all matches running on this
// game server: player connections, per-player unacknowledged message queues
// and a registry of recently finished matches.
package game

import (
	"context"
	"errors"
	"sync"
	"time"

	"github.com/ruby-sea-battle/game-server/internal/auth"
	"github.com/ruby-sea-battle/game-server/internal/protocol"
)

// ErrMatchFinished is returned by Join when the player tries to join a match
// that has already ended. The stored END message (if any) is replayed to the
// connection before Join returns.
var (
	ErrMatchFinished      = errors.New("match already finished")
	ErrInvalidSeq         = errors.New("invalid seq")
	ErrMatchNotStarted    = errors.New("match not started")
	ErrUnexpectedState    = errors.New("unexpected state for placing ships")
	ErrShipsAlreadyPlaced = errors.New("ships already placed")
	ErrInvalidSlot        = errors.New("invalid slot index")
	ErrNotYourTurn        = errors.New("not your turn")
)

// BoardSize is the width and height of the Sea Battle board.
const BoardSize = protocol.BoardSize

// Board represents a 10x10 Sea Battle grid.
// 0 indicates water/empty cell.
// 1..10 indicates a ship (ship index 0..9 + 1).
type Board [BoardSize][BoardSize]int8

// NewBoard constructs a Board from validated ships.
func NewBoard(ships []protocol.Ship) Board {
	var b Board
	for i, ship := range ships {
		id := int8(i + 1)
		for _, cell := range ship.Cells() {
			b[cell.Y][cell.X] = id
		}
	}
	return b
}

// Server-side match states.
const (
	StatePlacingShips = 1
	StatePlayer1Ready = 2
	StatePlayer2Ready = 3
	StatePlayer1Turn  = 4
	StatePlayer2Turn  = 5
)

// playerStatesFor maps a server-side state to the respective states seen by player 1 and player 2.
func playerStatesFor(serverState int) (p1State, p2State int) {
	switch serverState {
	case StatePlacingShips:
		return protocol.StatePlacingShips, protocol.StatePlacingShips
	case StatePlayer1Ready:
		return protocol.StateWaitOpponentReady, protocol.StateOpponentReady
	case StatePlayer2Ready:
		return protocol.StateOpponentReady, protocol.StateWaitOpponentReady
	case StatePlayer1Turn:
		return protocol.StateYourTurn, protocol.StateWaitOpponentTurn
	case StatePlayer2Turn:
		return protocol.StateWaitOpponentTurn, protocol.StateYourTurn
	default:
		return 0, 0
	}
}

// Sender abstracts a player connection so the game logic stays testable and
// independent of the websocket implementation.
type Sender interface {
	// Send writes one message to the connection. Implementations must be
	// safe for concurrent use and must bound the write with a timeout.
	Send(ctx context.Context, data []byte) error
	// Close closes the connection with the given websocket close code.
	Close(code int, reason string) error
}

// queuedMessage is a state-changing message that has been sent to a player
// but not yet acknowledged with an ACK.
type queuedMessage struct {
	seq  int64
	data []byte
}

// PlayerSlot is the state of one of the two players of a match.
type PlayerSlot struct {
	UserID int64

	conn          Sender // nil while the player is disconnected
	seq           int64  // last issued sequence number
	unacked       []queuedMessage
	lastState     int  // last state sent to this player, 0 if none
	everConnected bool // true after the first successful join

	ships          []protocol.Ship
	board          Board
	shipsPlaced    bool
	struck         [BoardSize][BoardSize]bool
	shipsRemaining int
}

// Match is the state of a single game of sea battle.
type Match struct {
	ID        string
	gameToken string    // game token used when match was created
	tokenExp  time.Time // game token expiration; bounds the join deadline
	playerIDs [2]int64  // slot 0 = player1, slot 1 = player2 (internal only)
	slots     [2]*PlayerSlot

	mu        sync.Mutex
	started   bool // both players joined at least once
	ended     bool
	state     int // current global server-side game state
	joinTimer *time.Timer
}

func (m *Match) GameToken() string {
	return m.gameToken
}

func (m *Match) PlayerIDs() [2]int64 {
	return m.playerIDs
}

func (m *Match) State() int {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.state
}

// Ships returns a copy of the ships placed by the player in slotIdx,
// or nil if ships have not been placed yet.
func (m *Match) Ships(slotIdx int) []protocol.Ship {
	m.mu.Lock()
	defer m.mu.Unlock()
	if slotIdx < 0 || slotIdx >= len(m.slots) || !m.slots[slotIdx].shipsPlaced {
		return nil
	}
	res := make([]protocol.Ship, len(m.slots[slotIdx].ships))
	copy(res, m.slots[slotIdx].ships)
	return res
}

// Board returns the board of the player in slotIdx and whether ships have been placed.
func (m *Match) Board(slotIdx int) (Board, bool) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if slotIdx < 0 || slotIdx >= len(m.slots) {
		return Board{}, false
	}
	return m.slots[slotIdx].board, m.slots[slotIdx].shipsPlaced
}

// Struck returns a copy of the struck grid for the player in slotIdx.
func (m *Match) Struck(slotIdx int) ([BoardSize][BoardSize]bool, bool) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if slotIdx < 0 || slotIdx >= len(m.slots) {
		return [BoardSize][BoardSize]bool{}, false
	}
	return m.slots[slotIdx].struck, true
}

// ShipsRemaining returns the number of unsunk ships for the player in slotIdx.
func (m *Match) ShipsRemaining(slotIdx int) int {
	m.mu.Lock()
	defer m.mu.Unlock()
	if slotIdx < 0 || slotIdx >= len(m.slots) {
		return -1
	}
	return m.slots[slotIdx].shipsRemaining
}

func (m *Match) slotIndex(userID int64) int {
	switch userID {
	case m.playerIDs[0]:
		return 0
	case m.playerIDs[1]:
		return 1
	default:
		return -1
	}
}

// finishedMatch remembers how a match ended so that late join attempts with
// a still-valid game token get the same END message instead of creating a
// phantom match (and, later, a duplicate result report to the main server).
type finishedMatch struct {
	ends      map[int64][]byte // userID -> END message sent to that player
	expiresAt time.Time        // game token expiration time
	gameToken string
}

// Manager owns all matches of this server instance.
type Manager struct {
	joinTimeout time.Duration

	mu       sync.Mutex
	matches  map[string]*Match
	finished map[string]*finishedMatch
}

func NewManager(joinTimeout time.Duration) *Manager {
	return &Manager{
		joinTimeout: joinTimeout,
		matches:     make(map[string]*Match),
		finished:    make(map[string]*finishedMatch),
	}
}

// GetMatch returns the match with the given ID, or nil if not found.
func (m *Manager) GetMatch(matchID string) *Match {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.matches[matchID]
}

// Join attaches conn to the match identified by the game token claims.
// userID must already be validated against the claims by the caller.
//
// If the player was already connected, the old connection is closed and
// replaced. Unacknowledged messages are replayed to the new connection; if
// there is nothing to replay but the game is already running, the current
// state is re-sent so the client can restore its UI.
//
// Returns the match and the player's slot index (0 or 1).
func (m *Manager) Join(ctx context.Context, conn Sender, gc *auth.GameClaims, userID int64) (*Match, int, error) {
	now := time.Now()

	m.mu.Lock()
	m.purgeFinishedLocked(now)
	if fin, ok := m.finished[gc.MatchID]; ok {
		m.mu.Unlock()
		if data, ok := fin.ends[userID]; ok {
			_ = conn.Send(ctx, data)
		}
		_ = conn.Close(protocol.CloseCodeNormal, "match already finished")
		return nil, -1, ErrMatchFinished
	}
	match, ok := m.matches[gc.MatchID]
	if !ok {
		token := gc.RawToken
		if token == "" {
			token = gc.MatchID
		}
		match = &Match{
			ID:        gc.MatchID,
			gameToken: token,
			tokenExp:  gc.ExpiresAt.Time,
			playerIDs: [2]int64{gc.Player1ID, gc.Player2ID},
			slots: [2]*PlayerSlot{
				{UserID: gc.Player1ID},
				{UserID: gc.Player2ID},
			},
		}
		// The second player must join before the join timeout expires, but
		// never later than the game token itself expires — after that a
		// join is impossible anyway.
		deadline := now.Add(m.joinTimeout)
		if !match.tokenExp.IsZero() && match.tokenExp.Before(deadline) {
			deadline = match.tokenExp
		}
		match.joinTimer = time.AfterFunc(time.Until(deadline), func() {
			m.onJoinTimeout(match.ID)
		})
		m.matches[gc.MatchID] = match
	}

	// Acquire match.mu before releasing m.mu to eliminate the race with onJoinTimeout / EndMatch
	match.mu.Lock()
	m.mu.Unlock()
	defer match.mu.Unlock()

	if match.ended {
		// The match has ended; make sure the END message is sent before closing.
		var endData []byte
		idx := match.slotIndex(userID)
		if idx >= 0 && len(match.slots[idx].unacked) > 0 {
			endData = match.slots[idx].unacked[len(match.slots[idx].unacked)-1].data
		}
		if endData == nil {
			m.mu.Lock()
			if fin := m.finished[gc.MatchID]; fin != nil {
				endData = fin.ends[userID]
			}
			m.mu.Unlock()
		}
		if endData != nil {
			_ = conn.Send(ctx, endData)
		}
		_ = conn.Close(protocol.CloseCodeNormal, "match already finished")
		return nil, -1, ErrMatchFinished
	}

	idx := match.slotIndex(userID)
	if idx < 0 {
		// Should be impossible: the caller validated membership.
		_ = conn.Close(protocol.CloseCodeBadUserToken, "not a participant of this match")
		return nil, -1, errors.New("user is not a participant")
	}
	slot := match.slots[idx]

	if slot.conn != nil && slot.conn != conn {
		// Only one live connection per player: replace the old one.
		_ = slot.conn.Close(protocol.CloseCodeReplaced, "connection replaced by a new one")
	}
	slot.conn = conn
	slot.everConnected = true

	switch {
	case !match.started && match.slots[0].everConnected && match.slots[1].everConnected:
		// Both players joined for the first time: the game begins.
		match.started = true
		if match.joinTimer != nil {
			match.joinTimer.Stop()
		}
		m.setStateLocked(ctx, match, StatePlacingShips)
	case match.started:
		// Reconnect of a player while the game is running.
		m.replayLocked(ctx, slot)
	}
	// Otherwise: first player waiting for the opponent, nothing to send.

	return match, idx, nil
}

// Leave detaches conn from the match if it is still the current connection
// of the player's slot. Messages queued for the player survive the
// disconnect and will be replayed on rejoin.
func (m *Manager) Leave(conn Sender, match *Match, idx int) {
	match.mu.Lock()
	defer match.mu.Unlock()
	if slot := match.slots[idx]; slot.conn == conn {
		slot.conn = nil
	}
}

// Ack confirms delivery of all messages with seq <= the given sequence
// number for the given player.
func (m *Manager) Ack(match *Match, idx int, seq int64) error {
	match.mu.Lock()
	defer match.mu.Unlock()
	if idx < 0 || idx >= len(match.slots) {
		return errors.New("invalid slot index")
	}
	slot := match.slots[idx]
	if seq <= 0 || seq > slot.seq {
		return ErrInvalidSeq
	}
	kept := slot.unacked[:0]
	for _, qm := range slot.unacked {
		if qm.seq > seq {
			kept = append(kept, qm)
		}
	}
	slot.unacked = kept
	return nil
}

// PlaceShips validates and stores the player's ship placement, advancing to the next match state.
// The state order depends on which player completed the placement first:
// - If Player 1 is first: StatePlacingShips -> StatePlayer1Ready -> StatePlayer1Turn
// - If Player 2 is first: StatePlacingShips -> StatePlayer2Ready -> StatePlayer2Turn
func (m *Manager) PlaceShips(ctx context.Context, match *Match, slotIdx int, ships []protocol.Ship) error {
	if err := protocol.ValidateShips(ships); err != nil {
		return err
	}

	match.mu.Lock()
	defer match.mu.Unlock()

	if match.ended {
		return ErrMatchFinished
	}
	if !match.started {
		return ErrMatchNotStarted
	}
	if slotIdx < 0 || slotIdx >= len(match.slots) {
		return ErrInvalidSlot
	}

	slot := match.slots[slotIdx]
	if slot.shipsPlaced {
		return ErrShipsAlreadyPlaced
	}

	var nextState int
	switch match.state {
	case StatePlacingShips:
		if slotIdx == 0 {
			nextState = StatePlayer1Ready
		} else {
			nextState = StatePlayer2Ready
		}
	case StatePlayer1Ready:
		if slotIdx != 1 {
			return ErrUnexpectedState
		}
		nextState = StatePlayer1Turn
	case StatePlayer2Ready:
		if slotIdx != 0 {
			return ErrUnexpectedState
		}
		nextState = StatePlayer2Turn
	default:
		return ErrUnexpectedState
	}

	slot.ships = make([]protocol.Ship, len(ships))
	copy(slot.ships, ships)
	slot.board = NewBoard(ships)
	slot.shipsPlaced = true
	slot.shipsRemaining = len(ships)
	slot.struck = [BoardSize][BoardSize]bool{}

	m.setStateLocked(ctx, match, nextState)
	return nil
}

// Fire processes a shot from the player at slotIdx towards opponent's board (1 - slotIdx).
// If the player does not have the right of turn, the message is ignored and ErrNotYourTurn is returned.
// If the position is out of bounds or already struck, FireResultInvalidPosition is sent to the shooter.
// Otherwise:
// - Miss: marks cell as struck, sends FireResultMiss to shooter, turn passes to opponent.
// - Hit: marks cell as struck, sends FireResultHit to shooter, turn does not pass.
// - Sunk: marks ship and all surrounding cells (including diagonals) as struck, sends FireResultSunk to shooter.
//   If all opponent ships are destroyed, match finishes: END messages (victory/defeat) are sent to both players,
//   and connections are closed immediately.
//   Otherwise, turn does not pass.
func (m *Manager) Fire(ctx context.Context, match *Match, slotIdx int, x, y int) (int, error) {
	res, won, matchID, err := m.fireLocked(ctx, match, slotIdx, x, y)
	if won {
		// TODO: Report match result to main-server once endpoint is ready.
		m.finishGame(matchID, slotIdx)
	}
	return res, err
}

func (m *Manager) fireLocked(ctx context.Context, match *Match, slotIdx int, x, y int) (res int, won bool, matchID string, err error) {
	match.mu.Lock()
	defer match.mu.Unlock()

	if match.ended {
		return 0, false, "", ErrMatchFinished
	}
	if !match.started {
		return 0, false, "", ErrMatchNotStarted
	}
	if slotIdx < 0 || slotIdx >= len(match.slots) {
		return 0, false, "", ErrInvalidSlot
	}

	var isTurn bool
	switch match.state {
	case StatePlayer1Turn:
		isTurn = (slotIdx == 0)
	case StatePlayer2Turn:
		isTurn = (slotIdx == 1)
	}
	if !isTurn {
		return 0, false, "", ErrNotYourTurn
	}

	shooter := match.slots[slotIdx]
	target := match.slots[1-slotIdx]

	if x < 0 || x >= BoardSize || y < 0 || y >= BoardSize || target.struck[y][x] {
		m.queueFireResultLocked(ctx, shooter, protocol.FireResultInvalidPosition)
		return protocol.FireResultInvalidPosition, false, "", nil
	}

	shipID := target.board[y][x]
	if shipID == 0 {
		target.struck[y][x] = true
		m.queueFireResultLocked(ctx, shooter, protocol.FireResultMiss)

		var nextState int
		if slotIdx == 0 {
			nextState = StatePlayer2Turn
		} else {
			nextState = StatePlayer1Turn
		}
		m.setStateLocked(ctx, match, nextState)
		return protocol.FireResultMiss, false, "", nil
	}

	target.struck[y][x] = true
	shipIdx := int(shipID - 1)
	ship := target.ships[shipIdx]

	sunk := true
	for _, c := range ship.Cells() {
		if !target.struck[c.Y][c.X] {
			sunk = false
			break
		}
	}

	if !sunk {
		m.queueFireResultLocked(ctx, shooter, protocol.FireResultHit)
		return protocol.FireResultHit, false, "", nil
	}

	for _, c := range ship.Cells() {
		for dy := -1; dy <= 1; dy++ {
			for dx := -1; dx <= 1; dx++ {
				nx, ny := c.X+dx, c.Y+dy
				if nx >= 0 && nx < BoardSize && ny >= 0 && ny < BoardSize {
					target.struck[ny][nx] = true
				}
			}
		}
	}
	target.shipsRemaining--

	m.queueFireResultLocked(ctx, shooter, protocol.FireResultSunk)

	if target.shipsRemaining == 0 {
		match.state = 0
		return protocol.FireResultSunk, true, match.ID, nil
	}

	return protocol.FireResultSunk, false, "", nil
}

// SetState moves the match into a new server-side state and notifies both players.
func (m *Manager) SetState(ctx context.Context, match *Match, state int) {
	match.mu.Lock()
	defer match.mu.Unlock()
	m.setStateLocked(ctx, match, state)
}

// setStateLocked moves the match into a new server-side state and notifies both players
// with appropriate player-side states. match.mu must be held.
func (m *Manager) setStateLocked(ctx context.Context, match *Match, state int) {
	match.state = state
	p1State, p2State := playerStatesFor(state)
	if p1State != 0 {
		m.queueStateLocked(ctx, match.slots[0], p1State)
	}
	if p2State != 0 {
		m.queueStateLocked(ctx, match.slots[1], p2State)
	}
}

// queueStateLocked appends a STATE message to the player's unacknowledged
// queue and sends it right away if the player is connected.
// match.mu must be held.
func (m *Manager) queueStateLocked(ctx context.Context, slot *PlayerSlot, state int) {
	slot.seq++
	data := protocol.BuildState(slot.seq, state)
	slot.unacked = append(slot.unacked, queuedMessage{seq: slot.seq, data: data})
	slot.lastState = state
	if slot.conn != nil {
		_ = slot.conn.Send(ctx, data)
	}
}

// queueFireResultLocked appends a FIRE_RESULT message to the player's unacknowledged
// queue and sends it right away if the player is connected.
// match.mu must be held.
func (m *Manager) queueFireResultLocked(ctx context.Context, slot *PlayerSlot, result int) {
	slot.seq++
	data := protocol.BuildFireResult(slot.seq, result)
	slot.unacked = append(slot.unacked, queuedMessage{seq: slot.seq, data: data})
	if slot.conn != nil {
		_ = slot.conn.Send(ctx, data)
	}
}

// replayLocked re-sends everything the connected player has not
// acknowledged yet; if there is nothing pending, the current state is sent
// anew so a reconnected client can catch up.
// match.mu must be held.
func (m *Manager) replayLocked(ctx context.Context, slot *PlayerSlot) {
	if len(slot.unacked) == 0 {
		if slot.lastState != 0 {
			m.queueStateLocked(ctx, slot, slot.lastState)
		}
		return
	}
	for _, qm := range slot.unacked {
		_ = slot.conn.Send(ctx, qm.data)
	}
}

// EndMatch finishes a match with the given code and reason for both players.
func (m *Manager) EndMatch(matchID string, code int, reason string) {
	m.mu.Lock()
	match, ok := m.matches[matchID]
	if !ok {
		m.mu.Unlock()
		return
	}
	match.mu.Lock()
	if match.ended {
		match.mu.Unlock()
		m.mu.Unlock()
		return
	}
	m.finishMatchLocked(match, code, reason)
}

// onJoinTimeout finishes a match whose second player never joined.
func (m *Manager) onJoinTimeout(matchID string) {
	m.mu.Lock()
	match, ok := m.matches[matchID]
	if !ok {
		m.mu.Unlock()
		return
	}
	match.mu.Lock()
	if match.started || match.ended {
		match.mu.Unlock()
		m.mu.Unlock()
		return
	}
	m.finishMatchLocked(match, protocol.EndCodeJoinTimeout, "join timeout exceeded")
}

type playerEnd struct {
	code   int
	reason string
}

// finishMatchLocked tears down a match, removes it from m.matches, records it in m.finished,
// and sends identical END notifications to all players. Both m.mu and match.mu must be held
// upon entry; both are unlocked by this function.
func (m *Manager) finishMatchLocked(match *Match, code int, reason string) {
	m.finishMatchWithEndsLocked(match, [2]playerEnd{
		{code: code, reason: reason},
		{code: code, reason: reason},
	})
}

// finishMatchWithEndsLocked tears down a match, removes it from m.matches, records it in m.finished,
// and sends per-player END notifications to connected players. Both m.mu and match.mu must be held
// upon entry; both are unlocked by this function.
func (m *Manager) finishMatchWithEndsLocked(match *Match, endsPerSlot [2]playerEnd) {
	match.ended = true
	if match.joinTimer != nil {
		match.joinTimer.Stop()
	}
	delete(m.matches, match.ID)

	ends := make(map[int64][]byte, 2)
	type notification struct {
		conn Sender
		data []byte
	}
	var notify []notification
	for i, slot := range match.slots {
		slot.seq++
		end := endsPerSlot[i]
		data := protocol.BuildEnd(slot.seq, end.code, end.reason)
		slot.unacked = append(slot.unacked, queuedMessage{seq: slot.seq, data: data})
		ends[slot.UserID] = data
		if slot.conn != nil {
			notify = append(notify, notification{conn: slot.conn, data: data})
			slot.conn = nil
		}
	}
	match.mu.Unlock()

	// Remember outcome in finished registry. Bound expiresAt to at least now + 5m if tokenExp is in the past.
	now := time.Now()
	expiresAt := match.tokenExp
	if expiresAt.Before(now.Add(5 * time.Minute)) {
		expiresAt = now.Add(5 * time.Minute)
	}

	m.finished[match.ID] = &finishedMatch{
		ends:      ends,
		expiresAt: expiresAt,
		gameToken: match.gameToken,
	}
	m.purgeFinishedLocked(now)
	m.mu.Unlock()

	// Network I/O happens outside of any lock.
	for _, n := range notify {
		_ = n.conn.Send(context.Background(), n.data)
	}
	var wg sync.WaitGroup
	for _, n := range notify {
		wg.Add(1)
		go func(c Sender) {
			defer wg.Done()
			_ = c.Close(protocol.CloseCodeNormal, "match ended")
		}(n.conn)
	}
	wg.Wait()
}

// finishGame completes a match due to all opponent ships being sunk.
// winnerSlot is the slot index of the winning player (0 or 1).
func (m *Manager) finishGame(matchID string, winnerSlot int) {
	m.mu.Lock()
	match, ok := m.matches[matchID]
	if !ok {
		m.mu.Unlock()
		return
	}
	match.mu.Lock()
	if match.ended {
		match.mu.Unlock()
		m.mu.Unlock()
		return
	}

	endsPerSlot := [2]playerEnd{
		{code: protocol.EndCodeVictory, reason: "victory"},
		{code: protocol.EndCodeDefeat, reason: "defeat"},
	}
	if winnerSlot == 1 {
		endsPerSlot[0], endsPerSlot[1] = endsPerSlot[1], endsPerSlot[0]
	}

	m.finishMatchWithEndsLocked(match, endsPerSlot)
}

// Shutdown closes all active player connections across all matches with code 1001 (going away)
// and stops all pending join timers.
func (m *Manager) Shutdown() {
	m.mu.Lock()
	var conns []Sender
	for _, match := range m.matches {
		match.mu.Lock()
		if match.joinTimer != nil {
			match.joinTimer.Stop()
		}
		for _, slot := range match.slots {
			if slot.conn != nil {
				conns = append(conns, slot.conn)
				slot.conn = nil
			}
		}
		match.mu.Unlock()
	}
	m.mu.Unlock()

	var wg sync.WaitGroup
	for _, conn := range conns {
		wg.Add(1)
		go func(c Sender) {
			defer wg.Done()
			_ = c.Close(protocol.CloseCodeGoingAway, "going away")
		}(conn)
	}
	wg.Wait()
}

// purgeFinishedLocked drops registry entries whose game token has expired.
// Called lazily whenever the registry is touched. m.mu must be held.
func (m *Manager) purgeFinishedLocked(now time.Time) {
	for id, fin := range m.finished {
		if !now.Before(fin.expiresAt) {
			delete(m.finished, id)
		}
	}
}