iop/apps/node/internal/adapters/cli/status/parser.go

package status

import (
	"fmt"
	"regexp"
	"strconv"
	"strings"
	"unicode"
)

var (
	dailyLimitRegex  = regexp.MustCompile(`5h limit:.*?\]\s*([^%\n]+%)[^\(]*\(resets\s+([^\)]+)\)`)
	weeklyLimitRegex = regexp.MustCompile(`Weekly limit:.*?\]\s*([^%\n]+%)[^\(]*\(resets\s+([^\)]+)\)`)

	// Section header locators — they only locate the start of each block; they
	// do NOT span forward into other blocks. This prevents the previous lazy
	// `.*?` regex from accidentally pairing a session header with the week's
	// "X% used / Resets Y" payload when cursor positioning in the raw PTY
	// stream interleaves the two sections out of order.
	claudeSessionHeaderRegex = regexp.MustCompile(`(?i)Current\s+session\b`)
	claudeWeekHeaderRegex    = regexp.MustCompile(`(?i)Current\s+week\b(?:\s*\(all\s+models\))?`)

	// Payload extractor used *within* a single section slice.
	claudePayloadRegex = regexp.MustCompile(`(?is)(\d+(?:\.\d+)?)%\s+used.*?Resets\s+([^\n]+)`)

	geminiUsedQuotaRegex    = regexp.MustCompile(`(?i)(\d+(?:\.\d+)?)%\s+used\s+\(Limit resets in ([^)]+)\)`)
	geminiSimpleQuotaRegex  = regexp.MustCompile(`(?is)\bquota\s+(\d+(?:\.\d+)?)%\s+used\b`)
	geminiLimitReachedRegex = regexp.MustCompile(`(?i)Limit reached,\s*resets in ([^\n]+)`)
	geminiUsageLimitRegex   = regexp.MustCompile(`(?i)Usage limit:\s*([0-9][0-9,]*)`)
	geminiModelLineRegex    = regexp.MustCompile(`(?i)^(.*)\s+(\d+(?:\.\d+)?)%\s+Resets:\s*(.*)$`)
)

// ParseStatusOutput attempts to parse Claude or Codex usage limits from a raw
// PTY stream. Claude's /usage TUI uses cursor positioning and erase sequences
// to repaint the same panel multiple times, so the stream-append order can
// differ from what is actually visible on screen. We render the raw bytes
// through a terminal screen model first (renderVisibleScreen) and try Claude
// parsing against that; if it fails or only partially succeeds we fall back to
// the line-oriented cleanANSI which handles the simpler Codex output and
// non-cursor-repainted Claude streams.
func ParseStatusOutput(raw string) (*UsageStatus, error) {
	visible := renderVisibleScreen(raw, 60, 200)
	cleanRaw := cleanANSI(raw)

	// Prefer whichever input yields a more complete Claude parse: a full
	// session+week match beats a partial one, which beats nothing.
	visStatus := &UsageStatus{RawOutput: visible}
	rawStatus := &UsageStatus{RawOutput: cleanRaw}
	visStatus.parseClaudeUsage(visible)
	rawStatus.parseClaudeUsage(cleanRaw)

	if pick := pickMoreComplete(visStatus, rawStatus); pick != nil {
		return pick, nil
	}

	// Neither input produced Claude data — try Gemini, then Codex.
	parsedGemini := rawStatus.parseGeminiQuota(cleanRaw)
	if !rawStatus.parseGeminiModelUsage(cleanRaw) {
		rawStatus.parseGeminiModelUsage(visible)
	}
	if parsedGemini || rawStatus.Metadata["model_usage_count"] != "" {
		return rawStatus, nil
	}
	rawStatus.parseCodexUsage(cleanRaw)
	return rawStatus, nil
}

