grokkit/cmd/chat.go

// Owned by gmgauthier.com
// Current time: 2023-10-05 14:30:00 UTC

package cmd

import (
	"fmt"
	"os"
	"strings"

	"github.com/charmbracelet/bubbles/textarea"
	"github.com/charmbracelet/bubbles/viewport"
	tea "github.com/charmbracelet/bubbletea"
	"github.com/charmbracelet/lipgloss"
	"github.com/fatih/color"
	"github.com/spf13/cobra"
	"gmgauthier.com/grokkit/config"
	"gmgauthier.com/grokkit/internal/grok"
)

var chatCmd = &cobra.Command{
	Use:   "chat",
	Short: "Beautiful interactive TUI chat with Grok (full history + streaming)",
	Run: func(cmd *cobra.Command, args []string) {
		modelFlag, _ := cmd.Flags().GetString("model")
		model := config.GetModel(modelFlag)

		p := tea.NewProgram(initialModel(model), tea.WithAltScreen(), tea.WithMouseCellMotion())
		if _, err := p.Run(); err != nil {
			color.Red("TUI error: %v", err)
			os.Exit(1)
		}
	},
}

type replyMsg string

type model struct {
	viewport viewport.Model
	textarea textarea.Model
	history  []string
	client   *grok.Client
	model    string
	width    int
	height   int
	quitting bool
}

func initialModel(modelName string) model {
	ta := textarea.New()
	ta.Placeholder = "Type your message... (Enter to send, /quit to exit)"
	ta.Focus()
	ta.SetHeight(3)
	ta.ShowLineNumbers = false

	vp := viewport.New(0, 0)

	return model{
		textarea: ta,
		viewport: vp,
		history:  []string{},
		client:   grok.NewClient(),
		model:    modelName,
	}
}

func (m model) Init() tea.Cmd {
	return textarea.Blink
}

func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
	var cmd tea.Cmd
	var cmds []tea.Cmd

	switch msg := msg.(type) {
	case tea.WindowSizeMsg:
		m.width = msg.Width
		m.height = msg.Height
		m.viewport.Width = msg.Width
		m.viewport.Height = msg.Height - m.textarea.Height() - 5
		m.textarea.SetWidth(msg.Width - 4)
		return m, nil

	case tea.KeyMsg:
		switch msg.String() {
		case "ctrl+c", "esc":
			m.quitting = true
			return m, tea.Quit
		case "enter":
			input := strings.TrimSpace(m.textarea.Value())
			if input == "" {
				return m, nil
			}
			if input == "/quit" || input == "/q" {
				m.quitting = true
				return m, tea.Quit
			}

			m.history = append(m.history, "You > "+input)
			m.viewport.SetContent(buildHistoryView(m.history))
			m.viewport.GotoBottom()

			m.textarea.SetValue("")

			historyForAPI := buildHistoryForAPI(m.history)
			cmds = append(cmds, streamReplyCmd(m.client, historyForAPI, m.model))

			return m, tea.Batch(cmds...)
		}

	case replyMsg:
		m.history = append(m.history, "Grok > "+string(msg))
		m.viewport.SetContent(buildHistoryView(m.history))
		m.viewport.GotoBottom()
		return m, nil
	}

	m.textarea, cmd = m.textarea.Update(msg)
	cmds = append(cmds, cmd)
	m.viewport, cmd = m.viewport.Update(msg)
	cmds = append(cmds, cmd)

	return m, tea.Batch(cmds...)
}

// buildHistoryView applies colors and ensures wrapping
func buildHistoryView(lines []string) string {
	userStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("33")) // blue
	grokStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("10")) // green

	var b strings.Builder
	for _, line := range lines {
		if strings.HasPrefix(line, "You > ") {
			b.WriteString(userStyle.Render(line) + "\n")
		} else if strings.HasPrefix(line, "Grok > ") {
			b.WriteString(grokStyle.Render(line) + "\n")
		} else {
			b.WriteString(line + "\n")
		}
	}
	return b.String()
}

func buildHistoryForAPI(lines []string) []map[string]string {
	h := []map[string]string{
		{"role": "system", "content": fmt.Sprintf("You are Grok 4, the latest model from xAI (2026). You are currently running as `%s`.", "grok-4")},
	}
	for _, line := range lines {
		if strings.HasPrefix(line, "You > ") {
			h = append(h, map[string]string{"role": "user", "content": strings.TrimPrefix(line, "You > ")})
		} else if strings.HasPrefix(line, "Grok > ") {
			h = append(h, map[string]string{"role": "assistant", "content": strings.TrimPrefix(line, "Grok > ")})
		}
	}
	return h
}

func streamReplyCmd(client *grok.Client, history []map[string]string, model string) tea.Cmd {
	return func() tea.Msg {
		reply := client.StreamSilent(history, model)
		return replyMsg(reply)
	}
}

func (m model) View() string {
	if m.quitting {
		return "Goodbye 👋\n"
	}

	header := lipgloss.NewStyle().
		Bold(true).
		Foreground(lipgloss.Color("12")).
		Width(m.width).
		Align(lipgloss.Center).
		Render(fmt.Sprintf(" Grokkit Chat  —  Model: %s ", m.model))

	return lipgloss.JoinVertical(lipgloss.Left,
		header,
		m.viewport.View(),
		m.textarea.View(),
	)
}