package main

import (
	"bufio"
	"encoding/binary"
	"fmt"
	"net"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"
)

type Client struct {
	socket   net.Conn
	username string
	inbox    chan string
}

type ActiveClients struct {
	mu      sync.RWMutex
	clients map[string]*Client
}

func (ac *ActiveClients) add(c *Client) bool {
	ac.mu.Lock()
	defer ac.mu.Unlock()
	if ac.clients == nil {
		ac.clients = make(map[string]*Client, 0)
	}
	if c.inbox == nil {
		c.inbox = make(chan string, 100)
	}

	for username := range ac.clients {
		if username == c.username {
			return false
		}
	}
	ac.clients[c.username] = c
	return true
}

func (ac *ActiveClients) remove(c Client) bool {
	ac.mu.Lock()
	defer ac.mu.Unlock()
	debug("Removing %v\n", c.username)

	if ac.clients == nil {
		return false
	}

	_, ok := ac.clients[c.username]

	if ok {
		delete(ac.clients, c.username)
	}
	return ok
}

func (ac *ActiveClients) forwardMessage(message string, from string) {
	ac.mu.RLock()
	defer ac.mu.RUnlock()
	msgWithFrom := from + ": " + message

	for username, c := range ac.clients {
		if username != from {
			debug("Attempting to forward to %v...", username)
			select {
			case c.inbox <- msgWithFrom:
				debug("Forwarding message to %v\n", username)
			default:
				// if buffer is full, drop the message
				debug("Message dropped because %v inbox is full.\n", c.username)
			}
		}
	}
}

var wg sync.WaitGroup
var shouldDebug bool = false

func debug(message string, vals ...any) {
	if shouldDebug {
		fmt.Printf(message, vals...)
	}
}

func uint16ToBytes(i uint16) []byte {
	// return []byte{byte(i >> 8), byte(i & 0xFF)}
	b := make([]byte, 2)
	binary.BigEndian.PutUint16(b, i)
	return b
}

func bytesToUint16(b []byte) uint16 {
	return binary.BigEndian.Uint16(b[:2])
}

func sendMessage(message string, connection net.Conn, milliseconds int) (int, error) {
	var timeout time.Duration
	msg_bytes := []byte(message)
	msg_size := len(msg_bytes)
	msg := append(uint16ToBytes(uint16(msg_size)), msg_bytes...)
	if milliseconds > 0 {
		timeout = time.Millisecond * time.Duration(milliseconds)
		connection.SetWriteDeadline(time.Now().Add(timeout))
		defer connection.SetWriteDeadline(time.Time{})
	}
	size, err := connection.Write(msg)
	if err != nil {
		if !strings.Contains(err.Error(), "i/o timeout") {
			fmt.Println(err)
		}
	}
	debug("Sent %v bytes to %v\n", size, connection.RemoteAddr())
	return size, err
}

func receiveMessage(connection net.Conn, milliseconds int) (string, int, error) {
	var msg_size [2]byte
	var msg []byte
	var timeout time.Duration
	if milliseconds > 0 {
		timeout = time.Millisecond * time.Duration(milliseconds)
		connection.SetReadDeadline(time.Now().Add(timeout))
		defer connection.SetReadDeadline(time.Time{})
	}
	size, err := connection.Read(msg_size[:])
	if err != nil {
		if !strings.Contains(err.Error(), "i/o timeout") {
			fmt.Println(err)
		}
	}
	if err != nil || size == 0 {
		return "", size, err
	}
	size = int(bytesToUint16(msg_size[:]))
	msg = make([]byte, size)
	size, err = connection.Read(msg)
	if err != nil {
		if !strings.Contains(err.Error(), "i/o timeout") {
			fmt.Println(err)
		}
	}
	if err != nil || size == 0 {
		return "", size, err
	}

	debug("Received message (%v bytes) from %v\n", size+2, connection.RemoteAddr())
	return string(msg), size, err
}