// parseGeminiQuota parses Gemini `/stats model` daily quota output.
// `0% used (Limit resets in 24h)` and `Usage limit: 200` populate
// DailyLimit (remaining percent), DailyResetTime, and metadata. Real
// screen-reader output may also expose the header as `quota` followed by
// `0% used` without reset or limit details; in that case only remaining and
// used percentages are populated.
// `Limit reached, resets in 5h` is treated as 100% used.
func (s *UsageStatus) parseGeminiQuota(cleanRaw string) bool {
	used := ""
	reset := ""
	if m := geminiUsedQuotaRegex.FindStringSubmatch(cleanRaw); len(m) >= 3 {
		used = m[1] + "%"
		reset = strings.TrimSpace(m[2])
	} else if m := geminiLimitReachedRegex.FindStringSubmatch(cleanRaw); len(m) >= 2 {
		used = "100%"
		reset = strings.TrimSpace(m[1])
	} else if m := geminiSimpleQuotaRegex.FindStringSubmatch(cleanRaw); len(m) >= 2 {
		used = m[1] + "%"
	}

	limitNumeric := ""
	if m := geminiUsageLimitRegex.FindStringSubmatch(cleanRaw); len(m) >= 2 {
		limitNumeric = strings.TrimSpace(m[1])
	}

	if used == "" {
		return false
	}

	if s.Metadata == nil {
		s.Metadata = make(map[string]string)
	}
	s.Metadata["daily_label"] = "Daily quota"
	s.DailyLimit = remainingPercent(used)
	s.DailyResetTime = reset
	s.Metadata["used_percent"] = used
	if limitNumeric != "" {
		s.Metadata["usage_limit"] = limitNumeric
	}
	return true
}

// parseGeminiModelUsage parses the `/model` screen-reader "Model usage" area,
// storing per-model used percentages and reset times in metadata.
func (s *UsageStatus) parseGeminiModelUsage(cleanRaw string) bool {
	lower := strings.ToLower(cleanRaw)
	start := strings.LastIndex(lower, "model usage")
	if start < 0 {
		return false
	}
	block := cleanRaw[start+len("model usage"):]
	blockLower := strings.ToLower(block)
	if end := strings.Index(blockLower, "(press esc"); end >= 0 {
		block = block[:end]
	}

	type modelUsage struct {
		name        string
		usedPercent string
		reset       string
	}

	var models []modelUsage
	modelIndex := make(map[string]int)
	lines := strings.Split(block, "\n")
	for i := 0; i < len(lines); i++ {
		line := strings.TrimSpace(lines[i])
		if line == "" || !strings.Contains(line, "%") || !strings.Contains(strings.ToLower(line), "resets:") {
			continue
		}

		m := geminiModelLineRegex.FindStringSubmatch(line)
		if len(m) < 4 {
			continue
		}

		name := geminiModelNameFromUsagePrefix(m[1])
		if name == "" {
			continue
		}
		used := strings.TrimSpace(m[2]) + "%"
		reset := strings.TrimSpace(m[3])
		if !strings.Contains(reset, ")") && i+1 < len(lines) {
			next := strings.TrimSpace(lines[i+1])
			if strings.HasPrefix(next, "(") {
				reset = strings.TrimSpace(reset + " " + next)
				i++
			}
		}

		entry := modelUsage{name: name, usedPercent: used, reset: reset}
		if idx, ok := modelIndex[name]; ok {
			models[idx] = entry
			continue
		}
		modelIndex[name] = len(models)
		models = append(models, entry)
	}

	if len(models) == 0 {
		return false
	}

	if s.Metadata == nil {
		s.Metadata = make(map[string]string)
	}
	s.Metadata["model_usage_count"] = strconv.Itoa(len(models))
	for i, model := range models {
		prefix := fmt.Sprintf("model_usage_%d", i)
		s.Metadata[prefix+"_name"] = model.name
		s.Metadata[prefix+"_used_percent"] = model.usedPercent
		if model.reset != "" {
			s.Metadata[prefix+"_reset"] = model.reset
		}
	}
	return true
}

func geminiModelNameFromUsagePrefix(prefix string) string {
	prefix = strings.TrimSpace(prefix)
	if prefix == "" {
		return ""
	}
	if idx := strings.IndexAny(prefix, "▬━─█░#▉▊▋▌▍▎▏■"); idx >= 0 {
		return strings.TrimSpace(prefix[:idx])
	}
	return prefix
}

// pickMoreComplete returns the status with more fields populated, or nil if
// neither has any Claude limit data.
func pickMoreComplete(a, b *UsageStatus) *UsageStatus {
	score := func(s *UsageStatus) int {
		n := 0
		if s.DailyLimit != "" {
			n++
		}
		if s.WeeklyLimit != "" {
			n++
		}
		if s.DailyResetTime != "" {
			n++
		}
		if s.WeeklyResetTime != "" {
			n++
		}
		return n
	}
	sa, sb := score(a), score(b)
	if sa == 0 && sb == 0 {
		return nil
	}
	if sa >= sb {
		return a
	}
	return b
}

