doormile_backend/internal/ws/chat.go

package ws

import (
	"encoding/json"
	"fmt"
	"strconv"
	"sync"
	"time"

	"doormile/db"
	"doormile/models"
	"doormile/utils"

	"github.com/gofiber/websocket/v2"
)

// chatMessage is the JSON frame broadcast to the receiving participant.
type chatMessage struct {
	Sender    string `json:"sender"`
	Message   string `json:"message"`
	Timestamp string `json:"timestamp"`
}

// chatConn wraps a single WebSocket connection with a write mutex so concurrent
// senders (broadcast from the other peer, poller close) never interleave frames.
type chatConn struct {
	conn *websocket.Conn
	mu   sync.Mutex
}

func (cc *chatConn) send(data []byte) {
	cc.mu.Lock()
	defer cc.mu.Unlock()
	cc.conn.WriteMessage(websocket.TextMessage, data) //nolint:errcheck — best-effort delivery
}

func (cc *chatConn) close() {
	cc.mu.Lock()
	defer cc.mu.Unlock()
	cc.conn.Close() //nolint:errcheck
}

// chatRoom holds at most two named slots: "customer" and "miler".
// All slot mutations are protected by mu.
type chatRoom struct {
	bookingID  int
	mu         sync.Mutex
	slots      map[string]*chatConn
	closeOnce  sync.Once  // ensures shutdown() body runs exactly once
	closeCh    chan struct{} // closed by shutdown()
	pollerOnce sync.Once  // ensures exactly one status-poller goroutine per room
}

// rooms is the process-wide registry of active chat rooms, keyed by bookingID.
var rooms sync.Map // map[int]*chatRoom

// ChatHandler is the WebSocket handler for the ephemeral per-booking chat room.
//
// Route:  GET /ws/bookings/:bookingid/chat
// Params: role=customer|miler   (query)
//         token=<JWT>           (query, validated by WsChatAuth middleware)
//
// At most two participants (one customer, one miler) may occupy a room.
// Messages received from one participant are forwarded to the other.
// The room closes automatically when the booking reaches a terminal status.
func ChatHandler(c *websocket.Conn) {
	bookingID, err := strconv.Atoi(c.Params("bookingid"))
	if err != nil {
		sendError(c, "invalid booking ID")
		return
	}

	role := c.Query("role")
	if role != "customer" && role != "miler" {
		sendError(c, "role must be 'customer' or 'miler'")
		return
	}

	// Verify the booking exists before letting anyone into the room.
	var booking models.PickupBooking
	if err := db.DB.First(&booking, bookingID).Error; err != nil {
		sendError(c, "booking not found")
		return
	}
	if isTerminalStatus(booking.Status) {
		sendError(c, "chat is closed — booking is already completed")
		return
	}

	room, err := joinRoom(bookingID, role, c)
	if err != nil {
		sendError(c, err.Error())
		return
	}
	defer room.leave(role)

	utils.Info("WS/Chat: participant joined", "booking_id", bookingID, "role", role)

	// Block here reading client messages. Exits on client disconnect or room close.
	for {
		_, raw, err := c.ReadMessage()
		if err != nil {
			break
		}
		room.broadcast(role, string(raw))
	}

	utils.Info("WS/Chat: participant left", "booking_id", bookingID, "role", role)
}

// ─── Room lifecycle ───────────────────────────────────────────────────────────

// joinRoom finds or creates the room for bookingID, registers the connection in
// the given role slot, and starts the booking-status poller on first join.
// Returns an error if the slot is already occupied or the room is shutting down.
func joinRoom(bookingID int, role string, c *websocket.Conn) (*chatRoom, error) {
	val, _ := rooms.LoadOrStore(bookingID, &chatRoom{
		bookingID: bookingID,
		slots:     make(map[string]*chatConn),
		closeCh:   make(chan struct{}),
	})
	room := val.(*chatRoom)

	room.mu.Lock()
	defer room.mu.Unlock()

	// If shutdown already started (e.g., race with a closing poller), reject.
	select {
	case <-room.closeCh:
		return nil, fmt.Errorf("chat room is closed")
	default:
	}

	if _, exists := room.slots[role]; exists {
		return nil, fmt.Errorf("%s is already connected to this room", role)
	}

	room.slots[role] = &chatConn{conn: c}

	// Start the status poller the first time any participant joins.
	room.pollerOnce.Do(func() { go room.pollBookingStatus() })

	return room, nil
}