func server(numberOfClients int, hostInterface string) {
	shouldDebug = true
	users := ActiveClients{}
	listener, err := net.Listen("tcp", hostInterface)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	for i := 0; i < numberOfClients; i++ {
		wg.Add(1)
		go listen(listener, &users)
	}

	wg.Wait()
}

func listen(listener net.Listener, users *ActiveClients) {
	defer wg.Done()
	var username string
	var thisClient Client
	var lg sync.WaitGroup
	signalShouldEnd := make(chan bool, 1)

loop:
	connection, err := listener.Accept()
	if err != nil {
		debug(err.Error() + "\n")
		return
	}

	debug("Received connection from %v\n", connection.RemoteAddr())

	hasSetUsername := false
	lg.Add(2)
	go func() {
		defer lg.Done()
		for {
			// wait up to 100ms for the client to send a message
			message, size, err := receiveMessage(connection, 100)
			if err != nil {
				if strings.Contains(err.Error(), "i/o timeout") {
					// timed out, so go back to start of loop
					// fmt.Println("i/o timeout")
					continue
				}
				debug(err.Error())
				select {
				case signalShouldEnd <- true:
					if users.remove(thisClient) {
						users.forwardMessage("Goodbye "+thisClient.username, "server")
					}
				default:
				}
				return
			}

			if size > 0 {
				if !hasSetUsername {
					username = message
					if username == "server" {
						sendMessage("server: error: username is invalid; please send new username", connection, 0)
						break
					}
					thisClient = Client{
						socket:   connection,
						username: username,
					}
					ok := users.add(&thisClient)
					if !ok {
						sendMessage("server: error: username is invalid; please send new username", connection, 0)
						break
					} else {
						users.forwardMessage("Welcome "+username, "server")
					}
					hasSetUsername = true
				} else {
					users.forwardMessage(message, username)
				}
			}

			select {
			case <-signalShouldEnd:
				return
			default:
			}
		}
	}()

	go func() {
		defer lg.Done()
		for {
			select {
			case incoming := <-thisClient.inbox:
				debug("Message found in inbox; sending to %v\n", thisClient.username)
				_, err := sendMessage(incoming, thisClient.socket, 0)
				if err != nil {
					if !strings.Contains(err.Error(), "i/o timeout") {
						debug(err.Error() + "\n")
						select {
						case signalShouldEnd <- true:
							if users.remove(thisClient) {
								users.forwardMessage("Goodbye "+thisClient.username, "server")
							}
						default:
						}
						return
					}
					continue
				}
			case <-signalShouldEnd:
				return
			default:
				time.Sleep(time.Millisecond * 10)
			}
		}
	}()

	lg.Wait()
	debug("Recycling server socket for new client\n")
	goto loop
}

func client(host string, username string) {
	connection, err := net.Dial("tcp", host+":1337")
	endSignals := []string{"bye", "quit", "exit"}

	if err != nil {
		debug("%v\n", err)
		return
	}

	go func() {
		for {
			message, size, err := receiveMessage(connection, 100)

			if err != nil {
				if strings.Contains(err.Error(), "i/o timeout") {
					continue
				}
				debug("%v\n", err)
				os.Exit(0)
			}

			if size == 0 {
				continue
			}

			// fmt.Printf("Received %v bytes: ", size)
			fmt.Println(message)
		}
	}()

	fmt.Println("Connected to server.")

	_, err = sendMessage(username, connection, 0)
	if err != nil {
		fmt.Println(err)
		return
	}

	for {
		message := askInput("")
		if contains(endSignals, message) {
			return
		}
		_, err := sendMessage(message, connection, 0)
		if err != nil {
			fmt.Println(err)
			return
		}
	}
}

