claude-statusline/main.go

package main

import (
	"bufio"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/go-git/go-git/v6"
	"github.com/shirou/gopsutil/v4/process"
	"golang.org/x/term"
)

const statuslineWidthOffset = 7

// TokenUsage tracks context window token consumption.
type TokenUsage struct {
	InputTokens          int `json:"input_tokens"`
	CacheCreationTokens  int `json:"cache_creation_input_tokens"`
	CacheReadInputTokens int `json:"cache_read_input_tokens"`
}

// StatusInput represents the JSON input from Claude Code.
type StatusInput struct {
	Model struct {
		DisplayName string `json:"display_name"`
	} `json:"model"`
	Workspace struct {
		CurrentDir string `json:"current_dir"`
	} `json:"workspace"`
	ContextWindow struct {
		ContextWindowSize int         `json:"context_window_size"`
		CurrentUsage      *TokenUsage `json:"current_usage"`
	} `json:"context_window"`
}

// ANSI color codes
const (
	reset     = "\033[0m"
	bold      = "\033[1m"
	red       = "\033[31m"
	green     = "\033[32m"
	yellow    = "\033[33m"
	magenta   = "\033[35m"
	cyan      = "\033[36m"
	boldGreen = "\033[1;32m"
)

// readInputFromStdin reads JSON input from stdin.
func readInputFromStdin(r *bufio.Reader) string {
	var input strings.Builder
	for {
		line, err := r.ReadString('\n')
		input.WriteString(line)
		if err != nil {
			break
		}
	}
	return input.String()
}

// parseStatusInput unmarshals JSON string into StatusInput.
func parseStatusInput(jsonStr string) (*StatusInput, error) {
	var data StatusInput
	if err := json.Unmarshal([]byte(jsonStr), &data); err != nil {
		return nil, err
	}
	return &data, nil
}

// buildStatusLine constructs the left and right parts of the status line.
func buildStatusLine(data *StatusInput) (left, right string) {
	contextInfo := formatContextInfo(data.ContextWindow.ContextWindowSize, data.ContextWindow.CurrentUsage)
	dirName := filepath.Base(data.Workspace.CurrentDir)
	giteaStatus := getGiteaStatus()
	gitInfo := getGitInfo(data.Workspace.CurrentDir)

	left = fmt.Sprintf("%s %s%s%s %s➜%s %s%s%s%s",
		giteaStatus,
		magenta, data.Model.DisplayName, reset,
		boldGreen, reset,
		cyan, dirName, reset,
		gitInfo)

	right = fmt.Sprintf("%s%s%s", yellow, contextInfo, reset)
	return left, right
}

// calculatePadding returns the number of spaces needed for padding.
func calculatePadding(leftVisible, rightVisible string, termWidth int) int {
	return max(termWidth-len(leftVisible)-len(rightVisible), 1)
}

// formatOutput combines left, right, and padding into final output.
func formatOutput(left, right string, padding int) string {
	return fmt.Sprintf("%s%s%s", left, strings.Repeat(" ", padding), right)
}

// run reads JSON from r, builds the statusline, and writes it to w.
// Returns an error if the input cannot be parsed.
func run(r *bufio.Reader, w *strings.Builder) error {
	jsonStr := readInputFromStdin(r)

	data, err := parseStatusInput(jsonStr)
	if err != nil {
		return fmt.Errorf("error parsing JSON: %w", err)
	}

	left, right := buildStatusLine(data)

	// Calculate visible lengths (strip ANSI)
	leftVisible := stripANSI(left)
	rightVisible := stripANSI(right)

	// Get terminal width
	termWidth := getTerminalWidth() - statuslineWidthOffset

	// Calculate and apply padding
	padding := calculatePadding(leftVisible, rightVisible, termWidth)
	output := formatOutput(left, right, padding)

	w.WriteString(output)
	return nil
}

func main() {
	var out strings.Builder
	if err := run(bufio.NewReader(os.Stdin), &out); err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}
	fmt.Print(out.String())
}

func formatContextInfo(contextSize int, usage *TokenUsage) string {
	totalK := contextSize / 1000

	if usage == nil {
		return fmt.Sprintf("0/%dk", totalK)
	}

	currentTokens := usage.InputTokens + usage.CacheCreationTokens + usage.CacheReadInputTokens
	currentK := currentTokens / 1000
	return fmt.Sprintf("%dk/%dk", currentK, totalK)
}

// processLister returns the list of running processes.
// Replaced in tests to avoid depending on real process state.
var processLister = process.Processes

func getGiteaStatus() string {
	// Check if gitea process is running using gopsutil (cross-platform)
	procs, err := processLister()
	if err != nil {
		return red + "●" + reset
	}

	for _, p := range procs {
		name, err := p.Name()
		if err != nil {
			continue
		}
		if name == "gitea" {
			return green + "●" + reset
		}
	}
	return red + "●" + reset
}

func getGitInfo(cwd string) string {
	// Open the repository (searches up from cwd)
	repo, err := git.PlainOpenWithOptions(cwd, &git.PlainOpenOptions{
		DetectDotGit: true,
	})
	if err != nil {
		return ""
	}

	// Get HEAD reference
	head, err := repo.Head()
	if err != nil {
		return ""
	}

	// Get branch name
	var branch string
	if head.Name().IsBranch() {
		branch = head.Name().Short()
	} else {
		// Detached HEAD - use short hash
		branch = head.Hash().String()[:7]
	}

	// Check if working tree is dirty
	worktree, err := repo.Worktree()
	if err != nil {
		return fmt.Sprintf(" git:(%s)", branch)
	}

	status, err := worktree.Status()
	if err != nil {
		return fmt.Sprintf(" git:(%s)", branch)
	}

	if !status.IsClean() {
		return fmt.Sprintf(" git:(%s) ✗", branch)
	}
	return fmt.Sprintf(" git:(%s)", branch)
}

// termWidthFunc returns the terminal width.
// Tries stderr and stdin as fallbacks since stdout is typically piped
// when Claude Code captures the status line output.
// Replaced in tests to avoid depending on a real TTY.
var termWidthFunc = func() (int, error) {
	for _, fd := range []uintptr{os.Stderr.Fd(), os.Stdin.Fd(), os.Stdout.Fd()} {
		if width, _, err := term.GetSize(int(fd)); err == nil {
			return width, nil
		}
	}
	// Fall back to COLUMNS env var
	if cols := os.Getenv("COLUMNS"); cols != "" {
		var w int
		if _, err := fmt.Sscanf(cols, "%d", &w); err == nil && w > 0 {
			return w, nil
		}
	}
	return 0, errors.New("no terminal detected")
}

func getTerminalWidth() int {
	width, err := termWidthFunc()
	if err != nil {
		return 80
	}
	return width
}

var ansiRegex = regexp.MustCompile(`\x1b\[[0-9;]*m`)

func stripANSI(s string) string {
	return ansiRegex.ReplaceAllString(s, "")
}