yoyo/internal/tui/model.go

package tui

import (
	"context"
	"fmt"
	"strconv"
	"strings"
	"time"

	"charm.land/bubbles/v2/help"
	"charm.land/bubbles/v2/key"
	"charm.land/bubbles/v2/spinner"
	"charm.land/bubbles/v2/textarea"
	"charm.land/bubbles/v2/viewport"
	"charm.land/bubbletea/v2"
	"charm.land/lipgloss/v2"
	"github.com/titor/fanyi/internal/config"
	"github.com/titor/fanyi/internal/logo"
	"github.com/titor/fanyi/internal/translator"
)

var supportedLangs = []string{"zh-CN", "en-US", "ja", "ko", "zh-TW", "es", "fr", "de"}

type translateMsg struct {
	result string
	tokens int
	err    error
}

type model struct {
	config     *config.Config
	translator *translator.Translator

	messages []ChatMessage
	input    textarea.Model
	viewport viewport.Model
	spinner  spinner.Model
	help     help.Model
	keys     KeyMap

	targetLang string
	langIndex  int
	loading    bool
	lastInput  string

	width       int
	height      int
	inputHeight int
}

func NewApp(cfg *config.Config, t *translator.Translator) *tea.Program {
	keys := NewKeyMap()

	ta := textarea.New()
	ta.Placeholder = "在这里输入你要翻译的内容。（Enter 翻译，Ctrl+J 启用换行）"
	ta.Focus()
	ta.Prompt = ""
	ta.ShowLineNumbers = false
	ta.SetWidth(60)
	ta.SetHeight(5)
	ta.SetStyles(textarea.DefaultStyles(true))

	ta.KeyMap.InsertNewline.SetKeys("ctrl+j")
	ta.KeyMap.InsertNewline.SetEnabled(true)

	vp := viewport.New(viewport.WithWidth(50), viewport.WithHeight(20))
	vp.SetContent("")

	sp := spinner.New()
	sp.Spinner = spinner.MiniDot
	sp.Style = lipgloss.NewStyle().Foreground(lipgloss.Color("#8B5CF6"))

	hp := help.New()

	m := model{
		config:     cfg,
		translator: t,
		messages:   make([]ChatMessage, 0),
		input:      ta,
		viewport:   vp,
		spinner:    sp,
		help:       hp,
		keys:       keys,
		targetLang: getDefaultLang(cfg),
	}

	p := tea.NewProgram(m)
	return p
}

func getDefaultLang(cfg *config.Config) string {
	if cfg != nil && cfg.DefaultTargetLang != "" {
		return cfg.DefaultTargetLang
	}
	return "zh-CN"
}

func (m model) Init() tea.Cmd {
	return m.spinner.Tick
}

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

	switch msg := msg.(type) {
	case tea.WindowSizeMsg:
		m.width = msg.Width
		m.height = msg.Height
		m.help.SetWidth(msg.Width)
		m.updateLayout()

	case translateMsg:
		m.loading = false
		if msg.err != nil {
			m.addErrorMessage(msg.err.Error())
		} else {
			m.addSuccessMessage(msg.result, msg.tokens)
		}
		m.updateViewportContent()

	case tea.KeyPressMsg:
		switch msg.String() {
		case "ctrl+c", "esc":
			return m, tea.Quit
		case "ctrl+l":
			m.input.Reset()
		case "ctrl+t":
			m.langIndex = (m.langIndex + 1) % len(supportedLangs)
			m.targetLang = supportedLangs[m.langIndex]
		case "enter":
			text := strings.TrimSpace(m.input.Value())
			if text != "" && !m.loading {
				m.input.Reset()
				m.lastInput = text
				m.loading = true
				cmds = append(cmds, m.doTranslate(text, m.targetLang))
			}
		case "alt+up":
			m.viewport.ScrollUp(3)
		case "alt+down":
			m.viewport.ScrollDown(3)
		case "home":
			m.viewport.GotoTop()
		case "end":
			m.viewport.GotoBottom()
		}
		if key.Matches(msg, m.keys.Help) {
			m.help.ShowAll = !m.help.ShowAll
		}
	}

	m.input, cmd = m.input.Update(msg)
	cmds = append(cmds, cmd)

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

	m.spinner, cmd = m.spinner.Update(msg)
	cmds = append(cmds, cmd)

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

func (m *model) addSuccessMessage(result string, tokens int) {
	msg := ChatMessage{
		Input:     m.lastInput,
		Output:    result,
		ToLang:    m.targetLang,
		Model:     getModelName(m.config),
		Tokens:    tokens,
		Timestamp: time.Now(),
	}
	m.messages = append(m.messages, msg)
}

func (m *model) addErrorMessage(err string) {
	msg := ChatMessage{
		Input:     m.lastInput,
		Error:     err,
		ToLang:    m.targetLang,
		Timestamp: time.Now(),
	}
	m.messages = append(m.messages, msg)
}

func (m *model) doTranslate(text, toLang string) tea.Cmd {
	return func() tea.Msg {
		result, err := m.translator.Translate(
			context.Background(),
			text,
			&translator.TranslateOptions{
				ToLang:     toLang,
				PromptName: "simple",
			},
		)
		if err != nil {
			return translateMsg{err: err}
		}
		tokens := 0
		if result.Usage != nil {
			tokens = result.Usage.TotalTokens
		}
		return translateMsg{result: result.Translated, tokens: tokens}
	}
}

