From 539835627b74eee1e19347031247e6199dd32021 Mon Sep 17 00:00:00 2001
From: juancwu <git@juancwu.dev>
Date: Mon, 26 Jan 2026 18:40:08 +0000
Subject: [PATCH] save all pages in map

---
 internal/ui/model.go | 55 ++++++++++++++++++--------------------------
 1 file changed, 23 insertions(+), 32 deletions(-)

diff --git a/internal/ui/model.go b/internal/ui/model.go
index cd225de..8032aa8 100644
--- a/internal/ui/model.go
+++ b/internal/ui/model.go
@@ -15,10 +15,7 @@ import (
 
 type MainModel struct {
 	currentPage messages.Page
-	login       tea.Model
-	menu        tea.Model
-	listDomains tea.Model
-	dnsRetrieve tea.Model
+	pages       map[messages.Page]tea.Model
 	isMenuInit  bool
 	width       int
 	height      int
@@ -31,14 +28,17 @@ func New() MainModel {
 		os.Exit(1)
 	}
 
+	pages := make(map[messages.Page]tea.Model)
+	pages[messages.PageLogin] = login.New(cfg)
+
 	return MainModel{
 		currentPage: messages.PageLogin,
-		login:       login.New(cfg),
+		pages:       pages,
 	}
 }
 
 func (m MainModel) Init() tea.Cmd {
-	return m.login.Init()
+	return m.pages[messages.PageLogin].Init()
 }
 
 func (m MainModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
@@ -57,47 +57,38 @@ func (m MainModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 		m.currentPage = msg.Page
 		return m, nil
 	case messages.SessionReadyMsg:
-		m.menu = menu.New(msg.Client)
-		m.listDomains = listdomains.New(msg.Client)
-		m.dnsRetrieve = dns.NewRetrieveModel(msg.Client)
+		m.pages[messages.PageMenu] = menu.New(msg.Client)
+		m.pages[messages.PageListDomains] = listdomains.New(msg.Client)
+		m.pages[messages.PageDNSRetrieve] = dns.NewRetrieveModel(msg.Client)
 		if m.width > 0 && m.height > 0 {
-			m.menu, _ = m.menu.Update(tea.WindowSizeMsg{Width: m.width, Height: m.height})
+			m.pages[messages.PageMenu], _ = m.pages[messages.PageMenu].Update(tea.WindowSizeMsg{Width: m.width, Height: m.height})
 		}
 		m.currentPage = messages.PageMenu
 		m.isMenuInit = true
-		return m, tea.Batch(m.menu.Init(), m.listDomains.Init(), m.dnsRetrieve.Init())
+		var cmds []tea.Cmd
+		for _, page := range m.pages {
+			cmds = append(cmds, page.Init())
+		}
+		return m, tea.Batch(cmds...)
 	case messages.ListDomainsMsg:
 		m.currentPage = messages.PageListDomains
 	case messages.DNSRetrieveMsg:
 		m.currentPage = messages.PageDNSRetrieve
 	}
 
-	switch m.currentPage {
-	case messages.PageLogin:
-		m.login, cmd = m.login.Update(msg)
-	case messages.PageMenu:
-		m.menu, cmd = m.menu.Update(msg)
-	case messages.PageListDomains:
-		m.listDomains, cmd = m.listDomains.Update(msg)
-	case messages.PageDNSRetrieve:
-		m.dnsRetrieve, cmd = m.dnsRetrieve.Update(msg)
+	page, ok := m.pages[m.currentPage]
+	if ok {
+		m.pages[m.currentPage], cmd = page.Update(msg)
+		cmds = append(cmds, cmd)
 	}
-	cmds = append(cmds, cmd)
 
 	return m, tea.Batch(cmds...)
 }
 
 func (m MainModel) View() string {
-	switch m.currentPage {
-	case messages.PageLogin:
-		return m.login.View()
-	case messages.PageMenu:
-		return m.menu.View()
-	case messages.PageListDomains:
-		return m.listDomains.View()
-	case messages.PageDNSRetrieve:
-		return m.dnsRetrieve.View()
-	default:
-		return "Unknown Page"
+	page, ok := m.pages[m.currentPage]
+	if ok {
+		return page.View()
 	}
+	return "Unknown Page"
 }