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