goonscape/network/network.go

package network

import (
	"bufio"
	"encoding/binary"
	"fmt"
	"io"
	"log"
	"net"
	"sync"
	"time"

	"gitea.boner.be/bdnugget/goonscape/types"
	pb "gitea.boner.be/bdnugget/goonserver/actions"
	rl "github.com/gen2brain/raylib-go/raylib"
	"google.golang.org/protobuf/proto"
)

const protoVersion = 1

var serverAddr = "boner.be:6969"       // Default server address
var lastSeenMessageTimestamp int64 = 0 // Track the last message timestamp seen by this client

func SetServerAddr(addr string) {
	serverAddr = addr
	log.Printf("Server address set to: %s", serverAddr)
}

// MessageHandler handles reading and writing protobuf messages
type MessageHandler struct {
	conn   net.Conn
	reader *bufio.Reader
}

// NewMessageHandler creates a new message handler
func NewMessageHandler(conn net.Conn) *MessageHandler {
	return &MessageHandler{
		conn:   conn,
		reader: bufio.NewReader(conn),
	}
}

// ReadMessage reads a single message from the network
func (mh *MessageHandler) ReadMessage() (*pb.ServerMessage, error) {
	// Read message length
	lengthBuf := make([]byte, 4)
	if _, err := io.ReadFull(mh.reader, lengthBuf); err != nil {
		return nil, fmt.Errorf("failed to read message length: %v", err)
	}

	messageLength := binary.BigEndian.Uint32(lengthBuf)

	// Sanity check message size
	if messageLength > 1024*1024 { // 1MB max message size
		return nil, fmt.Errorf("message size too large: %d bytes", messageLength)
	}

	// Read message body
	messageBuf := make([]byte, messageLength)
	if _, err := io.ReadFull(mh.reader, messageBuf); err != nil {
		return nil, fmt.Errorf("failed to read message body: %v", err)
	}

	// Unmarshal the message
	var message pb.ServerMessage
	if err := proto.Unmarshal(messageBuf, &message); err != nil {
		return nil, fmt.Errorf("failed to unmarshal message: %v", err)
	}

	return &message, nil
}

// WriteMessage writes a protobuf message to the network
func (mh *MessageHandler) WriteMessage(msg proto.Message) error {
	data, err := proto.Marshal(msg)
	if err != nil {
		return err
	}

	// Write length prefix
	lengthBuf := make([]byte, 4)
	binary.BigEndian.PutUint32(lengthBuf, uint32(len(data)))
	if _, err := mh.conn.Write(lengthBuf); err != nil {
		return err
	}

	// Write message body
	_, err = mh.conn.Write(data)
	return err
}

// UpdateGameState processes a server message and updates game state
func UpdateGameState(serverMessage *pb.ServerMessage, player *types.Player, otherPlayers map[int32]*types.Player) {
	playerID := player.ID

	player.Lock()
	player.CurrentTick = serverMessage.CurrentTick

	tickDiff := serverMessage.CurrentTick - player.CurrentTick
	if tickDiff > types.MaxTickDesync {
		for _, state := range serverMessage.Players {
			if state.PlayerId == playerID {
				player.ForceResync(state)
				break
			}
		}
	}
	player.Unlock()

	// Process player states
	validPlayerIds := make(map[int32]bool)
	for _, state := range serverMessage.Players {
		validPlayerIds[state.PlayerId] = true

		if state.PlayerId == playerID {
			player.Lock()
			// Update initial position if not set
			if player.PosActual.X == 0 && player.PosActual.Z == 0 {
				player.PosActual = rl.Vector3{
					X: float32(state.X * types.TileSize),
					Y: 0,
					Z: float32(state.Y * types.TileSize),
				}
				player.PosTile = types.Tile{X: int(state.X), Y: int(state.Y)}
			}
			player.Unlock()
			continue
		}

		// Update or create other players
		if otherPlayer, exists := otherPlayers[state.PlayerId]; exists {
			otherPlayer.UpdatePosition(state, types.ServerTickRate)
		} else {
			log.Printf("Creating new player with ID: %d", state.PlayerId)
			otherPlayers[state.PlayerId] = types.NewPlayer(state)
		}
	}

	// Remove players no longer in the server state
	for id := range otherPlayers {
		if id != playerID && !validPlayerIds[id] {
			log.Printf("Removing player with ID: %d", id)
			delete(otherPlayers, id)
		}
	}

	// Handle chat messages
	if handler, ok := player.UserData.(types.ChatMessageHandler); ok && len(serverMessage.ChatMessages) > 0 {
		log.Printf("Received %d chat messages from server", len(serverMessage.ChatMessages))
		handler.HandleServerMessages(serverMessage.ChatMessages)

		// Update the last seen message timestamp to the most recent message
		if len(serverMessage.ChatMessages) > 0 {
			lastMsg := serverMessage.ChatMessages[len(serverMessage.ChatMessages)-1]
			lastSeenMessageTimestamp = lastMsg.Timestamp
			log.Printf("Updated last seen message timestamp to %d", lastSeenMessageTimestamp)
		}
	}
}

