1package page
2
3import (
4 "fmt"
5 "os"
6 "path/filepath"
7 "strconv"
8
9 "github.com/charmbracelet/bubbles/key"
10 tea "github.com/charmbracelet/bubbletea"
11 "github.com/charmbracelet/huh"
12 "github.com/charmbracelet/lipgloss"
13 "github.com/kujtimiihoxha/termai/internal/llm/models"
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 "github.com/spf13/viper"
18)
19
20var InitPage PageID = "init"
21
22type configSaved struct{}
23
24type initPage struct {
25 form *huh.Form
26 width int
27 height int
28 saved bool
29 errorMsg string
30 statusMsg string
31 modelOpts []huh.Option[string]
32 bigModel string
33 smallModel string
34 openAIKey string
35 anthropicKey string
36 groqKey string
37 maxTokens string
38 dataDir string
39 agent string
40}
41
42func (i *initPage) Init() tea.Cmd {
43 return i.form.Init()
44}
45
46func (i *initPage) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
47 var cmds []tea.Cmd
48
49 switch msg := msg.(type) {
50 case tea.WindowSizeMsg:
51 i.width = msg.Width - 4 // Account for border
52 i.height = msg.Height - 4
53 i.form = i.form.WithWidth(i.width).WithHeight(i.height)
54 return i, nil
55
56 case configSaved:
57 i.saved = true
58 i.statusMsg = "Configuration saved successfully. Press any key to continue."
59 return i, nil
60 }
61
62 if i.saved {
63 switch msg.(type) {
64 case tea.KeyMsg:
65 return i, util.CmdHandler(PageChangeMsg{ID: ReplPage})
66 }
67 return i, nil
68 }
69
70 // Process the form
71 form, cmd := i.form.Update(msg)
72 if f, ok := form.(*huh.Form); ok {
73 i.form = f
74 cmds = append(cmds, cmd)
75 }
76
77 if i.form.State == huh.StateCompleted {
78 // Save configuration to file
79 configPath := filepath.Join(os.Getenv("HOME"), ".termai.yaml")
80 maxTokens, _ := strconv.Atoi(i.maxTokens)
81 config := map[string]any{
82 "models": map[string]string{
83 "big": i.bigModel,
84 "small": i.smallModel,
85 },
86 "providers": map[string]any{
87 "openai": map[string]string{
88 "key": i.openAIKey,
89 },
90 "anthropic": map[string]string{
91 "key": i.anthropicKey,
92 },
93 "groq": map[string]string{
94 "key": i.groqKey,
95 },
96 "common": map[string]int{
97 "max_tokens": maxTokens,
98 },
99 },
100 "data": map[string]string{
101 "dir": i.dataDir,
102 },
103 "agents": map[string]string{
104 "default": i.agent,
105 },
106 "log": map[string]string{
107 "level": "info",
108 },
109 }
110
111 // Write config to viper
112 for k, v := range config {
113 viper.Set(k, v)
114 }
115
116 // Save configuration
117 err := viper.WriteConfigAs(configPath)
118 if err != nil {
119 i.errorMsg = fmt.Sprintf("Failed to save configuration: %s", err)
120 return i, nil
121 }
122
123 // Return to main page
124 return i, util.CmdHandler(configSaved{})
125 }
126
127 return i, tea.Batch(cmds...)
128}
129
130func (i *initPage) View() string {
131 if i.saved {
132 return lipgloss.NewStyle().
133 Width(i.width).
134 Height(i.height).
135 Align(lipgloss.Center, lipgloss.Center).
136 Render(lipgloss.JoinVertical(
137 lipgloss.Center,
138 lipgloss.NewStyle().Foreground(styles.Green).Render("✓ Configuration Saved"),
139 "",
140 lipgloss.NewStyle().Foreground(styles.Blue).Render(i.statusMsg),
141 ))
142 }
143
144 view := i.form.View()
145 if i.errorMsg != "" {
146 errorBox := lipgloss.NewStyle().
147 Padding(1).
148 Border(lipgloss.RoundedBorder()).
149 BorderForeground(styles.Red).
150 Width(i.width - 4).
151 Render(i.errorMsg)
152 view = lipgloss.JoinVertical(lipgloss.Left, errorBox, view)
153 }
154 return view
155}
156
157func (i *initPage) GetSize() (int, int) {
158 return i.width, i.height
159}
160
161func (i *initPage) SetSize(width int, height int) {
162 i.width = width
163 i.height = height
164 i.form = i.form.WithWidth(width).WithHeight(height)
165}
166
167func (i *initPage) BindingKeys() []key.Binding {
168 if i.saved {
169 return []key.Binding{
170 key.NewBinding(
171 key.WithKeys("enter", "space", "esc"),
172 key.WithHelp("any key", "continue"),
173 ),
174 }
175 }
176 return i.form.KeyBinds()
177}
178
179func NewInitPage() tea.Model {
180 // Create model options
181 var modelOpts []huh.Option[string]
182 for id, model := range models.SupportedModels {
183 modelOpts = append(modelOpts, huh.NewOption(model.Name, string(id)))
184 }
185
186 // Create agent options
187 agentOpts := []huh.Option[string]{
188 huh.NewOption("Coder", "coder"),
189 huh.NewOption("Assistant", "assistant"),
190 }
191
192 // Init page with form
193 initModel := &initPage{
194 modelOpts: modelOpts,
195 bigModel: string(models.Claude37Sonnet),
196 smallModel: string(models.Claude37Sonnet),
197 maxTokens: "4000",
198 dataDir: ".termai",
199 agent: "coder",
200 }
201
202 // API Keys group
203 apiKeysGroup := huh.NewGroup(
204 huh.NewNote().
205 Title("API Keys").
206 Description("You need to provide at least one API key to use termai"),
207
208 huh.NewInput().
209 Title("OpenAI API Key").
210 Placeholder("sk-...").
211 Key("openai_key").
212 Value(&initModel.openAIKey),
213
214 huh.NewInput().
215 Title("Anthropic API Key").
216 Placeholder("sk-ant-...").
217 Key("anthropic_key").
218 Value(&initModel.anthropicKey),
219
220 huh.NewInput().
221 Title("Groq API Key").
222 Placeholder("gsk_...").
223 Key("groq_key").
224 Value(&initModel.groqKey),
225 )
226
227 // Model configuration group
228 modelsGroup := huh.NewGroup(
229 huh.NewNote().
230 Title("Model Configuration").
231 Description("Select which models to use"),
232
233 huh.NewSelect[string]().
234 Title("Big Model").
235 Options(modelOpts...).
236 Key("big_model").
237 Value(&initModel.bigModel),
238
239 huh.NewSelect[string]().
240 Title("Small Model").
241 Options(modelOpts...).
242 Key("small_model").
243 Value(&initModel.smallModel),
244
245 huh.NewInput().
246 Title("Max Tokens").
247 Placeholder("4000").
248 Key("max_tokens").
249 CharLimit(5).
250 Validate(func(s string) error {
251 var n int
252 _, err := fmt.Sscanf(s, "%d", &n)
253 if err != nil || n <= 0 {
254 return fmt.Errorf("must be a positive number")
255 }
256 initModel.maxTokens = s
257 return nil
258 }).
259 Value(&initModel.maxTokens),
260 )
261
262 // General settings group
263 generalGroup := huh.NewGroup(
264 huh.NewNote().
265 Title("General Settings").
266 Description("Configure general termai settings"),
267
268 huh.NewInput().
269 Title("Data Directory").
270 Placeholder(".termai").
271 Key("data_dir").
272 Value(&initModel.dataDir),
273
274 huh.NewSelect[string]().
275 Title("Default Agent").
276 Options(agentOpts...).
277 Key("agent").
278 Value(&initModel.agent),
279
280 huh.NewConfirm().
281 Title("Save Configuration").
282 Affirmative("Save").
283 Negative("Cancel"),
284 )
285
286 // Create form with theme
287 form := huh.NewForm(
288 apiKeysGroup,
289 modelsGroup,
290 generalGroup,
291 ).WithTheme(styles.HuhTheme()).
292 WithShowHelp(true).
293 WithShowErrors(true)
294
295 // Set the form in the model
296 initModel.form = form
297
298 return layout.NewSinglePane(
299 initModel,
300 layout.WithSinglePaneFocusable(true),
301 layout.WithSinglePaneBordered(true),
302 layout.WithSinglePaneBorderText(
303 map[layout.BorderPosition]string{
304 layout.TopMiddleBorder: "Welcome to termai - Initial Setup",
305 },
306 ),
307 )
308}