chatter/main.go

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"
	"time"

	"github.com/charmbracelet/bubbles/textarea"
	"github.com/charmbracelet/bubbles/textinput"
	"github.com/charmbracelet/bubbles/viewport"
	tea "github.com/charmbracelet/bubbletea"
	"github.com/charmbracelet/lipgloss"
	"github.com/charmbracelet/log"
	"github.com/charmbracelet/ssh"
	"github.com/charmbracelet/wish"
	"github.com/charmbracelet/wish/activeterm"
	"github.com/charmbracelet/wish/bubbletea"
	"github.com/charmbracelet/wish/logging"
	"github.com/muesli/termenv"
)

const (
	host              = "0.0.0.0"
	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
}

// app contains a wish server and manages chat rooms and programs.
type app struct {
	*ssh.Server
	activePrograms map[*tea.Program]programInfo // program -> programInfo
	rooms          map[string][]*tea.Program    // roomName -> list of programs in that room
	mu             sync.Mutex
}

// BroadcastToRoom sends a message to all running programs in a specific room,
// 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, excluding %p): %+v", roomName, len(progsInRoom), excludeProg, msg)
		for _, p := range progsInRoom {
			if excludeProg != nil && p == excludeProg {
				continue
			}
			go p.Send(msg)
		}
	} else {
		log.Warnf("Attempted to broadcast to non-existent or empty room: %s", roomName)
	}
}

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),
		rooms:          make(map[string][]*tea.Program),
	}
	s, err := wish.NewServer(
		wish.WithAddress(net.JoinHostPort(host, port)),
		wish.WithHostKeyPath(".ssh/id_ed25519"),
		wish.WithMiddleware(
			bubbletea.MiddlewareWithProgramHandler(a.ProgramHandler, termenv.ANSI256),
			activeterm.Middleware(),
			logging.Middleware(),
		),
	)
	if err != nil {
		log.Fatal("Could not initialize server", "error", err)
	}

	a.Server = s
	return a
}

func (a *app) Start() {
	var err error
	done := make(chan os.Signal, 1)
	signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
	log.Info("Starting SSH server", "host", host, "port", port)
	go func() {
		if err = a.ListenAndServe(); err != nil && !errors.Is(err, ssh.ErrServerClosed) {
			log.Error("Could not start server", "error", err)
			done <- syscall.SIGTERM
		}
	}()

	<-done
	log.Info("Stopping SSH server...")
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer func() {
		cancel()
		a.mu.Lock()
		a.activePrograms = make(map[*tea.Program]programInfo)
		a.rooms = make(map[string][]*tea.Program)
		a.mu.Unlock()
		log.Info("Cleared app program stores.")
	}()

	if err := a.Shutdown(ctx); err != nil {
		log.Error("Could not gracefully stop server", "error", err)
	}
	log.Info("SSH server stopped.")
}

func (a *app) ProgramHandler(s ssh.Session) *tea.Program {
	userID := s.User()
	m := initialModel(a, userID)
	p := tea.NewProgram(m, bubbletea.MakeOptions(s)...)
	m.SetProgram(p)
	a.addProgram(p, 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}
	log.Debugf("Program %p added for user %s (unassigned)", p, userID)
}

func (a *app) AssignProgramToRoom(prog *tea.Program, roomName string) {
	a.mu.Lock()
	defer a.mu.Unlock()

	currentInfo, exists := a.activePrograms[prog]
	if !exists {
		log.Warnf("Program %p not found in activePrograms, cannot assign to room %s", prog, roomName)
		return
	}
	userID := currentInfo.userID
	currentRoom := currentInfo.roomName

	if currentRoom != "" && currentRoom != roomName {
		var updatedRoomProgs []*tea.Program
		for _, pInRoom := range a.rooms[currentRoom] {
			if pInRoom != prog {
				updatedRoomProgs = append(updatedRoomProgs, pInRoom)
			}
		}
		if len(updatedRoomProgs) > 0 {
			a.rooms[currentRoom] = updatedRoomProgs
		} else {
			delete(a.rooms, currentRoom)
		}
		log.Debugf("Program %p (user: %s) moved from room %s", prog, userID, currentRoom)
	}

	if _, ok := a.rooms[roomName]; !ok {
		a.rooms[roomName] = []*tea.Program{}
	}

	foundInNewRoomList := false
	for _, pInList := range a.rooms[roomName] {
		if pInList == prog {
			foundInNewRoomList = true
			break
		}
	}
	if !foundInNewRoomList {
		a.rooms[roomName] = append(a.rooms[roomName], prog)
	}

	a.activePrograms[prog] = programInfo{roomName: roomName, userID: userID}
	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()

	info, exists := a.activePrograms[prog]
	if !exists {
		a.mu.Unlock()
		log.Warnf("Attempted to remove unknown or already removed program: %p", prog)
		return
	}
	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] {
			if pInRoom != prog {
				updatedRoomProgs = append(updatedRoomProgs, pInRoom)
			}
		}
		if len(updatedRoomProgs) > 0 {
			a.rooms[roomName] = updatedRoomProgs
		} else {
			delete(a.rooms, roomName)
		}
		log.Infof("Program %p (user: %s) removed from room %s tracking (full disconnect).", prog, userID, roomName)
	}
	delete(a.activePrograms, prog)
	a.mu.Unlock() // Unlock before broadcasting

	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() {
	// Optional: Setup file logging
	// file, err := os.OpenFile("chatter_server.log", os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
	// if err == nil {
	// 	defer file.Close()
	// 	log.Default().SetOutput(io.MultiWriter(os.Stderr, file)) // Log to both stderr and file
	// }
	// log.SetLevel(log.DebugLevel)
	// log.SetTimeFormat(time.Kitchen)

	app := newApp()
	app.Start()
}

