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