added commands, join messages

main.go

@@ -2,12 +2,15 @@ package main
 
 import (
 	"context"
+	"crypto/rand" // For more robust random generation
 	"errors"
 	"fmt"
+	"math/big" // For crypto/rand with charset
 	"net"
 	"os"
 	"os/signal"
 	"regexp"
+	"sort" // For /users command
 	"strings"
 	"sync"
 	"syscall"
@@ -28,10 +31,27 @@ import (
 )
 
 const (
-	host = "0.0.0.0"
+	host              = "0.0.0.0"
-	port = "23234"
+	port              = "23234"
+	maxRoomNameLength = 24
+	roomNameCharset   = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
 )
 
+// generateRandomRoomName creates a random alphanumeric string of a given length.
+func generateRandomRoomName(length int) string {
+	b := make([]byte, length)
+	for i := range b {
+		n, err := rand.Int(rand.Reader, big.NewInt(int64(len(roomNameCharset))))
+		if err != nil {
+			log.Errorf("Error generating random char for room name: %v, falling back for char %d", err, i)
+			b[i] = roomNameCharset[0] // Fallback to 'a'
+			continue
+		}
+		b[i] = roomNameCharset[n.Int64()]
+	}
+	return string(b)
+}
+
 type programInfo struct {
 	roomName string
 	userID   string
@@ -45,14 +65,18 @@ type app struct {
 	mu             sync.Mutex
 }
 
-// BroadcastToRoom sends a message to all running programs in a specific room.
+// BroadcastToRoom sends a message to all running programs in a specific room,
-func (a *app) BroadcastToRoom(roomName string, msg tea.Msg) {
+// optionally excluding one program.
+func (a *app) BroadcastToRoom(roomName string, msg tea.Msg, excludeProg *tea.Program) {
 	a.mu.Lock()
 	defer a.mu.Unlock()
 
 	if progsInRoom, ok := a.rooms[roomName]; ok {
-		log.Debugf("Broadcasting to room %s (%d programs): %+v", roomName, len(progsInRoom), msg)
+		log.Debugf("Broadcasting to room %s (%d programs, excluding %p): %+v", roomName, len(progsInRoom), excludeProg, msg)
 		for _, p := range progsInRoom {
+			if excludeProg != nil && p == excludeProg {
+				continue
+			}
 			go p.Send(msg)
 		}
 	} else {
@@ -60,14 +84,30 @@ func (a *app) BroadcastToRoom(roomName string, msg tea.Msg) {
 	}
 }
 
+func (a *app) GetUserIDsInRoom(roomName string) []string {
+	a.mu.Lock()
+	defer a.mu.Unlock()
+
+	var userIDs []string
+	if progsInRoom, ok := a.rooms[roomName]; ok {
+		for _, p := range progsInRoom {
+			if info, ok := a.activePrograms[p]; ok {
+				userIDs = append(userIDs, info.userID)
+			}
+		}
+	}
+	sort.Strings(userIDs) // Sort for consistent display
+	return userIDs
+}
+
 func newApp() *app {
 	a := &app{
-		activePrograms: make(map[*tea.Program]programInfo), // Initialize correctly
+		activePrograms: make(map[*tea.Program]programInfo),
 		rooms:          make(map[string][]*tea.Program),
 	}
 	s, err := wish.NewServer(
 		wish.WithAddress(net.JoinHostPort(host, port)),
-		wish.WithHostKeyPath(".ssh/id_ed25519"), // Ensure this key exists (e.g., run ssh-keygen -t ed25519 -f .ssh/id_ed25519 -N "")
+		wish.WithHostKeyPath(".ssh/id_ed25519"),
 		wish.WithMiddleware(
 			bubbletea.MiddlewareWithProgramHandler(a.ProgramHandler, termenv.ANSI256),
 			activeterm.Middleware(),
@@ -90,7 +130,7 @@ func (a *app) Start() {
 	go func() {
 		if err = a.ListenAndServe(); err != nil && !errors.Is(err, ssh.ErrServerClosed) {
 			log.Error("Could not start server", "error", err)
-			done <- syscall.SIGTERM // Signal main to exit if server fails
+			done <- syscall.SIGTERM
 		}
 	}()
 
@@ -99,9 +139,8 @@ func (a *app) Start() {
 	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
 	defer func() {
 		cancel()
-		// Clean up internal state after server shutdown
 		a.mu.Lock()
-		a.activePrograms = make(map[*tea.Program]programInfo) // Correct type
+		a.activePrograms = make(map[*tea.Program]programInfo)
 		a.rooms = make(map[string][]*tea.Program)
 		a.mu.Unlock()
 		log.Info("Cleared app program stores.")
@@ -115,20 +154,17 @@ func (a *app) Start() {
 
 func (a *app) ProgramHandler(s ssh.Session) *tea.Program {
 	userID := s.User()
-	m := initialModel(a, userID) // initialModel returns *model
+	m := initialModel(a, userID)
-
 	p := tea.NewProgram(m, bubbletea.MakeOptions(s)...)
-
+	m.SetProgram(p)
-	m.SetProgram(p)         // Give model a reference to its own program
+	a.addProgram(p, userID)
-	a.addProgram(p, userID) // Register program with the app, passing userID
-
 	return p
 }
 
 func (a *app) addProgram(p *tea.Program, userID string) {
 	a.mu.Lock()
 	defer a.mu.Unlock()
-	a.activePrograms[p] = programInfo{roomName: "", userID: userID} // Store userID
+	a.activePrograms[p] = programInfo{roomName: "", userID: userID}
 	log.Debugf("Program %p added for user %s (unassigned)", p, userID)
 }
 
@@ -138,12 +174,9 @@ func (a *app) AssignProgramToRoom(prog *tea.Program, roomName string) {
 
 	currentInfo, exists := a.activePrograms[prog]
 	if !exists {
-		// userID is NOT available here if the program isn't in activePrograms.
-		// NO prog.Model() call should be here.
 		log.Warnf("Program %p not found in activePrograms, cannot assign to room %s", prog, roomName)
 		return
 	}
-	// Correct: UserID is obtained from the stored currentInfo.
 	userID := currentInfo.userID
 	currentRoom := currentInfo.roomName
 
@@ -159,7 +192,6 @@ func (a *app) AssignProgramToRoom(prog *tea.Program, roomName string) {
 		} else {
 			delete(a.rooms, currentRoom)
 		}
-		// Correct: userID is used from the variable. NO prog.Model() here.
 		log.Debugf("Program %p (user: %s) moved from room %s", prog, userID, currentRoom)
 	}
 
@@ -179,25 +211,85 @@ func (a *app) AssignProgramToRoom(prog *tea.Program, roomName string) {
 	}
 
 	a.activePrograms[prog] = programInfo{roomName: roomName, userID: userID}
-	// Correct: userID is used from the variable. NO prog.Model() here.
-	// If your error is on this line (or similar), ensure 'userID' is not derived from prog.Model().
 	log.Infof("Program %p assigned/confirmed in room %s (user: %s)", prog, roomName, userID)
 }
 
+func (a *app) HandleUserExitFromRoom(prog *tea.Program, oldRoomName string) {
+	a.mu.Lock()
+
+	info, exists := a.activePrograms[prog]
+	if !exists {
+		a.mu.Unlock()
+		log.Warnf("Program %p tried to exit room but not found in activePrograms.", prog)
+		return
+	}
+	if info.roomName == "" || info.roomName != oldRoomName {
+		a.mu.Unlock()
+		log.Warnf("Program %p (user: %s) tried to exit room %s but was not in it (actual: '%s').", prog, info.userID, oldRoomName, info.roomName)
+		return
+	}
+	userID := info.userID
+
+	leaveMsgText := fmt.Sprintf("%s has left the room.", userID)
+	var recipients []*tea.Program
+	if progs, ok := a.rooms[info.roomName]; ok {
+		recipients = make([]*tea.Program, 0, len(progs))
+		for _, p := range progs {
+			if p != prog {
+				recipients = append(recipients, p)
+			}
+		}
+	}
+
+	var updatedRoomProgs []*tea.Program
+	for _, pInRoom := range a.rooms[info.roomName] {
+		if pInRoom != prog {
+			updatedRoomProgs = append(updatedRoomProgs, pInRoom)
+		}
+	}
+	if len(updatedRoomProgs) > 0 {
+		a.rooms[info.roomName] = updatedRoomProgs
+	} else {
+		delete(a.rooms, info.roomName)
+	}
+	a.activePrograms[prog] = programInfo{roomName: "", userID: userID}
+	a.mu.Unlock() // Unlock before broadcasting
+
+	if len(recipients) > 0 {
+		for _, p := range recipients {
+			go p.Send(chatMsg{id: "system", text: leaveMsgText, room: oldRoomName, isSystemMessage: true})
+		}
+		log.Infof("Broadcasted leave message for user %s from room %s (commanded exit)", userID, oldRoomName)
+	}
+	log.Infof("Program %p (user: %s) exited room %s via /exit command.", prog, userID, oldRoomName)
+}
+
 func (a *app) RemoveProgram(prog *tea.Program) {
 	a.mu.Lock()
-	defer a.mu.Unlock()
 
 	info, exists := a.activePrograms[prog]
 	if !exists {
-		// Correct: NO prog.Model() call here.
+		a.mu.Unlock()
 		log.Warnf("Attempted to remove unknown or already removed program: %p", prog)
 		return
 	}
-	// Correct: UserID and roomName are obtained from the stored info.
 	userID := info.userID
 	roomName := info.roomName
 
+	var leaveRecipients []*tea.Program
+	var leaveMsgText string
+	if roomName != "" {
+		leaveMsgText = fmt.Sprintf("%s has disconnected.", userID)
+		if progs, ok := a.rooms[roomName]; ok {
+			leaveRecipients = make([]*tea.Program, 0, len(progs))
+			for _, p := range progs {
+				if p != prog {
+					leaveRecipients = append(leaveRecipients, p)
+				}
+			}
+		}
+	}
+
 	if roomName != "" {
 		var updatedRoomProgs []*tea.Program
 		for _, pInRoom := range a.rooms[roomName] {
@@ -210,26 +302,30 @@ func (a *app) RemoveProgram(prog *tea.Program) {
 		} else {
 			delete(a.rooms, roomName)
 		}
-		// Correct: userID is used from the variable. NO prog.Model() here.
+		log.Infof("Program %p (user: %s) removed from room %s tracking (full disconnect).", prog, userID, roomName)
-		// If your error is on this line, ensure 'userID' is not derived from prog.Model().
-		log.Infof("Program %p removed from room %s (user: %s)", prog, roomName, userID)
-	} else {
-		// Correct: userID is used from the variable. NO prog.Model() here.
-		log.Debugf("Program %p was not in any room, removing from active list (user: %s).", prog, userID)
 	}
-
 	delete(a.activePrograms, prog)
-	// Correct: userID is used from the variable. NO prog.Model() here.
+	a.mu.Unlock() // Unlock before broadcasting
-	log.Infof("Program %p removed from active list (user: %s).", prog, userID)
+
+	if roomName != "" && len(leaveRecipients) > 0 {
+		for _, p := range leaveRecipients {
+			go p.Send(chatMsg{id: "system", text: leaveMsgText, room: roomName, isSystemMessage: true})
+		}
+		log.Infof("Broadcasted disconnect message for user %s from room %s.", userID, roomName)
+	}
+	log.Infof("Program %p (user: %s) removed from active list (full disconnect).", prog, userID)
 }
 
+// --- main, model types, and functions continue below ---
 func main() {
-	// For more detailed logs:
+	// Optional: Setup file logging
-	// file, err := os.OpenFile("chat_server.log", os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
+	// file, err := os.OpenFile("chatter_server.log", os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
 	// if err == nil {
-	// 	log.Default().SetOutput(file)
+	// 	defer file.Close()
+	// 	log.Default().SetOutput(io.MultiWriter(os.Stderr, file)) // Log to both stderr and file
 	// }
-	// log.SetLevel(log.DebugLevel) // Show debug messages
+	// log.SetLevel(log.DebugLevel)
+	// log.SetTimeFormat(time.Kitchen)
 
 	app := newApp()
 	app.Start()
@@ -238,9 +334,10 @@ func main() {
 type (
 	errMsg  error
 	chatMsg struct {
-		id   string
+		id              string
-		text string
+		text            string
-		room string // So message carries its intended room
+		room            string
+		isSystemMessage bool
 	}
 )
 
@@ -258,30 +355,29 @@ func isValidRoomName(name string) bool {
 }
 
 type model struct {
-	app         *app         // Pointer to the main application state
+	app                *app
-	program     *tea.Program // Pointer to this model's own tea.Program
+	program            *tea.Program
-	phase       modelPhase
+	phase              modelPhase
-	roomInput   textinput.Model
+	roomInput          textinput.Model
-	currentRoom string
+	currentRoom        string
-	id          string // User ID from SSH session
+	id                 string // User ID from SSH session
-
+	viewport           viewport.Model
-	// Chat fields
+	messages           []string
-	viewport    viewport.Model
+	textarea           textarea.Model
-	messages    []string
+	senderStyle        lipgloss.Style
-	textarea    textarea.Model
+	systemMessageStyle lipgloss.Style // For join/leave/command output
-	senderStyle lipgloss.Style
+	err                error
-	err         error
 }
 
 func initialModel(appInstance *app, userID string) *model {
 	roomTi := textinput.New()
-	roomTi.Placeholder = "e.g., 'dev-chat' or 'gaming'"
+	roomTi.Placeholder = "e.g., 'dev-chat' or 'random' (or press Enter for random)"
 	roomTi.Focus()
-	roomTi.CharLimit = 24
+	roomTi.CharLimit = maxRoomNameLength
-	roomTi.Width = 40
+	roomTi.Width = 50
 
 	ta := textarea.New()
-	ta.Placeholder = "Type your message here and press Enter..."
+	ta.Placeholder = "Type your message or a command (e.g. /help)..."
 	ta.Prompt = "┃ "
 	ta.CharLimit = 280
 	ta.SetWidth(60)
@@ -294,15 +390,16 @@ func initialModel(appInstance *app, userID string) *model {
 	vp.SetContent("Welcome! Please enter a room name to join or create.")
 
 	return &model{
-		app:         appInstance,
+		app:                appInstance,
-		id:          userID,
+		id:                 userID,
-		phase:       phaseRoomInput,
+		phase:              phaseRoomInput,
-		roomInput:   roomTi,
+		roomInput:          roomTi,
-		textarea:    ta,
+		textarea:           ta,
-		messages:    []string{},
+		messages:           []string{},
-		viewport:    vp,
+		viewport:           vp,
-		senderStyle: lipgloss.NewStyle().Foreground(lipgloss.Color("12")),
+		senderStyle:        lipgloss.NewStyle().Foreground(lipgloss.Color("12")),                                      // Blue for sender
-		err:         nil,
+		systemMessageStyle: lipgloss.NewStyle().Foreground(lipgloss.AdaptiveColor{Light: "#909090", Dark: "#6C6C6C"}), // Grey for system
+		err:                nil,
 	}
 }
 
@@ -337,23 +434,37 @@ func (m *model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 				return m, tea.Quit
 			case tea.KeyEnter:
 				roomName := strings.TrimSpace(m.roomInput.Value())
+				if roomName == "" {
+					roomName = generateRandomRoomName(maxRoomNameLength)
+					m.roomInput.SetValue(roomName) // Show generated name in input briefly
+					log.Infof("User %s generated random room name: %s", m.id, roomName)
+				}
+
 				if isValidRoomName(roomName) {
 					m.currentRoom = roomName
 					m.phase = phaseChat
 					m.err = nil
 
 					if m.program == nil {
-						log.Error("CRITICAL: m.program is nil when trying to assign to room. This should not happen.")
+						log.Error("CRITICAL: m.program is nil when trying to assign to room.")
 						m.err = errors.New("internal setup error: program reference missing")
 						return m, nil
 					}
 					m.app.AssignProgramToRoom(m.program, m.currentRoom)
 
+					// Send join message
+					joinMsgText := fmt.Sprintf("%s has joined the room.", m.id)
+					m.app.BroadcastToRoom(m.currentRoom, chatMsg{
+						id:              "system",
+						text:            joinMsgText,
+						room:            m.currentRoom,
+						isSystemMessage: true,
+					}, nil) // Broadcast to all, including self
+
 					m.textarea.Focus()
 					m.roomInput.Blur()
-
+					m.viewport.SetContent(fmt.Sprintf("Welcome to room '%s'! (User: %s)\nType /help for commands.", m.currentRoom, m.id))
-					m.viewport.SetContent(fmt.Sprintf("Welcome to room '%s'! (User: %s)\nMessages will appear below.", m.currentRoom, m.id))
+					m.messages = []string{} // Clear previous messages if any
-					m.messages = []string{}
 					cmds = append(cmds, textarea.Blink)
 				} else {
 					m.err = errors.New("invalid room name: 1-24 alphanumeric characters required")
@@ -366,7 +477,6 @@ func (m *model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 			m.err = msg
 			return m, nil
 		}
-
 		m.roomInput, cmd = m.roomInput.Update(msg)
 		cmds = append(cmds, cmd)
 
@@ -388,31 +498,96 @@ func (m *model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 				return m, tea.Quit
 			case tea.KeyEnter:
 				val := strings.TrimSpace(m.textarea.Value())
-				if val != "" {
+				m.textarea.Reset() // Reset textarea immediately
-					m.app.BroadcastToRoom(m.currentRoom, chatMsg{
+
-						id:   m.id,
+				if strings.HasPrefix(val, "/") {
-						text: val,
+					parts := strings.Fields(val)
-						room: m.currentRoom,
+					command := parts[0]
-					})
+					// args := parts[1:] // For future use if commands take arguments
-					m.textarea.Reset()
+
+					switch command {
+					case "/help":
+						helpText := m.systemMessageStyle.Render(`Available commands:
+  /help          Show this help message.
+  /exit          Leave the current room.
+  /users         List users in this room.`)
+						m.messages = append(m.messages, helpText)
+					case "/exit", "/leave":
+						if m.program != nil && m.currentRoom != "" {
+							// App handles broadcasting leave message
+							m.app.HandleUserExitFromRoom(m.program, m.currentRoom)
+						}
+						m.phase = phaseRoomInput
+						m.currentRoom = ""
+						m.messages = []string{}
+						m.err = nil
+						m.roomInput.Reset()
+						m.roomInput.Placeholder = "Enter another room or press Enter for random..." // Update placeholder
+						m.roomInput.Focus()
+						m.textarea.Blur()
+						m.viewport.SetContent("You have left the room. Enter a new room name.")
+						cmds = append(cmds, textinput.Blink)
+						// No further processing for this Update cycle after state change.
+						return m, tea.Batch(cmds...)
+					case "/users":
+						if m.currentRoom != "" {
+							userIDs := m.app.GetUserIDsInRoom(m.currentRoom)
+							var userListText string
+							if len(userIDs) > 0 {
+								userListText = "Users in '" + m.currentRoom + "':\n  " + strings.Join(userIDs, "\n  ")
+							} else {
+								userListText = "You are the only one here." // Should not happen if self is listed
+							}
+							m.messages = append(m.messages, m.systemMessageStyle.Render(userListText))
+						} else {
+							m.messages = append(m.messages, m.systemMessageStyle.Render("Error: Not in a room."))
+						}
+					default:
+						m.messages = append(m.messages, m.systemMessageStyle.Render("Unknown command: "+command))
+					}
+					m.viewport.SetContent(strings.Join(m.messages, "\n"))
 					m.viewport.GotoBottom()
+
+				} else if val != "" { // Regular chat message
+					m.app.BroadcastToRoom(m.currentRoom, chatMsg{
+						id:              m.id,
+						text:            val,
+						room:            m.currentRoom,
+						isSystemMessage: false,
+					}, nil) // Broadcast to all, including self (sender)
+					// Message will be added to m.messages when this client receives its own broadcast.
 				}
+				// After sending or command, ensure textarea is focused for next input
+				cmds = append(cmds, m.textarea.Focus(), textarea.Blink)
 			}
-		case chatMsg:
+
-			if msg.room == m.currentRoom {
+		case chatMsg: // Received a message (could be own, from others, or system)
-				styledId := m.senderStyle.Render(msg.id)
+			if msg.room == m.currentRoom || (msg.isSystemMessage && msg.room == m.currentRoom) { // Ensure for current room
-				if msg.id == m.id {
+				var styledMessage string
-					styledId = lipgloss.NewStyle().Foreground(lipgloss.Color("10")).Render(msg.id + " (you)")
+				if msg.isSystemMessage {
+					styledMessage = m.systemMessageStyle.Render(msg.text)
+				} else {
+					styledId := m.senderStyle.Render(msg.id)
+					if msg.id == m.id { // Own regular message
+						// Optionally style self messages differently, e.g., if not relying on system message for own display
+						// For now, senderStyle applies to all non-system messages.
+					}
+					styledMessage = styledId + ": " + msg.text
 				}
-				m.messages = append(m.messages, styledId+": "+msg.text)
+				m.messages = append(m.messages, styledMessage)
 				m.viewport.SetContent(strings.Join(m.messages, "\n"))
 				m.viewport.GotoBottom()
-			} else {
+			} else if msg.room != "" { // Message for a different room, should ideally not happen if app logic is correct
-				log.Warnf("Client in room '%s' (user '%s') incorrectly received msg for room '%s'. Ignored.", m.currentRoom, m.id, msg.room)
+				log.Warnf("Client '%s' in room '%s' received message for different room '%s': %s", m.id, m.currentRoom, msg.room, msg.text)
 			}
+
 		case errMsg:
 			m.err = msg
 			log.Error("Error in chat model", "user", m.id, "room", m.currentRoom, "error", msg)
+			// Optionally display this error in the viewport too
+			// m.messages = append(m.messages, m.systemMessageStyle.Render("An error occurred: "+msg.Error()))
+			// m.viewport.SetContent(strings.Join(m.messages, "\n"))
+			// m.viewport.GotoBottom()
 			return m, nil
 		}
 	}
@@ -423,21 +598,23 @@ func (m *model) View() string {
 	var viewBuilder strings.Builder
 
 	titleStyle := lipgloss.NewStyle().Bold(true).Foreground(lipgloss.Color("63"))
-	errorStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("9")).Bold(true)
+	errorStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("9")).Bold(true) // Red
 
 	switch m.phase {
 	case phaseRoomInput:
 		viewBuilder.WriteString(titleStyle.Render("chatter.sh") + "\n\n")
 		viewBuilder.WriteString("Please enter a room name to join or create.\n")
-		viewBuilder.WriteString("Rules: 1-24 characters, letters (a-z, A-Z) and numbers (0-9) only.\n\n")
+		viewBuilder.WriteString("Rules: 1-24 characters, letters (a-z, A-Z) and numbers (0-9) only.\n")
+		viewBuilder.WriteString("(Leave empty and press Enter for a random room name.)\n\n")
 		viewBuilder.WriteString(m.roomInput.View())
 		viewBuilder.WriteString("\n\n(Ctrl+C or Esc to quit)")
 	case phaseChat:
-		header := fmt.Sprintf("Room: %s / User: %s", m.currentRoom, m.id)
+		header := fmt.Sprintf("Room: %s / User: %s (type /help for commands)", m.currentRoom, m.id)
 		viewBuilder.WriteString(titleStyle.Render(header) + "\n")
-		viewBuilder.WriteString(lipgloss.NewStyle().BorderStyle(lipgloss.NormalBorder()).BorderBottom(true).Width(m.viewport.Width).Render(m.viewport.View()) + "\n\n")
+		// Simple border for viewport content
+		viewBuilder.WriteString(lipgloss.NewStyle().BorderStyle(lipgloss.RoundedBorder()).Padding(0, 1).Render(m.viewport.View()) + "\n")
 		viewBuilder.WriteString(m.textarea.View())
-		viewBuilder.WriteString("\n\n(Ctrl+C or Esc to quit)")
+		viewBuilder.WriteString("\n(Ctrl+C or Esc to quit)")
 	default:
 		viewBuilder.WriteString("Unknown application state. This is a bug.")
 	}