1package compact
  2
  3import (
  4	"context"
  5
  6	"github.com/charmbracelet/bubbles/v2/key"
  7	tea "github.com/charmbracelet/bubbletea/v2"
  8	"github.com/charmbracelet/lipgloss/v2"
  9
 10	"github.com/charmbracelet/crush/internal/client"
 11	"github.com/charmbracelet/crush/internal/llm/agent"
 12	"github.com/charmbracelet/crush/internal/proto"
 13	"github.com/charmbracelet/crush/internal/tui/components/core"
 14	"github.com/charmbracelet/crush/internal/tui/components/dialogs"
 15	"github.com/charmbracelet/crush/internal/tui/styles"
 16	"github.com/charmbracelet/crush/internal/tui/util"
 17)
 18
 19const CompactDialogID dialogs.DialogID = "compact"
 20
 21// CompactDialog interface for the session compact dialog
 22type CompactDialog interface {
 23	dialogs.DialogModel
 24}
 25
 26type compactDialogCmp struct {
 27	wWidth, wHeight int
 28	width, height   int
 29	selected        int
 30	keyMap          KeyMap
 31	sessionID       string
 32	state           compactState
 33	progress        string
 34	c               *client.Client
 35	ins             *proto.Instance
 36	noAsk           bool // If true, skip confirmation dialog
 37}
 38
 39type compactState int
 40
 41const (
 42	stateConfirm compactState = iota
 43	stateCompacting
 44	stateError
 45)
 46
 47// NewCompactDialogCmp creates a new session compact dialog
 48func NewCompactDialogCmp(c *client.Client, ins *proto.Instance, sessionID string, noAsk bool) CompactDialog {
 49	return &compactDialogCmp{
 50		sessionID: sessionID,
 51		keyMap:    DefaultKeyMap(),
 52		state:     stateConfirm,
 53		selected:  0,
 54		c:         c,
 55		ins:       ins,
 56		noAsk:     noAsk,
 57	}
 58}
 59
 60func (c *compactDialogCmp) Init() tea.Cmd {
 61	if c.noAsk {
 62		// If noAsk is true, skip confirmation and start compaction immediately
 63		return c.startCompaction()
 64	}
 65	return nil
 66}
 67
 68func (c *compactDialogCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 69	switch msg := msg.(type) {
 70	case tea.WindowSizeMsg:
 71		c.wWidth = msg.Width
 72		c.wHeight = msg.Height
 73		cmd := c.SetSize()
 74		return c, cmd
 75
 76	case tea.KeyPressMsg:
 77		switch c.state {
 78		case stateConfirm:
 79			switch {
 80			case key.Matches(msg, c.keyMap.ChangeSelection):
 81				c.selected = (c.selected + 1) % 2
 82				return c, nil
 83			case key.Matches(msg, c.keyMap.Select):
 84				if c.selected == 0 {
 85					return c, c.startCompaction()
 86				} else {
 87					return c, util.CmdHandler(dialogs.CloseDialogMsg{})
 88				}
 89			case key.Matches(msg, c.keyMap.Y):
 90				return c, c.startCompaction()
 91			case key.Matches(msg, c.keyMap.N):
 92				return c, util.CmdHandler(dialogs.CloseDialogMsg{})
 93			case key.Matches(msg, c.keyMap.Close):
 94				return c, util.CmdHandler(dialogs.CloseDialogMsg{})
 95			}
 96		case stateCompacting:
 97			switch {
 98			case key.Matches(msg, c.keyMap.Close):
 99				return c, util.CmdHandler(dialogs.CloseDialogMsg{})
100			}
101		case stateError:
102			switch {
103			case key.Matches(msg, c.keyMap.Select):
104				return c, util.CmdHandler(dialogs.CloseDialogMsg{})
105			case key.Matches(msg, c.keyMap.Close):
106				return c, util.CmdHandler(dialogs.CloseDialogMsg{})
107			}
108		}
109
110	case agent.AgentEvent:
111		switch msg.Type {
112		case agent.AgentEventTypeSummarize:
113			if msg.Error != nil {
114				c.state = stateError
115				c.progress = "Error: " + msg.Error.Error()
116			} else if msg.Done {
117				return c, util.CmdHandler(dialogs.CloseDialogMsg{})
118			} else {
119				c.progress = msg.Progress
120			}
121		case agent.AgentEventTypeError:
122			// Handle errors that occur during summarization but are sent as separate error events.
123			c.state = stateError
124			if msg.Error != nil {
125				c.progress = "Error: " + msg.Error.Error()
126			} else {
127				c.progress = "An unknown error occurred"
128			}
129		}
130		return c, nil
131	}
132
133	return c, nil
134}
135
136func (c *compactDialogCmp) startCompaction() tea.Cmd {
137	c.state = stateCompacting
138	c.progress = "Starting summarization..."
139	return func() tea.Msg {
140		err := c.c.AgentSummarizeSession(context.Background(), c.ins.ID, c.sessionID)
141		if err != nil {
142			c.state = stateError
143			c.progress = "Error: " + err.Error()
144		}
145		return nil
146	}
147}
148
149func (c *compactDialogCmp) renderButtons() string {
150	t := styles.CurrentTheme()
151	baseStyle := t.S().Base
152
153	buttons := []core.ButtonOpts{
154		{
155			Text:           "Yes",
156			UnderlineIndex: 0, // "Y"
157			Selected:       c.selected == 0,
158		},
159		{
160			Text:           "No",
161			UnderlineIndex: 0, // "N"
162			Selected:       c.selected == 1,
163		},
164	}
165
166	content := core.SelectableButtons(buttons, "  ")
167
168	return baseStyle.AlignHorizontal(lipgloss.Right).Width(c.width - 4).Render(content)
169}
170
171func (c *compactDialogCmp) renderContent() string {
172	t := styles.CurrentTheme()
173	baseStyle := t.S().Base
174
175	switch c.state {
176	case stateConfirm:
177		explanation := t.S().Text.
178			Width(c.width - 4).
179			Render("This will summarize the current session and reset the context. The conversation history will be condensed into a summary to free up context space while preserving important information.")
180
181		question := t.S().Text.
182			Width(c.width - 4).
183			Render("Do you want to continue?")
184
185		return baseStyle.Render(lipgloss.JoinVertical(
186			lipgloss.Left,
187			explanation,
188			"",
189			question,
190		))
191	case stateCompacting:
192		return baseStyle.Render(lipgloss.JoinVertical(
193			lipgloss.Left,
194			c.progress,
195			"",
196			"Please wait...",
197		))
198	case stateError:
199		return baseStyle.Render(lipgloss.JoinVertical(
200			lipgloss.Left,
201			c.progress,
202			"",
203			"Press Enter to close",
204		))
205	}
206	return ""
207}
208
209func (c *compactDialogCmp) render() string {
210	t := styles.CurrentTheme()
211	baseStyle := t.S().Base
212
213	var title string
214	switch c.state {
215	case stateConfirm:
216		title = "Compact Session"
217	case stateCompacting:
218		title = "Compacting Session"
219	case stateError:
220		title = "Compact Failed"
221	}
222
223	titleView := core.Title(title, c.width-4)
224	content := c.renderContent()
225
226	var dialogContent string
227	if c.state == stateConfirm {
228		buttons := c.renderButtons()
229		dialogContent = lipgloss.JoinVertical(
230			lipgloss.Top,
231			titleView,
232			"",
233			content,
234			"",
235			buttons,
236			"",
237		)
238	} else {
239		dialogContent = lipgloss.JoinVertical(
240			lipgloss.Top,
241			titleView,
242			"",
243			content,
244			"",
245		)
246	}
247
248	return baseStyle.
249		Padding(0, 1).
250		Border(lipgloss.RoundedBorder()).
251		BorderForeground(t.BorderFocus).
252		Width(c.width).
253		Render(dialogContent)
254}
255
256func (c *compactDialogCmp) View() string {
257	return c.render()
258}
259
260// SetSize sets the size of the component.
261func (c *compactDialogCmp) SetSize() tea.Cmd {
262	c.width = min(90, c.wWidth)
263	c.height = min(15, c.wHeight)
264	return nil
265}
266
267func (c *compactDialogCmp) Position() (int, int) {
268	row := (c.wHeight / 2) - (c.height / 2)
269	col := (c.wWidth / 2) - (c.width / 2)
270	return row, col
271}
272
273// ID implements CompactDialog.
274func (c *compactDialogCmp) ID() dialogs.DialogID {
275	return CompactDialogID
276}