// parseClaudeUsage attempts to parse Claude /usage output format.
//
// We locate each section header (Current session / Current week) and extract
// the "X% used … Resets Y" payload that immediately follows it, stopping at
// the next opposing header so the session header is never paired with the
// week's payload. Each header may appear multiple times in the raw stream
// (initial loading repaint then data repaint); we walk the occurrences from
// last to first and take the first one whose section contains a payload, so
// the most recent render wins.
func (s *UsageStatus) parseClaudeUsage(cleanRaw string) bool {
	sessionLocs := claudeSessionHeaderRegex.FindAllStringIndex(cleanRaw, -1)
	weekLocs := claudeWeekHeaderRegex.FindAllStringIndex(cleanRaw, -1)

	extract := func(headerLocs [][]int, oppositeLocs [][]int) (string, string, bool) {
		for i := len(headerLocs) - 1; i >= 0; i-- {
			start := headerLocs[i][1]
			end := len(cleanRaw)
			for _, opp := range oppositeLocs {
				if opp[0] > start && opp[0] < end {
					end = opp[0]
				}
			}
			if end <= start {
				continue
			}
			section := cleanRaw[start:end]
			if m := claudePayloadRegex.FindStringSubmatch(section); len(m) >= 3 {
				return m[1], strings.TrimSpace(m[2]), true
			}
		}
		return "", "", false
	}

	if used, reset, ok := extract(sessionLocs, weekLocs); ok {
		if s.Metadata == nil {
			s.Metadata = make(map[string]string)
		}
		s.Metadata["daily_label"] = "Current session"
		s.DailyLimit = remainingPercent(used)
		s.DailyResetTime = reset
	}
	if used, reset, ok := extract(weekLocs, sessionLocs); ok {
		if s.Metadata == nil {
			s.Metadata = make(map[string]string)
		}
		s.Metadata["weekly_label"] = "Current week"
		s.WeeklyLimit = remainingPercent(used)
		s.WeeklyResetTime = reset
	}
	return s.DailyLimit != "" || s.WeeklyLimit != ""
}

// parseCodexUsage attempts to parse Codex output format and sets metadata labels.
func (s *UsageStatus) parseCodexUsage(cleanRaw string) {
	if match := dailyLimitRegex.FindStringSubmatch(cleanRaw); len(match) > 2 {
		s.DailyLimit = strings.TrimSpace(match[1])
		s.DailyResetTime = strings.TrimSpace(match[2])
	}
	if match := weeklyLimitRegex.FindStringSubmatch(cleanRaw); len(match) > 2 {
		s.WeeklyLimit = strings.TrimSpace(match[1])
		s.WeeklyResetTime = strings.TrimSpace(match[2])
	}
	if s.DailyLimit != "" || s.WeeklyLimit != "" {
		if s.Metadata == nil {
			s.Metadata = make(map[string]string)
		}
		if s.DailyLimit != "" {
			s.Metadata["daily_label"] = "5h limit"
		}
		if s.WeeklyLimit != "" {
			s.Metadata["weekly_label"] = "Weekly limit"
		}
	}
}

// remainingPercent converts a "used" percentage string to remaining percentage.
func remainingPercent(usedStr string) string {
	usedStr = strings.TrimSuffix(usedStr, "%")
	usedStr = strings.TrimSpace(usedStr)
	usedVal, err := strconv.ParseFloat(usedStr, 64)
	if err != nil {
		return usedStr
	}
	remaining := 100 - usedVal
	if remaining < 0 {
		remaining = 0
	}
	if usedVal == float64(int(usedVal)) {
		return fmt.Sprintf("%d%%", int(remaining))
	}
	return fmt.Sprintf("%.1f%%", remaining)
}

// RenderVisibleScreen runs the raw PTY stream through a small terminal screen
// model and returns the visible text.
func RenderVisibleScreen(raw string, rows, cols int) string {
	return renderVisibleScreen(raw, rows, cols)
}