type (
	errMsg  error
	chatMsg struct {
		id              string
		text            string
		room            string
		isSystemMessage bool
	}
)

type modelPhase int

const (
	phaseRoomInput modelPhase = iota
	phaseChat
)

var roomNameRegex = regexp.MustCompile("^[a-zA-Z0-9]{1,24}$")

func isValidRoomName(name string) bool {
	return roomNameRegex.MatchString(name)
}

type model struct {
	app                *app
	program            *tea.Program
	phase              modelPhase
	roomInput          textinput.Model
	currentRoom        string
	id                 string // User ID from SSH session
	viewport           viewport.Model
	messages           []string
	textarea           textarea.Model
	senderStyle        lipgloss.Style
	systemMessageStyle lipgloss.Style // For join/leave/command output
	err                error
}

func initialModel(appInstance *app, userID string) *model {
	roomTi := textinput.New()
	roomTi.Placeholder = "e.g., 'dev-chat' or 'random' (or press Enter for random)"
	roomTi.Focus()
	roomTi.CharLimit = maxRoomNameLength
	roomTi.Width = 50

	ta := textarea.New()
	ta.Placeholder = "Type your message or a command (e.g. /help)..."
	ta.Prompt = "┃ "
	ta.CharLimit = 280
	ta.SetWidth(60)
	ta.SetHeight(3)
	ta.FocusedStyle.CursorLine = lipgloss.NewStyle()
	ta.ShowLineNumbers = false
	ta.KeyMap.InsertNewline.SetEnabled(false)

	vp := viewport.New(60, 15)
	vp.SetContent("Welcome! Please enter a room name to join or create.")

	return &model{
		app:                appInstance,
		id:                 userID,
		phase:              phaseRoomInput,
		roomInput:          roomTi,
		textarea:           ta,
		messages:           []string{},
		viewport:           vp,
		senderStyle:        lipgloss.NewStyle().Foreground(lipgloss.Color("12")),                                      // Blue for sender
		systemMessageStyle: lipgloss.NewStyle().Foreground(lipgloss.AdaptiveColor{Light: "#909090", Dark: "#6C6C6C"}), // Grey for system
		err:                nil,
	}
}

func (m *model) SetProgram(p *tea.Program) {
	m.program = p
}

func (m *model) Init() tea.Cmd {
	switch m.phase {
	case phaseRoomInput:
		return textinput.Blink
	case phaseChat:
		return textarea.Blink
	default:
		return nil
	}
}

func (m *model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
	var cmds []tea.Cmd
	var cmd tea.Cmd

	switch m.phase {
	case phaseRoomInput:
		switch msg := msg.(type) {
		case tea.KeyMsg:
			switch msg.Type {
			case tea.KeyCtrlC, tea.KeyEsc:
				if m.program != nil {
					m.app.RemoveProgram(m.program)
				}
				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.")
						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.messages = []string{} // Clear previous messages if any
					cmds = append(cmds, textarea.Blink)
				} else {
					m.err = errors.New("invalid room name: 1-24 alphanumeric characters required")
					m.roomInput.Focus()
					cmds = append(cmds, textinput.Blink)
				}
				return m, tea.Batch(cmds...)
			}
		case errMsg:
			m.err = msg
			return m, nil
		}
		m.roomInput, cmd = m.roomInput.Update(msg)
		cmds = append(cmds, cmd)

	case phaseChat:
		var tiCmd tea.Cmd
		var vpCmd tea.Cmd

		m.textarea, tiCmd = m.textarea.Update(msg)
		m.viewport, vpCmd = m.viewport.Update(msg)
		cmds = append(cmds, tiCmd, vpCmd)

		switch msg := msg.(type) {
		case tea.KeyMsg:
			switch msg.Type {
			case tea.KeyCtrlC, tea.KeyEsc:
				if m.program != nil {
					m.app.RemoveProgram(m.program)
				}
				return m, tea.Quit
			case tea.KeyEnter:
				val := strings.TrimSpace(m.textarea.Value())
				m.textarea.Reset() // Reset textarea immediately

				if strings.HasPrefix(val, "/") {
					parts := strings.Fields(val)
					command := parts[0]
					// args := parts[1:] // For future use if commands take arguments

					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: // Received a message (could be own, from others, or system)
			if msg.room == m.currentRoom || (msg.isSystemMessage && msg.room == m.currentRoom) { // Ensure for current room
				var styledMessage string
				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, styledMessage)
				m.viewport.SetContent(strings.Join(m.messages, "\n"))
				m.viewport.GotoBottom()
			} else if msg.room != "" { // Message for a different room, should ideally not happen if app logic is correct
				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
		}
	}
	return m, tea.Batch(cmds...)
}

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) // 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")
		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 (type /help for commands)", m.currentRoom, m.id)
		viewBuilder.WriteString(titleStyle.Render(header) + "\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(Ctrl+C or Esc to quit)")
	default:
		viewBuilder.WriteString("Unknown application state. This is a bug.")
	}

	if m.err != nil {
		viewBuilder.WriteString("\n\n" + errorStyle.Render("Error: "+m.err.Error()))
	}

	return viewBuilder.String() + "\n"
}