func ConnectToServer(username, password string, isRegistering bool) (net.Conn, int32, error) {
	log.Printf("Connecting to server at %s...", serverAddr)

	var err error
	var conn net.Conn

	// Try connecting with a timeout
	connChan := make(chan net.Conn, 1)
	errChan := make(chan error, 1)

	go func() {
		c, e := net.Dial("tcp", serverAddr)
		if e != nil {
			errChan <- e
			return
		}
		connChan <- c
	}()

	// Wait for connection with timeout
	select {
	case conn = <-connChan:
		// Connection successful, continue
	case err = <-errChan:
		return nil, 0, fmt.Errorf("failed to dial server: %v", err)
	case <-time.After(5 * time.Second):
		return nil, 0, fmt.Errorf("connection timeout after 5 seconds")
	}

	log.Println("Connected to server. Authenticating...")

	// Create a message handler
	msgHandler := NewMessageHandler(conn)

	// Send auth message
	authAction := &pb.Action{
		Type:     pb.Action_LOGIN,
		Username: username,
		Password: password,
	}
	if isRegistering {
		authAction.Type = pb.Action_REGISTER
	}

	authBatch := &pb.ActionBatch{
		Actions:         []*pb.Action{authAction},
		ProtocolVersion: protoVersion,
	}

	if err := msgHandler.WriteMessage(authBatch); err != nil {
		conn.Close()
		return nil, 0, fmt.Errorf("failed to send auth: %v", err)
	}

	// Set a read deadline for authentication
	conn.SetReadDeadline(time.Now().Add(10 * time.Second))

	// Read server response
	response, err := msgHandler.ReadMessage()
	if err != nil {
		conn.Close()
		return nil, 0, fmt.Errorf("failed to read auth response: %v", err)
	}

	// Clear read deadline after authentication
	conn.SetReadDeadline(time.Time{})

	if response.ProtocolVersion > protoVersion {
		conn.Close()
		return nil, 0, fmt.Errorf("server requires newer protocol version (server: %d, client: %d)",
			response.ProtocolVersion, protoVersion)
	}

	if !response.AuthSuccess {
		conn.Close()
		return nil, 0, fmt.Errorf(response.ErrorMessage)
	}

	playerID := response.GetPlayerId()
	log.Printf("Successfully authenticated with player ID: %d", playerID)

	// Reset the lastSeenMessageTimestamp when reconnecting
	lastSeenMessageTimestamp = 0

	return conn, playerID, nil
}

