inp2025/stages/lobby.go

package stages

import (
	"fmt"
	"net"
	"strings"
	"time"

	"gitea.konchin.com/ytshih/inp2025/types"
	"gitea.konchin.com/ytshih/inp2025/utils"
	"gitea.konchin.com/ytshih/inp2025/workflows"
	tea "github.com/charmbracelet/bubbletea"
	"github.com/spf13/viper"
	"github.com/vmihailenco/msgpack/v5"
)

const (
	REFRESH_TIME = time.Second
)

type lobbyOperationType int

const (
	lobbyOperationChoose lobbyOperationType = iota
	lobbyOperationServerWaiting
	lobbyOperationServerChoose
	lobbyOperationClientScannning
)

type LobbyModel struct {
	*BaseModel
	op lobbyOperationType

	shutdown types.ShutdownFunc
	local    string
	remote   string

	// server
	remoteUser string
	listener   net.Listener

	// client
	cursor    int
	endpoints []string

	info string
	err  error
}

func NewLobbyModel(base *BaseModel) *LobbyModel {
	return &LobbyModel{
		BaseModel: base,
		op:        lobbyOperationChoose,

		shutdown: func() {},
	}
}

type serverListenMsg struct{}

func (m *LobbyModel) serverListen() tea.Cmd {
	return func() tea.Msg {
		dataCh := make(chan string)
		m.local, m.shutdown, m.err = utils.ListenUDPData(
			viper.GetInt("udp-listen-port"), dataCh)
		if m.err != nil {
			m.err = fmt.Errorf("failed to listen, %w", m.err)
			return serverListenMsg{}
		}

		req := <-dataCh
		m.shutdown()
		m.shutdown = func() {}

		var joinRequest types.JoinRequest
		if err := msgpack.Unmarshal([]byte(req), &joinRequest); err != nil {
			m.err = fmt.Errorf("failed to unmarshal msgpack, %w", err)
			return serverListenMsg{}
		}

		m.remote = joinRequest.Endpoint
		m.remoteUser = joinRequest.Username
		return serverListenMsg{}
	}
}

type serverSendReplyMsg struct{}

func (m *LobbyModel) serverSendReply(response bool) tea.Cmd {
	return func() tea.Msg {
		if response {
			// Start Wordle Server listener
			m.listener, m.err = net.Listen("tcp4", ":0")
			if m.err != nil {
				m.err = fmt.Errorf("failed to listen on anonymous port, %w", m.err)
				return serverSendReplyMsg{}
			}

			local := fmt.Sprintf("%s:%d",
				m.listener.Addr().(*net.TCPAddr).IP.String(),
				m.listener.Addr().(*net.TCPAddr).Port)
			m.err = utils.SendPayload(local, m.remote,
				types.JoinResponse{Endpoint: local})

			// Store wordle server endpoint
			m.remote = local
		} else {
			m.err = utils.SendPayload("", m.remote,
				types.JoinResponse{Endpoint: ""})
		}
		return serverSendReplyMsg{}
	}
}

type clientScanMsg time.Time

func (m *LobbyModel) clientScan() tea.Cmd {
	return tea.Tick(REFRESH_TIME, func(t time.Time) tea.Msg {
		m.endpoints = []string{}
		for _, endpoint := range viper.GetStringSlice("udp-endpoints") {
			if err := utils.Ping(endpoint); err == nil {
				m.endpoints = append(m.endpoints, endpoint)
			}
		}
		return clientScanMsg(t)
	})
}

type clientJoinMsg struct{}

func (m *LobbyModel) clientJoin() tea.Cmd {
	return func() tea.Msg {
		dataCh := make(chan string)
		m.local, m.shutdown, m.err = utils.ListenUDPData(0, dataCh)
		if m.err != nil {
			m.err = fmt.Errorf("failed to listen udp data, %w", m.err)
			return clientJoinMsg{}
		}

		m.err = utils.SendPayload(m.local, m.remote, types.JoinRequest{
			Endpoint: m.local,
			Username: m.client.UserInfo.Username,
		})
		if m.err != nil {
			m.err = fmt.Errorf("failed to send invitation, %w", m.err)
			return clientJoinMsg{}
		}

		data := <-dataCh
		m.shutdown()
		m.shutdown = func() {}

		var joinResponse types.JoinResponse
		if err := msgpack.Unmarshal([]byte(data), &joinResponse); err != nil {
			m.err = fmt.Errorf("failed to unmarshal msgpack, %w", err)
			return clientJoinMsg{}
		}

		m.remote = joinResponse.Endpoint
		return clientJoinMsg{}
	}
}

