server.go

package main

import (
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"sync"

	"golang.org/x/net/websocket"
)

type Room struct {
	ID      string
	Name    string
	Clients int
}

type User struct {
	ID       string
	Username string
	Conn     *websocket.Conn
	Send     chan string
	Room     string
}

type Message struct {
	Sender  string
	Content string
}

type ChangeRoom struct {
	roomID string
	user   *User
}

type Hub struct {
	Mutex      sync.Mutex
	Users      map[string]*User
	Conns      map[*websocket.Conn]*User
	Register   chan *User
	Unregister chan *User
	JoinRoom   chan *ChangeRoom
	Messages   chan *Message
	Rooms      map[string]*Room
}

const (
	DEFAULT_ROOM_ID = "Default"
)

func NewHub() *Hub {
	return &Hub{
		Users:      make(map[string]*User),
		Conns:      make(map[*websocket.Conn]*User),
		Register:   make(chan *User),
		Unregister: make(chan *User),
		JoinRoom:   make(chan *ChangeRoom),
		Messages:   make(chan *Message),
		Rooms:      make(map[string]*Room),
	}
}

func (h *Hub) Run() {
	for {
		select {
		case user := <-h.Register:
			h.Users[user.ID] = user
			h.Conns[user.Conn] = user
			fmt.Println("User register:", user.ID)

		case user := <-h.Unregister:
			delete(h.Users, user.ID)
			delete(h.Conns, user.Conn)
			close(user.Send)
			fmt.Println("User unregister:", user.ID)

		case ChangeRoom := <-h.JoinRoom:
			user := ChangeRoom.user
			currentRoom := user.Room
			nextRoom := ChangeRoom.roomID

			// join next room
			h.Users[user.ID].Room = nextRoom
			if _, ok := h.Rooms[nextRoom]; !ok {
				h.Rooms[nextRoom] = &Room{ID: nextRoom, Name: nextRoom}
			}
			h.Rooms[nextRoom].Clients++

			// leave current room
			if currentRoom == "" {
				// case when you want to enter DEFAULT_ROOM and you aren't in a ROOM
				continue
			}
			h.Rooms[currentRoom].Clients--
			if h.Rooms[currentRoom].Clients == 0 && currentRoom != DEFAULT_ROOM_ID {
				delete(h.Rooms, currentRoom)
			}

		case msg := <-h.Messages:
			fmt.Printf("Broadcasting message from %s to %d users", msg.Sender, len(h.Users))
			senderRoom := h.Users[msg.Sender].Room
			for _, user := range h.Users {
				if user.Room == senderRoom {
					user.Send <- fmt.Sprintf("%s: %s", msg.Sender, msg.Content)
					fmt.Println("Sent to", user.ID)
				}
			}
		}
	}
}

func (h *Hub) HandleWebSocket(ws *websocket.Conn) {
	// Read initial username
	var username string
	if err := websocket.Message.Receive(ws, &username); err != nil {
		fmt.Println("Error reading username:", err)
		return
	}
	username = strings.TrimSpace(username)

	// Create user
	user := &User{
		ID:       username,
		Username: username,
		Conn:     ws,
		Send:     make(chan string),
	}
	h.Register <- user
	defer func() { h.Unregister <- user }()
	h.JoinRoom <- &ChangeRoom{user: user, roomID: DEFAULT_ROOM_ID}
	h.Messages <- &Message{Sender: user.ID, Content: fmt.Sprintf("%s joined the room", user.ID)}

	// Send messages to client
	go func() {
		for msg := range user.Send {
			if err := websocket.Message.Send(ws, msg); err != nil {
				fmt.Println("Error sending to", user.ID, err)
				return
			}
		}
	}()

	// Handle incoming messages
	for {
		var data string
		if err := websocket.Message.Receive(ws, &data); err != nil {
			fmt.Println("Client disconnect:", user.ID, err)
			return
		}
		data = strings.TrimSpace(data)

		switch {
		case data == "/list":
			user := h.Conns[ws]
			currentChar := ' '
			user.Send <- "Chat Rooms:"
			for name, room := range h.Rooms {
				if name == user.Room {
					currentChar = '*'
				}
				user.Send <- fmt.Sprintf("%c %s (%d)", currentChar, name, room.Clients)
			}

		case strings.HasPrefix(data, "/join"):
			roomID := data[6:]
			h.Messages <- &Message{Sender: user.ID, Content: fmt.Sprintf("%s left the room", user.ID)}
			h.JoinRoom <- &ChangeRoom{user: user, roomID: roomID}
			h.Messages <- &Message{Sender: user.ID, Content: fmt.Sprintf("%s joined the room", user.ID)}

		case strings.HasPrefix(data, "FILE:"):
			parts := strings.Split(data, ":")
			filename, sizeStr, recieverID := parts[1], parts[2], parts[3]
			size, _ := strconv.ParseInt(sizeStr, 10, 64)
			h.HandleFileTransfer(user, filename, size, recieverID)

		default:
			// Broadcast messages
			h.Messages <- &Message{Sender: user.ID, Content: data}
		}
	}
}

func (h *Hub) HandleFileTransfer(sender *User, filename string, size int64, recieverID string) {
	// Notify reciever
	reciever, exists := h.Users[recieverID]
	if !exists {
		sender.Send <- "Reciever not found"
		return
	}

	// Signal reciever to expect file
	reciever.Send <- fmt.Sprintf("FILE:%s:%d:%s", filename, size, sender.ID)

	// Recieve file data (asumes bin data follows)
	fileData := make([]byte, size)
	if err := websocket.Message.Receive(sender.Conn, &fileData); err != nil {
		sender.Send <- fmt.Sprintf("Error recieving file: %v", err)
		return
	}

	// Forward file to reciever
	if err := websocket.Message.Send(reciever.Conn, fileData); err != nil {
		sender.Send <- fmt.Sprintf("Error sending file to %s: %v", recieverID, err)
		return
	}

	sender.Send <- fmt.Sprintf("File %s sent to %s", filename, recieverID)
	reciever.Send <- fmt.Sprintf("File %s recieved from %s", filename, sender.ID)
}

func main() {
	hub := NewHub()
	go hub.Run()

	http.Handle("/ws", websocket.Handler(hub.HandleWebSocket))
	http.ListenAndServe(":8000", nil)
}