Implemented proper connection and retry mechanisms

2 files changed, 77 insertions, 6 deletions

diff --git a/internal/client/ui/core/core.go b/internal/client/ui/core/core.go
index 12b8b42..790c52e 100644
--- a/internal/client/ui/core/core.go
+++ b/internal/client/ui/core/core.go
@@ -2,37 +2,51 @@ package core
 
 import (
 	"crypto/ed25519"
+	"fmt"
 	"time"
 
+	"github.com/charmbracelet/bubbles/spinner"
+	"github.com/charmbracelet/bubbles/timer"
 	tea "github.com/charmbracelet/bubbletea"
 	"github.com/kyren223/eko/internal/client/gateway"
+	"github.com/kyren223/eko/internal/client/ui"
 	"github.com/kyren223/eko/internal/client/ui/loadscreen"
+	"github.com/kyren223/eko/pkg/snowflake"
 )
 
 var (
 	connectingToServer = "Connecting to server.."
 	connectionFailed   = "Connection failed - retrying in %d sec..."
+	connectionTimeout  = 5 * time.Second
+	initialTimeout     = 3750 * time.Millisecond
+	timerInterval      = 50 * time.Millisecond
 )
 
 type Model struct {
-	privKey ed25519.PrivateKey
 	name    string
+	privKey ed25519.PrivateKey
 
 	loading   loadscreen.Model
+	timeout   time.Duration
+	timer     timer.Model
 	connected bool
+
+	id snowflake.ID
 }
 
 func New(privKey ed25519.PrivateKey, name string) Model {
 	return Model{
-		privKey:   privKey,
 		name:      name,
+		privKey:   privKey,
 		loading:   loadscreen.New(connectingToServer),
+		timeout:   initialTimeout,
+		timer:     newTimer(initialTimeout),
 		connected: false,
 	}
 }
 
 func (m Model) Init() tea.Cmd {
-	return tea.Batch(gateway.Connect(m.privKey, 5*time.Second), m.loading.Init())
+	return tea.Batch(gateway.Connect(m.privKey, connectionTimeout), m.loading.Init())
 }
 
 func (m Model) View() string {
@@ -45,9 +59,62 @@ func (m Model) View() string {
 
 func (m Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 	if !m.connected {
-		var loadscreenCmd tea.Cmd
-		m.loading, loadscreenCmd = m.loading.Update(msg)
-		return m, loadscreenCmd
+		switch msg := msg.(type) {
+		case gateway.ConnectionEstablished:
+			m.id = snowflake.ID(msg)
+			m.connected = true
+			m.timeout = initialTimeout
+			return m, m.timer.Stop()
+
+		case gateway.ConnectionFailed:
+			m.timer = newTimer(m.timeout)
+			m.updateLoadScreenContent()
+			return m, m.timer.Start()
+
+		case timer.TimeoutMsg:
+			m.timeout = min(m.timeout*2, time.Minute)
+			m.loading.SetContent(connectingToServer)
+			return m, gateway.Connect(m.privKey, connectionTimeout)
+
+		case timer.StartStopMsg:
+			var cmd tea.Cmd
+			m.timer, cmd = m.timer.Update(msg)
+			return m, cmd
+
+		case timer.TickMsg:
+			m.updateLoadScreenContent()
+			var cmd tea.Cmd
+			m.timer, cmd = m.timer.Update(msg)
+			return m, cmd
+
+		case spinner.TickMsg:
+			var loadscreenCmd tea.Cmd
+			m.loading, loadscreenCmd = m.loading.Update(msg)
+			return m, loadscreenCmd
+
+		default:
+			return m, nil
+		}
 	}
+
+	switch msg.(type) {
+	case gateway.ConnectionLost:
+		m.connected = false
+		m.timeout = initialTimeout
+		return m, tea.Batch(gateway.Connect(m.privKey, connectionTimeout), m.loading.Init())
+
+	case ui.QuitMsg:
+		gateway.Disconnect()
+	}
+
 	return m, nil
 }
+
+func (m *Model) updateLoadScreenContent() {
+	seconds := m.timer.Timeout.Round(time.Second) / time.Second
+	m.loading.SetContent(fmt.Sprintf(connectionFailed, seconds))
+}
+
+func newTimer(timeout time.Duration) timer.Model {
+	return timer.NewWithInterval(timeout.Truncate(time.Second)+(time.Second/2), timerInterval)
+}
diff --git a/internal/client/ui/ui.go b/internal/client/ui/ui.go
index 8cda302..762d0b6 100644
--- a/internal/client/ui/ui.go
+++ b/internal/client/ui/ui.go
@@ -15,6 +15,8 @@ import (
 	"github.com/kyren223/eko/pkg/assert"
 )
 
+const DEBUG = true
+
 var Width int
 var Height int
 var BackgroundColor = "#1E1E2E"
@@ -32,6 +34,8 @@ func Transition(model tea.Model) tea.Cmd {
 	}
 }
 
+type QuitMsg struct{}
+
 func AddBorderHeader(header string, headerOffset int, style lipgloss.Style, render string) string {
 	b := style.GetBorderStyle()
 	body := style.UnsetBorderTop().Render(render)