func (m *LobbyModel) Init() tea.Cmd {
	return tea.ClearScreen
}

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

	switch msg := msg.(type) {
	case tea.KeyMsg:
		switch msg.String() {
		case "ctrl+c":
			return m, tea.Quit
		}
	}

	switch m.op {
	case lobbyOperationChoose:
		switch msg := msg.(type) {
		case tea.KeyMsg:
			switch msg.String() {
			case "S", "s":
				m.op = lobbyOperationServerWaiting
				cmds = append(cmds, m.serverListen())
			case "C", "c":
				m.op = lobbyOperationClientScannning
				cmds = append(cmds, m.clientScan())
			}
		}
	case lobbyOperationServerWaiting:
		switch msg.(type) {
		case serverListenMsg:
			m.op = lobbyOperationServerChoose
		}
	case lobbyOperationServerChoose:
		switch msg := msg.(type) {
		case tea.KeyMsg:
			switch msg.String() {
			case "Y", "y", "enter":
				cmds = append(cmds, m.serverSendReply(true))
			case "N", "n":
				m.op = lobbyOperationServerWaiting
				cmds = append(cmds, m.serverSendReply(false))
			}
		case serverSendReplyMsg:
			if m.err == nil {
				m.shutdown()
				m.client.SetBaseURL("http://" + m.remote)
				shutdown := workflows.WordleServer(m.listener)
				*m.queue = append(*m.queue,
					tea.NewProgram(NewWordleClientModel(m.BaseModel, shutdown)))
				return m, tea.Quit
			} else {
				m.op = lobbyOperationServerWaiting
			}
		}
	case lobbyOperationClientScannning:
		switch msg := msg.(type) {
		case tea.KeyMsg:
			switch msg.String() {
			case "shift+tab", "up":
				if n := len(m.endpoints); n > 0 {
					m.cursor = (m.cursor - 1 + n) % n
				}
			case "tab", "down":
				if n := len(m.endpoints); n > 0 {
					m.cursor = (m.cursor + 1) % n
				}
			case "enter":
				m.remote = m.endpoints[m.cursor]
				cmds = append(cmds, m.clientJoin())
			}
		case clientScanMsg:
			cmds = append(cmds, m.clientScan())
		case clientJoinMsg:
			if m.err == nil && m.remote != "" {
				m.shutdown()
				m.client.SetBaseURL("http://" + m.remote)
				*m.BaseModel.queue = append(*m.BaseModel.queue,
					tea.NewProgram(NewWordleClientModel(m.BaseModel, func() {})))
				return m, tea.Quit
			}
		}
	}

	return m, tea.Batch(cmds...)
}

func (m *LobbyModel) View() string {
	var b strings.Builder

	switch m.op {
	case lobbyOperationChoose:
		b.WriteString("Choose Role\n(S)erver / (C)lient\n")
	case lobbyOperationServerWaiting:
		b.WriteString("Wait for user to join...\n")
	case lobbyOperationServerChoose:
		fmt.Fprintf(&b, "Receive Join Request by '%s'\nAccept (Y)/N\n",
			m.remoteUser)
	case lobbyOperationClientScannning:
		b.WriteString("Scanning server...\nChoose one to join\n")
		for i, endpoint := range m.endpoints {
			if i == m.cursor {
				fmt.Fprintf(&b, "(•) %s\n", endpoint)
			} else {
				fmt.Fprintf(&b, "( ) %s\n", endpoint)
			}
		}
	}

	if m.info != "" {
		b.WriteString(m.info + "\n")
	}

	if m.err != nil {
		b.WriteString("----------\n")
		b.WriteString(m.err.Error() + "\n")
	}

	return b.String()
}