WebSockets & Real-time

package main

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"github.com/google/uuid"
	"github.com/gorilla/websocket"
)

// --- Types ---
type ChatMessage struct {
	ID        string `json:"id"`
	Room      string `json:"room"`
	UserID    string `json:"userId"`
	Content   string `json:"content"`
	Timestamp int64  `json:"timestamp"`
}

type ClientMessage struct {
	Type    string `json:"type"`
	Room    string `json:"room,omitempty"`
	Content string `json:"content,omitempty"`
}

type ServerMessage struct {
	Type      string        `json:"type"`
	Room      string        `json:"room,omitempty"`
	UserID    string        `json:"userId,omitempty"`
	Content   string        `json:"content,omitempty"`
	Messages  []ChatMessage `json:"messages,omitempty"`
	Users     []string      `json:"users,omitempty"`
	Timestamp int64         `json:"timestamp,omitempty"`
}

// --- Ring Buffer ---
type RingBuffer struct {
	buf   []ChatMessage
	head  int
	count int
	cap   int
}

func NewRingBuffer(capacity int) *RingBuffer {
	return &RingBuffer{buf: make([]ChatMessage, capacity), cap: capacity}
}

func (rb *RingBuffer) Push(msg ChatMessage) {
	rb.buf[rb.head] = msg
	rb.head = (rb.head + 1) % rb.cap
	if rb.count < rb.cap {
		rb.count++
	}
}

func (rb *RingBuffer) GetAll() []ChatMessage {
	if rb.count == 0 {
		return nil
	}
	result := make([]ChatMessage, 0, rb.count)
	start := 0
	if rb.count == rb.cap {
		start = rb.head
	}
	for i := 0; i < rb.count; i++ {
		idx := (start + i) % rb.cap
		result = append(result, rb.buf[idx])
	}
	return result
}

// --- Connected Client ---
type ConnectedClient struct {
	conn    *websocket.Conn
	userID  string
	rooms   map[string]bool
	mu      sync.Mutex
	isAlive bool
}

func (c *ConnectedClient) Send(msg ServerMessage) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
	return c.conn.WriteJSON(msg)
}

// --- Room Manager ---
type RoomManager struct {
	mu      sync.RWMutex
	members map[string]map[string]bool    // room -> set of userIDs
	history map[string]*RingBuffer        // room -> message history
}

func NewRoomManager() *RoomManager {
	return &RoomManager{
		members: make(map[string]map[string]bool),
		history: make(map[string]*RingBuffer),
	}
}

func (rm *RoomManager) Join(room, userID string) []string {
	rm.mu.Lock()
	defer rm.mu.Unlock()
	if rm.members[room] == nil {
		rm.members[room] = make(map[string]bool)
		rm.history[room] = NewRingBuffer(100)
	}
	rm.members[room][userID] = true
	return rm.memberList(room)
}

func (rm *RoomManager) Leave(room, userID string) []string {
	rm.mu.Lock()
	defer rm.mu.Unlock()
	if rm.members[room] == nil {
		return nil
	}
	delete(rm.members[room], userID)
	if len(rm.members[room]) == 0 {
		delete(rm.members, room)
		delete(rm.history, room)
		return nil
	}
	return rm.memberList(room)
}

func (rm *RoomManager) memberList(room string) []string {
	members := rm.members[room]
	list := make([]string, 0, len(members))
	for id := range members {
		list = append(list, id)
	}
	return list
}

func (rm *RoomManager) GetMembers(room string) []string {
	rm.mu.RLock()
	defer rm.mu.RUnlock()
	return rm.memberList(room)
}

func (rm *RoomManager) AddMessage(room string, msg ChatMessage) {
	rm.mu.Lock()
	defer rm.mu.Unlock()
	if buf := rm.history[room]; buf != nil {
		buf.Push(msg)
	}
}