func HandleServerCommunication(conn net.Conn, playerID int32, player *types.Player, otherPlayers map[int32]*types.Player, quitChan <-chan struct{}) {
	msgHandler := NewMessageHandler(conn)

	// Create channels for coordinating goroutines
	errChan := make(chan error, 1)
	done := make(chan struct{})

	// Create a WaitGroup to track both sender and receiver goroutines
	var wg sync.WaitGroup
	wg.Add(2) // One for sender, one for receiver

	// Set up a deferred cleanup function
	defer func() {
		if r := recover(); r != nil {
			log.Printf("Recovered from panic in HandleServerCommunication: %v", r)
		}

		// Close the done channel to signal both goroutines to exit
		close(done)

		// Wait for both goroutines to finish
		wg.Wait()

		// Close the connection
		conn.Close()

		// Close the player's QuitDone channel if it exists
		if player.QuitDone != nil {
			select {
			case <-player.QuitDone: // Check if it's already closed
				// Already closed, do nothing
			default:
				close(player.QuitDone)
			}
		}
	}()

	actionTicker := time.NewTicker(types.ClientTickRate)
	defer actionTicker.Stop()

	// Add a heartbeat ticker to detect connection issues
	heartbeatTicker := time.NewTicker(5 * time.Second)
	defer heartbeatTicker.Stop()

	lastMessageTime := time.Now()

	// Start message sending goroutine
	go func() {
		defer func() {
			if r := recover(); r != nil {
				log.Printf("Recovered from panic in message sender: %v", r)
				select {
				case errChan <- fmt.Errorf("message sender panic: %v", r):
				default:
					// Channel already closed or full, just log
					log.Printf("Unable to send error: %v", r)
				}
			}
			wg.Done() // Mark this goroutine as done
		}()

		for {
			select {
			case <-quitChan:
				// Send disconnect message to server
				disconnectMsg := &pb.ActionBatch{
					PlayerId: playerID,
					Actions: []*pb.Action{{
						Type:     pb.Action_DISCONNECT,
						PlayerId: playerID,
					}},
				}

				// Try to send disconnect message, ignoring errors
				_ = msgHandler.WriteMessage(disconnectMsg)

				// No need to signal done channel here, the main goroutine handles this
				return
			case <-done:
				return
			case <-heartbeatTicker.C:
				// If no message has been sent for a while, send a heartbeat
				timeSinceLastMessage := time.Since(lastMessageTime)
				if timeSinceLastMessage > 5*time.Second {
					// Send an empty batch as a heartbeat
					emptyBatch := &pb.ActionBatch{
						PlayerId:                 playerID,
						LastSeenMessageTimestamp: lastSeenMessageTimestamp,
					}
					if err := msgHandler.WriteMessage(emptyBatch); err != nil {
						log.Printf("Failed to send heartbeat: %v", err)
						select {
						case errChan <- err:
						case <-done:
							return
						}
					}
					lastMessageTime = time.Now()
				}
			case <-actionTicker.C:
				player.Lock()
				if len(player.ActionQueue) > 0 {
					actions := make([]*pb.Action, len(player.ActionQueue))
					copy(actions, player.ActionQueue)
					batch := &pb.ActionBatch{
						PlayerId:                 playerID,
						Actions:                  actions,
						Tick:                     player.CurrentTick,
						LastSeenMessageTimestamp: lastSeenMessageTimestamp,
					}
					player.ActionQueue = player.ActionQueue[:0]
					player.Unlock()

					if err := msgHandler.WriteMessage(batch); err != nil {
						select {
						case errChan <- err:
						case <-done:
							return
						}
					}
					lastMessageTime = time.Now()
				} else {
					player.Unlock()
				}
			}
		}
	}()

	// Main message receiving loop
	go func() {
		defer func() {
			if r := recover(); r != nil {
				log.Printf("Recovered from panic in message receiver: %v", r)
				select {
				case errChan <- fmt.Errorf("message receiver panic: %v", r):
				default:
					// Channel already closed or full, just log
					log.Printf("Unable to send error: %v", r)
				}
			}
			wg.Done() // Mark this goroutine as done
		}()

		for {
			select {
			case <-quitChan:
				return
			case <-done:
				return
			default:
				serverMessage, err := msgHandler.ReadMessage()
				if err != nil {
					if err, ok := err.(net.Error); ok && err.Timeout() {
						log.Printf("Network timeout: %v", err)
					} else if err != io.EOF {
						log.Printf("Network read error: %v", err)
						select {
						case errChan <- err:
						case <-done:
							return
						}
					} else {
						log.Printf("Connection closed by server")
					}
					return
				}

				// Process the server message
				UpdateGameState(serverMessage, player, otherPlayers)
			}
		}
	}()

	// Wait for error or quit signal
	select {
	case <-quitChan:
		log.Printf("Received quit signal, sending disconnect message")
		// The cleanup will happen in the deferred function
		return
	case err := <-errChan:
		log.Printf("Network error: %v", err)
		// The cleanup will happen in the deferred function
		return
	}
}

// Helper function to write length-prefixed messages
func writeMessage(conn net.Conn, msg proto.Message) error {
	msgHandler := NewMessageHandler(conn)
	return msgHandler.WriteMessage(msg)
}

type Connection struct {
	conn      net.Conn
	playerID  int32
	quitChan  chan struct{}
	quitDone  chan struct{}
	closeOnce sync.Once
}

func NewConnection(username, password string, isRegistering bool) (*Connection, error) {
	conn, playerID, err := ConnectToServer(username, password, isRegistering)
	if err != nil {
		return nil, err
	}
	return &Connection{
		conn:     conn,
		playerID: playerID,
		quitChan: make(chan struct{}),
		quitDone: make(chan struct{}),
	}, nil
}

func (c *Connection) Close() {
	c.closeOnce.Do(func() {
		select {
		case <-c.quitChan: // Check if it's already closed
			// Already closed, do nothing
		default:
			close(c.quitChan)
		}

		// Wait with timeout for network cleanup
		select {
		case <-c.quitDone:
			// Clean shutdown completed
		case <-time.After(500 * time.Millisecond):
			log.Println("Network cleanup timed out")
		}

		// Make sure the connection is closed
		c.conn.Close()
	})
}

func (c *Connection) PlayerID() int32 {
	return c.playerID
}

func (c *Connection) Start(player *types.Player, otherPlayers map[int32]*types.Player) {
	go HandleServerCommunication(c.conn, c.playerID, player, otherPlayers, c.quitChan)
}

func (c *Connection) QuitChan() <-chan struct{} {
	return c.quitChan
}