import "github.com/gorilla/websocket"
 
var upgrader = websocket.Upgrader{
    ReadBufferSize:  1024,
    WriteBufferSize: 1024,
    CheckOrigin: func(r *http.Request) bool {
        // CSRF protection: only allow your domain
        origin := r.Header.Get("Origin")
        return origin == "https://app.example.com"
    },
}
 
type Client struct {
    conn   *websocket.Conn
    userID string
    send   chan []byte  // buffered channel for outgoing messages
}
 
func HandleWebSocket(w http.ResponseWriter, r *http.Request) {
    // Authenticate before upgrade (JWT in query param or cookie)
    userID, err := authenticateWS(r)
    if err != nil {
        http.Error(w, "Unauthorized", 401)
        return
    }
    
    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        return
    }
    
    client := &Client{
        conn:   conn,
        userID: userID,
        send:   make(chan []byte, 256),  // buffer for slow clients
    }
    
    hub.Register(client)
    defer hub.Unregister(client)
    
    // Two goroutines per connection
    go client.writePump()  // sends messages to client
    client.readPump()      // reads messages from client (blocks)
}
 
func (c *Client) readPump() {
    defer c.conn.Close()
    
    c.conn.SetReadLimit(512 * 1024)  // max message: 512KB
    c.conn.SetReadDeadline(time.Now().Add(60 * time.Second))
    c.conn.SetPongHandler(func(string) error {
        c.conn.SetReadDeadline(time.Now().Add(60 * time.Second))
        return nil
    })
    
    for {
        _, message, err := c.conn.ReadMessage()
        if err != nil {
            if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
                log.Println("ws error:", err)
            }
            return  // exits, defer closes connection
        }
        hub.Broadcast(message)
    }
}
 
func (c *Client) writePump() {
    ticker := time.NewTicker(30 * time.Second)  // ping every 30s
    defer func() {
        ticker.Stop()
        c.conn.Close()
    }()
    
    for {
        select {
        case message, ok := <-c.send:
            c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
            if !ok {
                c.conn.WriteMessage(websocket.CloseMessage, []byte{})
                return
            }
            if err := c.conn.WriteMessage(websocket.TextMessage, message); err != nil {
                return
            }
            
        case <-ticker.C:
            c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
            if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
                return  // client gone
            }
        }
    }
}

type Hub struct {
    clients    map[*Client]bool
    rooms      map[string]map[*Client]bool  // room → clients
    broadcast  chan []byte
    register   chan *Client
    unregister chan *Client
    mu         sync.RWMutex
}
 
func (h *Hub) Run() {
    for {
        select {
        case client := <-h.register:
            h.mu.Lock()
            h.clients[client] = true
            h.mu.Unlock()
            
        case client := <-h.unregister:
            h.mu.Lock()
            if _, ok := h.clients[client]; ok {
                delete(h.clients, client)
                close(client.send)
            }
            h.mu.Unlock()
            
        case message := <-h.broadcast:
            h.mu.RLock()
            for client := range h.clients {
                select {
                case client.send <- message:
                    // sent
                default:
                    // client's send buffer full (slow client)
                    // drop message OR disconnect client
                    close(client.send)
                    delete(h.clients, client)
                }
            }
            h.mu.RUnlock()
        }
    }
}

// Subscribe to Redis channel on server startup
func (h *Hub) SubscribeToRedis(ctx context.Context, rdb *redis.Client) {
    pubsub := rdb.Subscribe(ctx, "ws:broadcast", "ws:room:*")
    defer pubsub.Close()
    
    for msg := range pubsub.Channel() {
        var wsMsg WSMessage
        json.Unmarshal([]byte(msg.Payload), &wsMsg)
        
        // Deliver to local clients only
        h.DeliverLocal(wsMsg)
    }
}
 
// Send message from any service → publish to Redis → all WS servers receive
func (s *ChatService) SendMessage(ctx context.Context, roomID string, msg *Message) {
    payload, _ := json.Marshal(WSMessage{Room: roomID, Data: msg})
    rdb.Publish(ctx, "ws:room:"+roomID, payload)
}

// When bidirectional not needed (order tracking, notifications)
func OrderUpdatesSSE(w http.ResponseWriter, r *http.Request) {
    // Check SSE support
    flusher, ok := w.(http.Flusher)
    if !ok {
        http.Error(w, "SSE not supported", 500)
        return
    }
    
    w.Header().Set("Content-Type", "text/event-stream")
    w.Header().Set("Cache-Control", "no-cache")
    w.Header().Set("Connection", "keep-alive")
    w.Header().Set("Access-Control-Allow-Origin", "https://app.example.com")
    
    orderID := r.URL.Query().Get("order_id")
    ctx := r.Context()
    
    updateCh := orderService.Subscribe(ctx, orderID)
    
    for {
        select {
        case update := <-updateCh:
            // SSE format: "data: {json}\n\n"
            fmt.Fprintf(w, "id: %d\n", update.ID)
            fmt.Fprintf(w, "event: order_update\n")
            fmt.Fprintf(w, "data: %s\n\n", mustJSON(update))
            flusher.Flush()  // send immediately, don't buffer
            
        case <-ctx.Done():
            return  // client disconnected
        }
    }
}