1package repl
2
3import (
4 "fmt"
5 "strings"
6
7 "github.com/charmbracelet/bubbles/key"
8 "github.com/charmbracelet/bubbles/list"
9 tea "github.com/charmbracelet/bubbletea"
10 "github.com/charmbracelet/lipgloss"
11 "github.com/kujtimiihoxha/termai/internal/app"
12 "github.com/kujtimiihoxha/termai/internal/pubsub"
13 "github.com/kujtimiihoxha/termai/internal/session"
14 "github.com/kujtimiihoxha/termai/internal/tui/layout"
15 "github.com/kujtimiihoxha/termai/internal/tui/styles"
16 "github.com/kujtimiihoxha/termai/internal/tui/util"
17)
18
19type SessionsCmp interface {
20 tea.Model
21 layout.Sizeable
22 layout.Focusable
23 layout.Bordered
24 layout.Bindings
25}
26type sessionsCmp struct {
27 app *app.App
28 list list.Model
29 focused bool
30}
31
32type listItem struct {
33 id, title, desc string
34}
35
36func (i listItem) Title() string { return i.title }
37func (i listItem) Description() string { return i.desc }
38func (i listItem) FilterValue() string { return i.title }
39
40type InsertSessionsMsg struct {
41 sessions []session.Session
42}
43
44type SelectedSessionMsg struct {
45 SessionID string
46}
47
48type sessionsKeyMap struct {
49 Select key.Binding
50}
51
52var sessionKeyMapValue = sessionsKeyMap{
53 Select: key.NewBinding(
54 key.WithKeys("enter", " "),
55 key.WithHelp("enter/space", "select session"),
56 ),
57}
58
59func (i *sessionsCmp) Init() tea.Cmd {
60 existing, err := i.app.Sessions.List()
61 if err != nil {
62 return util.ReportError(err)
63 }
64 if len(existing) == 0 || existing[0].MessageCount > 0 {
65 newSession, err := i.app.Sessions.Create(
66 "New Session",
67 )
68 if err != nil {
69 return util.ReportError(err)
70 }
71 existing = append([]session.Session{newSession}, existing...)
72 }
73 return tea.Batch(
74 util.CmdHandler(InsertSessionsMsg{existing}),
75 util.CmdHandler(SelectedSessionMsg{existing[0].ID}),
76 )
77}
78
79func (i *sessionsCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
80 switch msg := msg.(type) {
81 case InsertSessionsMsg:
82 items := make([]list.Item, len(msg.sessions))
83 for i, s := range msg.sessions {
84 items[i] = listItem{
85 id: s.ID,
86 title: s.Title,
87 desc: formatTokensAndCost(s.PromptTokens+s.CompletionTokens, s.Cost),
88 }
89 }
90 return i, i.list.SetItems(items)
91 case pubsub.Event[session.Session]:
92 if msg.Type == pubsub.CreatedEvent && msg.Payload.ParentSessionID == "" {
93 // Check if the session is already in the list
94 items := i.list.Items()
95 for _, item := range items {
96 s := item.(listItem)
97 if s.id == msg.Payload.ID {
98 return i, nil
99 }
100 }
101 // insert the new session at the top of the list
102 items = append([]list.Item{listItem{
103 id: msg.Payload.ID,
104 title: msg.Payload.Title,
105 desc: formatTokensAndCost(msg.Payload.PromptTokens+msg.Payload.CompletionTokens, msg.Payload.Cost),
106 }}, items...)
107 return i, i.list.SetItems(items)
108 } else if msg.Type == pubsub.UpdatedEvent {
109 // update the session in the list
110 items := i.list.Items()
111 for idx, item := range items {
112 s := item.(listItem)
113 if s.id == msg.Payload.ID {
114 s.title = msg.Payload.Title
115 s.desc = formatTokensAndCost(msg.Payload.PromptTokens+msg.Payload.CompletionTokens, msg.Payload.Cost)
116 items[idx] = s
117 break
118 }
119 }
120 return i, i.list.SetItems(items)
121 }
122
123 case tea.KeyMsg:
124 switch {
125 case key.Matches(msg, sessionKeyMapValue.Select):
126 i.app.Logger.PersistInfo("Session selected")
127 selected := i.list.SelectedItem()
128 if selected == nil {
129 return i, nil
130 }
131 return i, util.CmdHandler(SelectedSessionMsg{selected.(listItem).id})
132 }
133 }
134 if i.focused {
135 u, cmd := i.list.Update(msg)
136 i.list = u
137 return i, cmd
138 }
139 return i, nil
140}
141
142func (i *sessionsCmp) View() string {
143 return i.list.View()
144}
145
146func (i *sessionsCmp) Blur() tea.Cmd {
147 i.focused = false
148 return nil
149}
150
151func (i *sessionsCmp) Focus() tea.Cmd {
152 i.focused = true
153 return nil
154}
155
156func (i *sessionsCmp) GetSize() (int, int) {
157 return i.list.Width(), i.list.Height()
158}
159
160func (i *sessionsCmp) IsFocused() bool {
161 return i.focused
162}
163
164func (i *sessionsCmp) SetSize(width int, height int) {
165 i.list.SetSize(width, height)
166}
167
168func (i *sessionsCmp) BorderText() map[layout.BorderPosition]string {
169 totalCount := len(i.list.Items())
170 itemsPerPage := i.list.Paginator.PerPage
171 currentPage := i.list.Paginator.Page
172
173 current := min(currentPage*itemsPerPage+itemsPerPage, totalCount)
174
175 pageInfo := fmt.Sprintf(
176 "%d-%d of %d",
177 currentPage*itemsPerPage+1,
178 current,
179 totalCount,
180 )
181
182 title := "Sessions"
183 if i.focused {
184 title = lipgloss.NewStyle().Foreground(styles.Primary).Render(title)
185 }
186 return map[layout.BorderPosition]string{
187 layout.TopMiddleBorder: title,
188 layout.BottomMiddleBorder: pageInfo,
189 }
190}
191
192func (i *sessionsCmp) BindingKeys() []key.Binding {
193 return append(layout.KeyMapToSlice(i.list.KeyMap), sessionKeyMapValue.Select)
194}
195
196func formatTokensAndCost(tokens int64, cost float64) string {
197 // Format tokens in human-readable format (e.g., 110K, 1.2M)
198 var formattedTokens string
199 switch {
200 case tokens >= 1_000_000:
201 formattedTokens = fmt.Sprintf("%.1fM", float64(tokens)/1_000_000)
202 case tokens >= 1_000:
203 formattedTokens = fmt.Sprintf("%.1fK", float64(tokens)/1_000)
204 default:
205 formattedTokens = fmt.Sprintf("%d", tokens)
206 }
207
208 // Remove .0 suffix if present
209 if strings.HasSuffix(formattedTokens, ".0K") {
210 formattedTokens = strings.Replace(formattedTokens, ".0K", "K", 1)
211 }
212 if strings.HasSuffix(formattedTokens, ".0M") {
213 formattedTokens = strings.Replace(formattedTokens, ".0M", "M", 1)
214 }
215
216 // Format cost with $ symbol and 2 decimal places
217 formattedCost := fmt.Sprintf("$%.2f", cost)
218
219 return fmt.Sprintf("Tokens: %s, Cost: %s", formattedTokens, formattedCost)
220}
221
222func NewSessionsCmp(app *app.App) SessionsCmp {
223 listDelegate := list.NewDefaultDelegate()
224 defaultItemStyle := list.NewDefaultItemStyles()
225 defaultItemStyle.SelectedTitle = defaultItemStyle.SelectedTitle.BorderForeground(styles.Secondary).Foreground(styles.Primary)
226 defaultItemStyle.SelectedDesc = defaultItemStyle.SelectedDesc.BorderForeground(styles.Secondary).Foreground(styles.Primary)
227
228 defaultStyle := list.DefaultStyles()
229 defaultStyle.FilterPrompt = defaultStyle.FilterPrompt.Foreground(styles.Secondary)
230 defaultStyle.FilterCursor = defaultStyle.FilterCursor.Foreground(styles.Flamingo)
231
232 listDelegate.Styles = defaultItemStyle
233
234 listComponent := list.New([]list.Item{}, listDelegate, 0, 0)
235 listComponent.FilterInput.PromptStyle = defaultStyle.FilterPrompt
236 listComponent.FilterInput.Cursor.Style = defaultStyle.FilterCursor
237 listComponent.SetShowTitle(false)
238 listComponent.SetShowPagination(false)
239 listComponent.SetShowHelp(false)
240 listComponent.SetShowStatusBar(false)
241 listComponent.DisableQuitKeybindings()
242
243 return &sessionsCmp{
244 app: app,
245 list: listComponent,
246 focused: false,
247 }
248}