claude-statusline/main.go

package main

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

	"github.com/go-git/go-git/v6"
	"github.com/shirou/gopsutil/v4/process"
)

// 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 formatted status line.
func buildStatusLine(data *StatusInput) 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
}

// run reads JSON from r, builds the statusline, and writes it to w.
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)
	}

	w.WriteString(buildStatusLine(data))
	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)
}

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

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