grokway/internal/tui/model.go

package tui

import (
	"fmt"
	"grokway/internal/tunnel"
	"strings"
	"time"

	"github.com/atotto/clipboard"
	"github.com/charmbracelet/bubbles/viewport"
	tea "github.com/charmbracelet/bubbletea"
	"github.com/charmbracelet/lipgloss"
)

// Styles
var (
	primaryColor   = lipgloss.Color("#7D56F4")
	secondaryColor = lipgloss.Color("#FAFAFA")
	subtleColor    = lipgloss.Color("#626262")
	accentColor    = lipgloss.Color("#FF79C6") // Pinkish
	successColor   = lipgloss.Color("#04B575") // Green
	errorColor     = lipgloss.Color("#FF0000") // Red

	headerStyle = lipgloss.NewStyle().
			Bold(true).
			Foreground(secondaryColor).
			Background(primaryColor).
			Padding(1, 2).
			Align(lipgloss.Center)

	urlLabelStyle = lipgloss.NewStyle().
			Foreground(successColor).
			Bold(true).
			MarginRight(1)

	urlValueStyle = lipgloss.NewStyle().
			Foreground(accentColor).
			Bold(true).
			Underline(true)

	boxStyle = lipgloss.NewStyle().
			Border(lipgloss.RoundedBorder()).
			BorderForeground(primaryColor).
			Padding(0, 1)

	statLabelStyle = lipgloss.NewStyle().
			Foreground(subtleColor).
			MarginRight(1)

	statValueStyle = lipgloss.NewStyle().
			Foreground(secondaryColor).
			Bold(true)
)

type Model struct {
	Client     *tunnel.Client
	LogLines   []string
	TotalBytes int64
	Requests   int
	Viewport   viewport.Model
	Ready      bool
	Width      int
	Height     int
	Copied     bool
}

type LogMsg string
type MetricMsg int64
type ClearCopiedMsg struct{}

func InitialModel(localPort, serverAddr, authToken string) Model {
	c := tunnel.NewClient(serverAddr, localPort, authToken)
	return Model{
		Client:   c,
		LogLines: []string{},
	}
}

func (m Model) Init() tea.Cmd {
	return tea.Batch(
		startTunnel(m.Client),
		waitForLog(m.Client),
		waitForMetric(m.Client),
		tea.EnableMouseAllMotion, // Enable mouse support
	)
}

func startTunnel(c *tunnel.Client) tea.Cmd {
	return func() tea.Msg {
		if err := c.Start(); err != nil {
			return LogMsg(fmt.Sprintf("Error starting tunnel: %v", err))
		}
		return LogMsg("Tunnel started successfully!")
	}
}

func waitForLog(c *tunnel.Client) tea.Cmd {
	return func() tea.Msg {
		msg := <-c.Events
		return LogMsg(msg)
	}
}

func waitForMetric(c *tunnel.Client) tea.Cmd {
	return func() tea.Msg {
		bytes := <-c.Metrics
		return MetricMsg(bytes)
	}
}

func clearCopiedMsg() tea.Cmd {
	return tea.Tick(time.Second*2, func(_ time.Time) tea.Msg {
		return ClearCopiedMsg{}
	})
}

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

	switch msg := msg.(type) {
	// Handle key presses
	case tea.KeyMsg:
		switch msg.String() {
		case "q", "ctrl+c", "esc":
			return m, tea.Quit
		case "c":
			if m.Client.PublicURL != "" {
				clipboard.WriteAll(m.Client.PublicURL)
				m.Copied = true
				cmds = append(cmds, clearCopiedMsg())
			}
		}

	case ClearCopiedMsg:
		m.Copied = false

	case LogMsg:
		// Colorize log line based on content
		line := string(msg)
		styledLine := line
		timestamp := time.Now().Format("15:04:05")
		prefix := lipgloss.NewStyle().Foreground(subtleColor).Render(timestamp + " | ")

		if strings.Contains(strings.ToLower(line), "error") {
			styledLine = lipgloss.NewStyle().Foreground(errorColor).Render(line)
		} else if strings.Contains(strings.ToLower(line), "success") || strings.Contains(strings.ToLower(line), "connected") {
			styledLine = lipgloss.NewStyle().Foreground(successColor).Render(line)
		} else {
			styledLine = lipgloss.NewStyle().Foreground(secondaryColor).Render(line)
		}

		m.LogLines = append(m.LogLines, prefix+styledLine)

		// Keep log buffer reasonable
		if len(m.LogLines) > 1000 {
			m.LogLines = m.LogLines[len(m.LogLines)-1000:]
		}

		cmds = append(cmds, waitForLog(m.Client))

		// Update viewport content and scroll if near bottom
		m.Viewport.SetContent(strings.Join(m.LogLines, "\n"))
		m.Viewport.GotoBottom()

	case MetricMsg:
		m.TotalBytes += int64(msg)
		m.Requests++
		cmds = append(cmds, waitForMetric(m.Client))

	case tea.WindowSizeMsg:
		m.Width = msg.Width
		m.Height = msg.Height

		headerHeight := 10 // Approximate height of header + metrics + borders

		if !m.Ready {
			m.Viewport = viewport.New(msg.Width-4, msg.Height-headerHeight) // -4 for borders/padding
			m.Viewport.YPosition = headerHeight
			m.Ready = true
		} else {
			m.Viewport.Width = msg.Width - 4
			m.Viewport.Height = msg.Height - headerHeight
		}

		// Re-render content with new width if needed
		m.Viewport.SetContent(strings.Join(m.LogLines, "\n"))
	}

	// Handle viewport updates (scrolling, mouse events)
	m.Viewport, cmd = m.Viewport.Update(msg)
	cmds = append(cmds, cmd)

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