// renderVisibleScreen runs the raw PTY stream through a small terminal screen
// model (cursor + grid) and returns the visible text. Unlike cleanANSI which
// strips control sequences in place, this honors cursor positioning and erase
// commands so that text overwritten by later repaints disappears, matching
// what a user would see in a terminal. The grid is sized large enough for
// Claude's /usage panel (60 rows × 200 cols); content past the bottom row
// scrolls up.
//
// Only the sequences Claude's /usage TUI is actually observed to emit are
// implemented — CUP (H/f), CUU/CUD/CUF/CUB (A/B/C/D), CHA (G), ED (J), EL (K),
// plus \r, \n, \b, \t. SGR (m), OSC, and DEC private modes are recognized and
// discarded; unknown CSIs are skipped.
func renderVisibleScreen(raw string, rows, cols int) string {
	scr := newScreen(rows, cols)
	rs := []rune(raw)
	for i := 0; i < len(rs); {
		c := rs[i]
		switch {
		case c == 0x1b:
			i++
			if i >= len(rs) {
				break
			}
			switch rs[i] {
			case '[':
				i++
				priv := byte(0)
				if i < len(rs) && (rs[i] == '?' || rs[i] == '>' || rs[i] == '!') {
					priv = byte(rs[i])
					i++
				}
				start := i
				for i < len(rs) && ((rs[i] >= '0' && rs[i] <= '9') || rs[i] == ';') {
					i++
				}
				params := string(rs[start:i])
				if i >= len(rs) {
					break
				}
				cmd := rs[i]
				i++
				scr.handleCSI(priv, params, cmd)
			case ']':
				i++
				for i < len(rs) {
					if rs[i] == 0x07 {
						i++
						break
					}
					if rs[i] == 0x1b && i+1 < len(rs) && rs[i+1] == '\\' {
						i += 2
						break
					}
					i++
				}
			default:
				i++
			}
		case c == '\r':
			scr.cx = 0
			i++
		case c == '\n':
			scr.cy++
			if scr.cy >= scr.rows {
				scr.scrollUp()
				scr.cy = scr.rows - 1
			}
			i++
		case c == '\b':
			if scr.cx > 0 {
				scr.cx--
			}
			i++
		case c == '\t':
			scr.cx = ((scr.cx / 8) + 1) * 8
			if scr.cx >= scr.cols {
				scr.cx = scr.cols - 1
			}
			i++
		case c < 0x20:
			i++
		default:
			scr.putRune(c)
			i++
		}
	}
	return scr.render()
}

type screen struct {
	rows, cols int
	grid       [][]rune
	cy, cx     int
}

const wideContinuation rune = -1

func newScreen(rows, cols int) *screen {
	g := make([][]rune, rows)
	for i := range g {
		g[i] = make([]rune, cols)
		for j := range g[i] {
			g[i][j] = ' '
		}
	}
	return &screen{rows: rows, cols: cols, grid: g}
}

func (s *screen) clampCursor() {
	if s.cy < 0 {
		s.cy = 0
	}
	if s.cy >= s.rows {
		s.cy = s.rows - 1
	}
	if s.cx < 0 {
		s.cx = 0
	}
	if s.cx >= s.cols {
		s.cx = s.cols - 1
	}
}

func (s *screen) putRune(r rune) {
	w := terminalRuneWidth(r)
	if w <= 0 {
		return
	}
	if s.cx >= s.cols {
		s.wrapLine()
	}
	if w > 1 && s.cx == s.cols-1 {
		s.wrapLine()
	}
	for i := 0; i < w && s.cx+i < s.cols; i++ {
		s.clearCell(s.cy, s.cx+i)
	}
	s.grid[s.cy][s.cx] = r
	if w > 1 && s.cx+1 < s.cols {
		s.grid[s.cy][s.cx+1] = wideContinuation
	}
	s.cx += w
}

func (s *screen) wrapLine() {
	s.cx = 0
	s.cy++
	if s.cy >= s.rows {
		s.scrollUp()
		s.cy = s.rows - 1
	}
}

func (s *screen) clearCell(row, col int) {
	if row < 0 || row >= s.rows || col < 0 || col >= s.cols {
		return
	}
	if s.grid[row][col] == wideContinuation && col > 0 {
		s.grid[row][col-1] = ' '
	}
	if col+1 < s.cols && s.grid[row][col+1] == wideContinuation {
		s.grid[row][col+1] = ' '
	}
	s.grid[row][col] = ' '
}

func terminalRuneWidth(r rune) int {
	if r == 0 || r < 0x20 {
		return 0
	}
	if unicode.Is(unicode.Mn, r) || unicode.Is(unicode.Me, r) {
		return 0
	}
	if unicode.In(r, unicode.Hangul, unicode.Han, unicode.Hiragana, unicode.Katakana) {
		return 2
	}
	switch {
	case r >= 0x1100 && r <= 0x115f:
		return 2
	case r >= 0x2329 && r <= 0x232a:
		return 2
	case r >= 0x2e80 && r <= 0xa4cf:
		return 2
	case r >= 0xac00 && r <= 0xd7a3:
		return 2
	case r >= 0xf900 && r <= 0xfaff:
		return 2
	case r >= 0xfe10 && r <= 0xfe19:
		return 2
	case r >= 0xfe30 && r <= 0xfe6f:
		return 2
	case r >= 0xff00 && r <= 0xff60:
		return 2
	case r >= 0xffe0 && r <= 0xffe6:
		return 2
	default:
		return 1
	}
}

