yoyo/internal/tui/model.go

package tui

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

	"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/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
	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 = "输入要翻译的文本... (Ctrl+J 换行)"
	ta.Focus()
	ta.Prompt = ""
	ta.ShowLineNumbers = false
	ta.SetWidth(60)
	ta.SetHeight(5)
	ta.SetStyles(textarea.DefaultStyles(false))

	ta.KeyMap.InsertNewline.SetEnabled(true)

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

	return tea.NewProgram(model{
		config:     cfg,
		translator: t,
		messages:   make([]ChatMessage, 0),
		input:      ta,
		viewport:   vp,
		keys:       keys,
		targetLang: getDefaultLang(cfg),
	})
}

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

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

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.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()
		}
	}

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

	m.viewport, cmd = m.viewport.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)
	m.viewport.SetHeight(m.height - 12)
	if m.viewport.Height() < 5 {
		m.viewport.SetHeight(10)
	}

	m.updateViewportContent()
}

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

	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 {
	metaContent := lipgloss.JoinHorizontal(
		lipgloss.Left,
		CardMetaStyle.Render(fmt.Sprintf("Tokens: %d", msg.Tokens)),
		CardMetaSeparatorStyle,
		CardMetaStyle.Render(fmt.Sprintf("耗时: %s", msg.Timestamp.Format("15:04:05"))),
		CardMetaSeparatorStyle,
		CardMetaStyle.Render(fmt.Sprintf("模型: %s", msg.Model)),
	)

	metaBlock := lipgloss.NewStyle().
		BorderStyle(lipgloss.Border{
			Top:    "─",
			Bottom: "─",
			Left:   "│",
			Right:  "│",
		}).
		BorderForeground(lipgloss.Color("#374151")).
		Width(m.viewport.Width() - 2).
		Render(metaContent)

	inputContent := lipgloss.NewStyle().
		Width(m.viewport.Width() - 2).
		Render(msg.Input)

	inputBlock := lipgloss.NewStyle().
		Background(lipgloss.Color("#1A1A1A")).
		Width(m.viewport.Width()).
		Render(inputContent)

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

	return CardStyle.Render(
		lipgloss.JoinVertical(
			lipgloss.Top,
			metaBlock,
			inputBlock,
			outputBlock,
		),
	) + "\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()
	statusBar := m.renderStatusBar()

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

func (m model) renderHeader() string {
	title := lipgloss.NewStyle().
		Foreground(lipgloss.Color("#8B5CF6")).
		Bold(true).
		Render("✦ YOYO 翻译")

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

	right := lipgloss.NewStyle().
		Foreground(lipgloss.Color("#6B7280")).
		Render("[Ctrl+C 退出]")

	return lipgloss.NewStyle().
		Width(width).
		Render(title + strings.Repeat(" ", width-len(title)-len(right)-1) + right)
}

func (m model) renderInputArea() string {
	inputView := m.input.View()

	container := lipgloss.NewStyle().
		Width(m.input.Width() + 1).
		BorderStyle(lipgloss.Border{
			Top:    "─",
			Bottom: "─",
			Left:   "│",
			Right:  "│",
		}).
		BorderForeground(lipgloss.Color("#60A5FA"))

	return "\n" + container.Render(inputView) + "\n"
}

func (m model) renderStatusBar() string {
	langInfo := "目标: " + m.targetLang
	modelInfo := "模型: " + getModelName(m.config)
	tokensInfo := "Tokens: -"
	if len(m.messages) > 0 {
		lastMsg := m.messages[len(m.messages)-1]
		if lastMsg.Tokens > 0 {
			tokensInfo = fmt.Sprintf("Tokens: %d", lastMsg.Tokens)
		}
	}

	statusDot := StatusDotStyle.Render("●")
	if m.loading {
		statusDot = LoadingStyle.Render("○")
	}

	sep := StatusItemStyle.Render(" │ ")

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

	status := StatusItemStyle.Render(langInfo) +
		sep + StatusItemStyle.Render(modelInfo) +
		sep + StatusItemStyle.Render(tokensInfo) +
		sep + statusDot + " " + StatusValueStyle.Render(m.getStatusText())

	return lipgloss.NewStyle().
		Width(width).
		Background(lipgloss.Color("#1F2937")).
		Render(" " + status)
}

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