dttt/cmd/chat.go

package cmd

import (
	"bufio"
	"fmt"
	"log"
	"net"
	"os"
	"sync"

	"github.com/spf13/cobra"
)

var (
	nickname string
	clients  []*client
	mu       sync.Mutex
)

type client struct {
	conn     net.Conn
	nickname string
}

// chatCmd represents the chat command
var chatCmd = &cobra.Command{
	Use:   "chat",
	Short: "Start a chat server or connect to other peers",
	Long: `Start a chat server or connect to other peers.
If no arguments are provided, it starts a chat server.
If an address is provided, it connects to the specified chat server.
You can specify your nickname using the --nickname (-n) flag.`,
	Args: cobra.MaximumNArgs(1),
	Run: func(cmd *cobra.Command, args []string) {
		if len(args) == 0 {
			startChatServer()
		} else {
			address := args[0]
			startChatClient(address)
		}
	},
}

func init() {
	rootCmd.AddCommand(chatCmd)
	chatCmd.Flags().StringVarP(&nickname, "nickname", "n", "", "Nickname for the user")
}

func startChatServer() {
	listener, err := net.Listen("tcp", "localhost:8080")
	if err != nil {
		log.Fatalf("Failed to start chat server: %v", err)
	}
	defer listener.Close()
	fmt.Println("Chat server started. Listening on localhost:8080")

	// Start a goroutine to handle server input
	go func() {
		scanner := bufio.NewScanner(os.Stdin)
		for scanner.Scan() {
			msg := scanner.Text()
			broadcastMessage(nil, "Server: "+msg+"\n")
		}
		if err := scanner.Err(); err != nil {
			log.Fatalf("Error reading server input: %v", err)
		}
	}()

	for {
		conn, err := listener.Accept()
		if err != nil {
			log.Printf("Failed to accept connection: %v", err)
			continue
		}
		go handleClient(conn)
	}
}

func handleClient(conn net.Conn) {
	defer conn.Close()

	reader := bufio.NewReader(conn)

	nickname, err := reader.ReadString('\n')
	if err != nil {
		log.Printf("Error reading nickname: %v", err)
		return
	}

	fmt.Printf("New connection from %s with nickname: %s", conn.RemoteAddr().String(), nickname)

	for {
		message, err := reader.ReadString('\n')
		if err != nil {
			log.Printf("Error reading message: %v", err)
			return
		}

		fmt.Printf("[%s]: %s", nickname, message)
		broadcastMessage(nil, fmt.Sprintf("[%s]: %s", nickname, message))
	}
}

func removeClient(c *client) {
	for i, cl := range clients {
		if cl == c {
			clients = append(clients[:i], clients[i+1:]...)
			return
		}
	}
}

func broadcastMessage(sender *client, message string) {
	mu.Lock()
	defer mu.Unlock()

	for _, cl := range clients {
		if cl != sender {
			_, err := cl.conn.Write([]byte(message))
			if err != nil {
				log.Printf("Error broadcasting message: %v", err)
			}
		}
	}
}

func startChatClient(address string) {
	conn, err := net.Dial("tcp", address)
	if err != nil {
		log.Fatalf("Failed to connect to chat server: %v", err)
	}
	defer conn.Close()

	fmt.Print("Enter your nickname: ")
	scanner := bufio.NewScanner(os.Stdin)
	if !scanner.Scan() {
		log.Fatal("Failed to read nickname from stdin")
	}
	nickname := scanner.Text() + ": "

	// Send nickname to server
	if _, err := fmt.Fprintf(conn, "%s\n", nickname); err != nil {
		log.Fatalf("Failed to send nickname to server: %v", err)
	}

	go func() {
		reader := bufio.NewReader(conn)
		for {
			message, err := reader.ReadString('\n')
			if err != nil {
				log.Printf("Error reading message: %v", err)
				return
			}
			fmt.Print(message)
		}
	}()

	for scanner.Scan() {
		msg := scanner.Text()
		if _, err := fmt.Fprintf(conn, "%s\n", msg); err != nil {
			log.Fatalf("Failed to send message to server: %v", err)
		}
	}
	if err := scanner.Err(); err != nil {
		log.Fatalf("Error reading standard input: %v", err)
	}
}