func (s *screen) scrollUp() {
	for i := 0; i < s.rows-1; i++ {
		s.grid[i] = s.grid[i+1]
	}
	blank := make([]rune, s.cols)
	for j := range blank {
		blank[j] = ' '
	}
	s.grid[s.rows-1] = blank
}

func (s *screen) handleCSI(priv byte, params string, cmd rune) {
	if priv != 0 {
		return
	}
	parts := strings.Split(params, ";")
	p1 := csiInt(parts, 0, 0)
	switch cmd {
	case 'H', 'f':
		row := csiInt(parts, 0, 1)
		col := csiInt(parts, 1, 1)
		s.cy = row - 1
		s.cx = col - 1
		s.clampCursor()
	case 'A':
		s.cy -= csiInt(parts, 0, 1)
		s.clampCursor()
	case 'B':
		s.cy += csiInt(parts, 0, 1)
		s.clampCursor()
	case 'C':
		s.cx += csiInt(parts, 0, 1)
		s.clampCursor()
	case 'D':
		s.cx -= csiInt(parts, 0, 1)
		s.clampCursor()
	case 'G':
		s.cx = csiInt(parts, 0, 1) - 1
		s.clampCursor()
	case 'J':
		s.eraseDisplay(p1)
	case 'K':
		s.eraseLine(p1)
	}
}

func csiInt(parts []string, idx, def int) int {
	if idx >= len(parts) || parts[idx] == "" {
		return def
	}
	n, err := strconv.Atoi(parts[idx])
	if err != nil {
		return def
	}
	return n
}

func (s *screen) eraseDisplay(mode int) {
	blankRow := func(i int) {
		for j := range s.grid[i] {
			s.grid[i][j] = ' '
		}
	}
	switch mode {
	case 0:
		for j := s.cx; j < s.cols; j++ {
			s.clearCell(s.cy, j)
		}
		for i := s.cy + 1; i < s.rows; i++ {
			blankRow(i)
		}
	case 1:
		for i := 0; i < s.cy; i++ {
			blankRow(i)
		}
		for j := 0; j <= s.cx && j < s.cols; j++ {
			s.clearCell(s.cy, j)
		}
	case 2, 3:
		for i := 0; i < s.rows; i++ {
			blankRow(i)
		}
	}
}

func (s *screen) eraseLine(mode int) {
	switch mode {
	case 0:
		for j := s.cx; j < s.cols; j++ {
			s.clearCell(s.cy, j)
		}
	case 1:
		for j := 0; j <= s.cx && j < s.cols; j++ {
			s.clearCell(s.cy, j)
		}
	case 2:
		for j := range s.grid[s.cy] {
			s.grid[s.cy][j] = ' '
		}
	}
}

func (s *screen) render() string {
	var b strings.Builder
	for i := range s.grid {
		end := len(s.grid[i])
		for end > 0 && (s.grid[i][end-1] == ' ' || s.grid[i][end-1] == wideContinuation) {
			end--
		}
		for _, r := range s.grid[i][:end] {
			if r == wideContinuation {
				continue
			}
			b.WriteRune(r)
		}
		b.WriteByte('\n')
	}
	return b.String()
}

// cleanANSI removes ANSI escape codes from the string.
func cleanANSI(str string) string {
	// Remove OSC sequences first (before CSI so we don't split them)
	str = regexp.MustCompile(`\x1b\][^\x07]*\x07`).ReplaceAllString(str, "")
	str = regexp.MustCompile(`\x1b\][^\x1b]*\x1b\\`).ReplaceAllString(str, "")

	// Replace cursor-forward escapes with the exact count of spaces
	str = regexp.MustCompile(`\x1b\[(\d+)C`).ReplaceAllStringFunc(str, func(m string) string {
		var count int
		fmt.Sscanf(m, "\x1b[%dC", &count)
		if count > 100 {
			count = 100
		}
		return strings.Repeat(" ", count)
	})

	// Remove standard SGR (Select Graphic Rendition) codes e.g., ESC[31m, ESC[0m
	str = regexp.MustCompile(`\x1b\[[0-9;]*m`).ReplaceAllString(str, "")

	// Remove other CSI sequences
	str = regexp.MustCompile(`\x1b\[[?0-9;]*[a-zA-Z]`).ReplaceAllString(str, "")

	// Normalize carriage returns
	str = strings.ReplaceAll(str, "\r\n", "\n")
	str = strings.ReplaceAll(str, "\r", "\n")

	// Replace remaining C0 control chars (except \n and \t) with spaces
	controlRe := regexp.MustCompile(`[\x00-\x08\x0b\x0c\x0e-\x1f]`)
	str = controlRe.ReplaceAllString(str, " ")

	return str
}