func (m *model) updateLayout() {
	if m.width <= 0 || m.height <= 0 {
		return
	}

	contentWidth := m.width - 4
	if contentWidth < 20 {
		contentWidth = 60
	}

	m.input.SetWidth(contentWidth)
	m.viewport.SetWidth(contentWidth)

	helpLines := 2
	if m.help.ShowAll {
		helpLines = 4
	}
	m.viewport.SetHeight(m.height - 12 - helpLines)
	if m.viewport.Height() < 5 {
		m.viewport.SetHeight(10)
	}

	m.updateViewportContent()
}

func (m *model) updateViewportContent() {
	var b strings.Builder

	if len(m.messages) > 0 {
		b.WriteString("\n")
	}

	for _, msg := range m.messages {
		b.WriteString(m.renderTranslationCard(msg))
	}

	m.viewport.SetContent(b.String())
	m.viewport.GotoBottom()
}

func (m *model) renderTranslationCard(msg ChatMessage) string {
	contentWidth := m.viewport.Width() - 2

	timeStr := msg.Timestamp.Format("15:04")
	timeLabel := CardTimeStyle.Render("# " + timeStr)

	inputText := lipgloss.NewStyle().
		Width(contentWidth).
		Render(msg.Input)

	inputBlock := lipgloss.NewStyle().
		Background(lipgloss.Color("#1A1A1A")).
		Padding(1, 3).
		Width(m.viewport.Width()).
		Render(lipgloss.JoinVertical(
			lipgloss.Top,
			timeLabel,
			inputText,
		))

	var outputBlock string
	if msg.Error != "" {
		outputContent := lipgloss.NewStyle().
			Width(m.viewport.Width() - 2).
			Render(msg.Error)
		outputBlock = lipgloss.NewStyle().
			Foreground(lipgloss.Color("#F87171")).
			Padding(1, 3, 1, 3).
			Width(m.viewport.Width()).
			Render(outputContent)
	} else {
		outputContent := lipgloss.NewStyle().
			Width(m.viewport.Width() - 2).
			Render(msg.Output)
		outputBlock = lipgloss.NewStyle().
			Padding(1, 3, 1, 3).
			Width(m.viewport.Width()).
			Render(outputContent)
	}

	footerContent := lipgloss.NewStyle().
		MarginLeft(3).
		Render(lipgloss.JoinHorizontal(
			lipgloss.Left,
			CardFooterIconStyle.Render("▣"),
			CardFooterTextStyle.Render(msg.Model+" · "+strconv.Itoa(msg.Tokens)),
		))

	return CardStyle.Render(
		lipgloss.JoinVertical(
			lipgloss.Top,
			inputBlock,
			outputBlock,
			footerContent,
		),
	) + "\n"
}

func getModelName(cfg *config.Config) string {
	if cfg != nil && cfg.DefaultModel != "" {
		return cfg.DefaultModel
	}
	return "gpt-3.5-turbo"
}

func (m model) View() tea.View {
	if m.width == 0 {
		return tea.NewView("正在加载...")
	}

	header := m.renderHeader()
	messages := m.viewport.View()
	inputArea := m.renderInputArea()
	infoBar := m.renderInfoBar()
	helpView := m.help.View(m.keys)

	content := header + "\n" + messages + inputArea + infoBar + helpView
	v := tea.NewView(content)
	v.AltScreen = true
	return v
}

func (m model) renderHeader() string {
	title := logo.GradientText(logo.GetLogoPattern(), "#B413DC", "#00C8C8")
	titleWithVersion := title + logo.GetVersionSuffix()

	width := m.width - 4
	if width < 20 {
		width = 60
	}

	return lipgloss.NewStyle().
		Width(width).
		Render(titleWithVersion)
}

func (m model) renderInputArea() string {
	inputView := m.input.View()
	separator := lipgloss.NewStyle().
		Foreground(lipgloss.Color("#8B5CF6")).
		Render(":::")
	return "\n" + separator + "\n" + inputView + "\n"
}

func (m model) renderInfoBar() string {
	var recordCount int
	if m.translator != nil && m.translator.GetCache() != nil {
		stats, _ := m.translator.GetCache().Stats(context.Background())
		if stats != nil {
			recordCount = stats.TotalRecords
		}
	}

	var countInfo string
	if recordCount == 0 {
		countInfo = "暂无记录"
	} else {
		countInfo = fmt.Sprintf("%d条", recordCount)
	}

	sep := lipgloss.NewStyle().Foreground(lipgloss.Color("#8B5CF6")).Render("  ")
	separator := lipgloss.NewStyle().Foreground(lipgloss.Color("#8B5CF6")).Render(":::")
	lang := lipgloss.NewStyle().Foreground(lipgloss.Color("#F87171")).Render(m.targetLang)
	model := lipgloss.NewStyle().Foreground(lipgloss.Color("#FAFAFA")).Render(getModelName(m.config))
	count := lipgloss.NewStyle().Foreground(lipgloss.Color("#6B7280")).Render(countInfo)

	result := separator + sep + lang + sep + model + sep + count

	if m.loading {
		result += sep + m.spinner.View() + " 翻译中..."
	}

	return result + "\n"
}

func (m model) renderSpinner() string {
	return ""
}

func (m model) getStatusText() string {
	if m.loading {
		return "翻译中..."
	}
	return "就绪"
}