func (rm *RoomManager) GetHistory(room string) []ChatMessage {
	rm.mu.RLock()
	defer rm.mu.RUnlock()
	if buf := rm.history[room]; buf != nil {
		return buf.GetAll()
	}
	return nil
}

// --- Connection Manager ---
type ConnectionManager struct {
	mu      sync.RWMutex
	clients map[*websocket.Conn]*ConnectedClient
}

func NewConnectionManager() *ConnectionManager {
	return &ConnectionManager{
		clients: make(map[*websocket.Conn]*ConnectedClient),
	}
}

func (cm *ConnectionManager) Add(conn *websocket.Conn, userID string) *ConnectedClient {
	cm.mu.Lock()
	defer cm.mu.Unlock()
	client := &ConnectedClient{
		conn:    conn,
		userID:  userID,
		rooms:   make(map[string]bool),
		isAlive: true,
	}
	cm.clients[conn] = client
	return client
}

func (cm *ConnectionManager) Remove(conn *websocket.Conn) *ConnectedClient {
	cm.mu.Lock()
	defer cm.mu.Unlock()
	client := cm.clients[conn]
	delete(cm.clients, conn)
	return client
}

func (cm *ConnectionManager) GetByRoom(room string) []*ConnectedClient {
	cm.mu.RLock()
	defer cm.mu.RUnlock()
	var result []*ConnectedClient
	for _, c := range cm.clients {
		if c.rooms[room] {
			result = append(result, c)
		}
	}
	return result
}

func (cm *ConnectionManager) GetAll() []*ConnectedClient {
	cm.mu.RLock()
	defer cm.mu.RUnlock()
	result := make([]*ConnectedClient, 0, len(cm.clients))
	for _, c := range cm.clients {
		result = append(result, c)
	}
	return result
}

// --- Broadcast ---
func broadcast(clients []*ConnectedClient, msg ServerMessage, exclude *websocket.Conn) {
	for _, c := range clients {
		if c.conn != exclude {
			c.Send(msg)
		}
	}
}

// --- Globals ---
var (
	upgrader = websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
		CheckOrigin:     func(r *http.Request) bool { return true },
	}
	roomMgr = NewRoomManager()
	connMgr = NewConnectionManager()
)