// leave removes the participant from the room and shuts the room down if empty.
// Called via defer in ChatHandler — safe even after a force-close.
func (r *chatRoom) leave(role string) {
	r.mu.Lock()
	delete(r.slots, role)
	isEmpty := len(r.slots) == 0
	r.mu.Unlock()

	if isEmpty {
		r.shutdown()
	}
}

// broadcast wraps the raw text in a chatMessage frame and sends it to every
// participant other than the sender.
func (r *chatRoom) broadcast(senderRole, text string) {
	frame := chatMessage{
		Sender:    senderRole,
		Message:   text,
		Timestamp: time.Now().UTC().Format(time.RFC3339),
	}
	data, err := json.Marshal(frame)
	if err != nil {
		return
	}

	r.mu.Lock()
	targets := make([]*chatConn, 0, 1)
	for role, cc := range r.slots {
		if role != senderRole {
			targets = append(targets, cc)
		}
	}
	r.mu.Unlock()

	for _, cc := range targets {
		cc.send(data)
	}
}

// shutdown closes the done channel and removes the room from the global registry.
// Protected by closeOnce — safe to call from both the poller and connection handlers.
func (r *chatRoom) shutdown() {
	r.closeOnce.Do(func() {
		close(r.closeCh)
		rooms.Delete(r.bookingID)
		utils.Info("WS/Chat: room removed from registry", "booking_id", r.bookingID)
	})
}

// forceClose closes every open WebSocket connection then shuts the room down.
// Called by the status poller when the booking reaches a terminal status.
// Each handler's c.ReadMessage() will return an error, causing it to break
// its read loop and call room.leave() — which hits shutdown() again (no-op).
func (r *chatRoom) forceClose() {
	r.mu.Lock()
	conns := make([]*chatConn, 0, len(r.slots))
	for _, cc := range r.slots {
		conns = append(conns, cc)
	}
	r.mu.Unlock()

	for _, cc := range conns {
		cc.close()
	}

	r.shutdown()
}

// ─── Status poller ────────────────────────────────────────────────────────────

// pollBookingStatus runs in a goroutine for the room's lifetime.
// Every 5 s it reloads the booking; when it hits a terminal status it publishes
// the NATS event, force-closes all connections, and exits.
func (r *chatRoom) pollBookingStatus() {
	t := time.NewTicker(5 * time.Second)
	defer t.Stop()

	for {
		select {
		case <-r.closeCh:
			return

		case <-t.C:
			var booking models.PickupBooking
			if err := db.DB.First(&booking, r.bookingID).Error; err != nil {
				continue
			}

			if isTerminalStatus(booking.Status) {
				utils.Info("WS/Chat: closing room — booking reached terminal status",
					"booking_id", r.bookingID,
					"status", booking.Status,
				)
				publishChatRoomClosed(r.bookingID)
				r.forceClose()
				return
			}
		}
	}
}

// publishChatRoomClosed publishes a NATS event so downstream services know the
// chat session has ended. Non-fatal — logs a warning on failure.
func publishChatRoomClosed(bookingID int) {
	if db.Js == nil {
		return
	}

	payload := map[string]interface{}{
		"booking_id": bookingID,
		"closed_at":  time.Now().UnixMilli(),
	}
	data, err := json.Marshal(payload)
	if err != nil {
		return
	}

	subject := fmt.Sprintf("chat.room.closed.%d", bookingID)
	if _, err := db.Js.Publish(subject, data); err != nil {
		utils.Warn("WS/Chat: failed to publish room-closed event",
			"booking_id", bookingID,
			"error", err,
		)
	}
}