func benchmark(host string, numberOfClients int, numberOfMessages int, baseMessage string) {
	wg.Add(numberOfClients)
	for i := 0; i < numberOfClients; i++ {
		go func(i int) {
			defer wg.Done()
			var bg sync.WaitGroup
			username := "benchmark" + strconv.Itoa(i)
			connection, err := net.Dial("tcp", host+":1337")
			if err != nil {
				fmt.Println(err)
				return
			}
			// current limitation: must wait for username to process
			_, _ = sendMessage(username, connection, 100)
			time.Sleep(time.Millisecond * 100)

			bg.Add(numberOfMessages)
			for d := 0; d < numberOfMessages; d++ {
				go func(message string) {
					defer bg.Done()
					_, err = sendMessage(message, connection, 100)
					if err != nil {
						fmt.Println(err)
						return
					}
				}(baseMessage + strconv.Itoa(d))
			}

			bg.Add(1)
			go func() {
				defer bg.Done()
				for {
					message, size, err := receiveMessage(connection, 2000)
					if err != nil {
						fmt.Println(err)
						return
					}
					if size == 0 {
						fmt.Println("aborting from empty receive")
						return
					}
					fmt.Println(message)
				}
			}()

			bg.Wait()
		}(i)
	}
	wg.Wait()
}

func askInput(query string) string {
	if query != "" {
		fmt.Printf("%v ", query)
	}
	in := bufio.NewReader(os.Stdin)
	data, _ := in.ReadString('\n')
	data, _, _ = strings.Cut(data, "\n")
	data, _, _ = strings.Cut(data, "\r")
	return data
}

func contains[T comparable](list []T, query T) bool {
	// contains function adapted from https://stackoverflow.com/a/10485970
	for _, item := range list {
		if item == query {
			return true
		}
	}
	return false
}

func usage(programName string) {
	fmt.Printf("usage: %v server numerOfThreads [hostInterface]\n", programName)
	fmt.Println("\thostInterface is of form hostname:port or IP:port")
	fmt.Println("\tuse 0.0.0.0:port to listen on every available IP interface")
	fmt.Println("\tdefault is localhost:1337")
	fmt.Printf("usage: %v client hostIP username\n", programName)
	fmt.Printf("usage: %v benchmark hostIP numberOfClients numberOfMessages baseMessage\n", programName)
	fmt.Println("arguments must be in order shown")
}

func main() {
	programName := os.Args[0]
	nameParts := strings.Split(programName, "\\")
	programName = nameParts[len(nameParts)-1]

	if len(os.Args) < 2 {
		usage(programName)
		return
	}

	mode := os.Args[1]

	switch mode {
	case "server", "serve":
		var numberOfClients int
		if len(os.Args) < 3 {
			usage(programName)
			return
		}
		n, _ := strconv.ParseInt(os.Args[2], 10, 16)
		numberOfClients = int(n)
		var hostInterface string
		if len(os.Args) > 3 {
			hostInterface = os.Args[3]
		} else {
			hostInterface = "localhost:1337"
		}
		server(numberOfClients, hostInterface)
	case "client", "connect":
		var host string
		var username string
		if len(os.Args) < 4 {
			usage(programName)
			return
		}
		host = os.Args[2]
		username = os.Args[3]
		client(host, username)
	case "benchmark":
		var host string
		var nClients string
		var numberOfClients int
		var nMsgs string
		var numberOfMessages int
		var message string
		if len(os.Args) < 6 {
			usage(programName)
			return
		}
		host = os.Args[2]
		nClients = os.Args[3]
		nMsgs = os.Args[4]
		message = os.Args[5]
		n, _ := strconv.ParseInt(nClients, 10, 16)
		numberOfClients = int(n)
		n, _ = strconv.ParseInt(nMsgs, 10, 16)
		numberOfMessages = int(n)
		benchmark(host, numberOfClients, numberOfMessages, message)
	default:
		usage(programName)
	}
}

// ISC License

// Copyleft (c) 2023 k98kurz

// Permission to use, copy, modify, and/or distribute this software
// for any purpose with or without fee is hereby granted, provided
// that the above copyleft notice and this permission notice appear in
// all copies.

// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL
// WARRANTIES WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED
// WARRANTIES OF MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE
// AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, INDIRECT, OR
// CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS
// OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT,
// NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
// CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.