// --- WebSocket handler ---
func handleWebSocket(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("[ERROR] Upgrade failed: %v", err)
		return
	}

	userID := r.URL.Query().Get("userId")
	if userID == "" {
		userID = fmt.Sprintf("anon-%s", uuid.New().String()[:8])
	}

	log.Printf("[CONNECT] User %s connected", userID)
	client := connMgr.Add(conn, userID)

	// Configure connection
	conn.SetReadLimit(4096)
	conn.SetReadDeadline(time.Now().Add(60 * time.Second))
	conn.SetPongHandler(func(string) error {
		client.isAlive = true
		conn.SetReadDeadline(time.Now().Add(60 * time.Second))
		return nil
	})

	defer func() {
		log.Printf("[DISCONNECT] User %s disconnected", userID)
		connMgr.Remove(conn)
		conn.Close()

		// Leave all rooms
		for room := range client.rooms {
			members := roomMgr.Leave(room, userID)
			broadcast(connMgr.GetByRoom(room), ServerMessage{
				Type:      "leave",
				Room:      room,
				UserID:    userID,
				Users:     members,
				Timestamp: time.Now().UnixMilli(),
			}, nil)
		}
	}()

	// Read loop
	for {
		_, rawMsg, err := conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
				log.Printf("[WS_ERROR] User %s: %v", userID, err)
			}
			return
		}

		var msg ClientMessage
		if err := json.Unmarshal(rawMsg, &msg); err != nil {
			client.Send(ServerMessage{Type: "error", Content: "Invalid JSON"})
			continue
		}

		switch msg.Type {
		case "join":
			if msg.Room == "" {
				client.Send(ServerMessage{Type: "error", Content: "room is required"})
				continue
			}
			client.rooms[msg.Room] = true
			members := roomMgr.Join(msg.Room, userID)
			log.Printf("[JOIN] %s joined room %s (%d members)", userID, msg.Room, len(members))

			broadcast(connMgr.GetByRoom(msg.Room), ServerMessage{
				Type:      "join",
				Room:      msg.Room,
				UserID:    userID,
				Users:     members,
				Timestamp: time.Now().UnixMilli(),
			}, nil)

			if history := roomMgr.GetHistory(msg.Room); len(history) > 0 {
				client.Send(ServerMessage{Type: "history", Room: msg.Room, Messages: history})
			}

		case "leave":
			if msg.Room == "" {
				continue
			}
			delete(client.rooms, msg.Room)
			members := roomMgr.Leave(msg.Room, userID)
			log.Printf("[LEAVE] %s left room %s", userID, msg.Room)

			broadcast(connMgr.GetByRoom(msg.Room), ServerMessage{
				Type:      "leave",
				Room:      msg.Room,
				UserID:    userID,
				Users:     members,
				Timestamp: time.Now().UnixMilli(),
			}, nil)

		case "message":
			if msg.Room == "" || msg.Content == "" || !client.rooms[msg.Room] {
				client.Send(ServerMessage{Type: "error", Content: "Must join room first"})
				continue
			}
			content := msg.Content
			if len(content) > 4096 {
				content = content[:4096]
			}

			chatMsg := ChatMessage{
				ID:        uuid.New().String(),
				Room:      msg.Room,
				UserID:    userID,
				Content:   content,
				Timestamp: time.Now().UnixMilli(),
			}
			roomMgr.AddMessage(msg.Room, chatMsg)

			broadcast(connMgr.GetByRoom(msg.Room), ServerMessage{
				Type:      "message",
				Room:      msg.Room,
				UserID:    userID,
				Content:   chatMsg.Content,
				Timestamp: chatMsg.Timestamp,
			}, nil)

		case "history":
			if msg.Room == "" {
				continue
			}
			history := roomMgr.GetHistory(msg.Room)
			client.Send(ServerMessage{Type: "history", Room: msg.Room, Messages: history})

		default:
			client.Send(ServerMessage{Type: "error", Content: "Unknown message type"})
		}
	}
}

// --- Heartbeat goroutine ---
func startHeartbeat(stop chan struct{}) {
	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			for _, client := range connMgr.GetAll() {
				if !client.isAlive {
					log.Printf("[TIMEOUT] Terminating dead connection: %s", client.userID)
					client.conn.Close()
					continue
				}
				client.isAlive = false
				client.mu.Lock()
				client.conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
				client.conn.WriteMessage(websocket.PingMessage, nil)
				client.mu.Unlock()
			}
		case <-stop:
			return
		}
	}
}

// --- Main ---
func main() {
	port := os.Getenv("PORT")
	if port == "" {
		port = "3000"
	}

	mux := http.NewServeMux()
	mux.HandleFunc("/ws", handleWebSocket)
	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/plain")
		w.Write([]byte("WebSocket Chat Server"))
	})

	srv := &http.Server{
		Addr:         ":" + port,
		Handler:      mux,
		ReadTimeout:  10 * time.Second,
		WriteTimeout: 10 * time.Second,
	}

	stopHeartbeat := make(chan struct{})
	go startHeartbeat(stopHeartbeat)

	go func() {
		log.Printf("Chat server listening on ws://localhost:%s/ws", port)
		if err := srv.ListenAndServe(); err != http.ErrServerClosed {
			log.Fatalf("Server error: %v", err)
		}
	}()

	quit := make(chan os.Signal, 1)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
	<-quit

	log.Println("Shutting down...")
	close(stopHeartbeat)

	// Close all WebSocket connections
	for _, client := range connMgr.GetAll() {
		client.conn.WriteMessage(websocket.CloseMessage,
			websocket.FormatCloseMessage(websocket.CloseGoingAway, "Server shutting down"))
		client.conn.Close()
	}

	srv.Close()
	log.Println("Server closed.")
}