func (m Model) View() string {
	if !m.Ready {
		return "\n  Initializing Grokway..."
	}

	// 1. Header
	headerStyle := lipgloss.NewStyle().
		Bold(true).
		Foreground(lipgloss.Color("#FFFFFF")).
		Background(primaryColor).
		Width(m.Width).
		Align(lipgloss.Center)

	header := headerStyle.Render("GROKWAY V2")

	// 2. Connection Info (URL under header)
	leftSide := fmt.Sprintf("%s localhost:%s", urlLabelStyle.Render("Local:"), m.Client.LocalPort)
	arrow := lipgloss.NewStyle().Foreground(subtleColor).Render(" ➜ ")

	displayAddr := m.Client.ServerAddr
	if m.Client.PublicURL != "" {
		displayAddr = m.Client.PublicURL
	}
	rightSide := fmt.Sprintf("%s %s", urlLabelStyle.Render("Remote:"), urlValueStyle.Render(displayAddr))

	// Add copy hint
	copyHint := " (Press 'c' to copy)"
	if m.Copied {
		copyHint = lipgloss.NewStyle().Foreground(successColor).Bold(true).Render(" (COPIED!)")
	} else {
		copyHint = lipgloss.NewStyle().Foreground(subtleColor).Render(copyHint)
	}

	connectionBar := lipgloss.NewStyle().
		Width(m.Width).
		Align(lipgloss.Center).
		Padding(1, 0).
		Render(leftSide + arrow + rightSide + copyHint)

	// 3. Stats Row
	stats := fmt.Sprintf("%s %s   %s   %s %s   %s   %s %s",
		statLabelStyle.Render("Requests:"), statValueStyle.Render(fmt.Sprintf("%d", m.Requests)),
		lipgloss.NewStyle().Foreground(subtleColor).Render("•"),
		statLabelStyle.Render("Data:"), statValueStyle.Render(formatBytes(m.TotalBytes)),
		lipgloss.NewStyle().Foreground(subtleColor).Render("•"),
		statLabelStyle.Render("Status:"), lipgloss.NewStyle().Foreground(successColor).Render("Active"),
	)

	statsBar := lipgloss.NewStyle().
		Width(m.Width).
		Align(lipgloss.Center).
		PaddingBottom(1).
		Render(stats)

	// 4. Log Container
	// Ensure viewport fits. REDUCING HEIGHT BY EXTRA MARGIN (2 lines) to prevent scroll-off
	availableHeight := m.Height - lipgloss.Height(header) - lipgloss.Height(connectionBar) - lipgloss.Height(statsBar) - 2
	if availableHeight < 0 {
		availableHeight = 0
	}

	logBox := boxStyle.
		Width(m.Width - 2).
		Height(availableHeight).
		Render(m.Viewport.View())

	return lipgloss.JoinVertical(lipgloss.Left,
		header,
		connectionBar,
		statsBar,
		logBox,
	)
}

func formatBytes(b int64) string {
	const unit = 1024
	if b < unit {
		return fmt.Sprintf("%d B", b)
	}
	div, exp := int64(unit), 0
	for n := b / unit; n >= unit; n /= unit {
		div *= unit
		exp++
	}
	return fmt.Sprintf("%.1f %cB", float64(b)/float64(div), "KMGTPE"[exp])
}