iop/apps/edge/internal/opsconsole/status.go

package opsconsole

import (
	"context"
	"fmt"
	"io"
	"sort"
	"strconv"
	"strings"
	"time"

	edgeservice "iop/apps/edge/internal/service"
	iop "iop/proto/gen/iop"
)

func BuildNodeCommandRequest(adapter, targetName, sessionID string, timeoutSec int) (*iop.NodeCommandRequest, string) {
	req := edgeservice.BuildUsageStatusRequest(adapter, targetName, sessionID, timeoutSec)
	return req, req.GetRequestId()
}

func StatusWaitTimeout(req *iop.NodeCommandRequest) time.Duration {
	return edgeservice.StatusWaitTimeout(req)
}

func SendStatus(ctx context.Context, edgeSvc *edgeservice.Service, out io.Writer, target *TargetState) error {
	snap, err := edgeSvc.ResolveNodeSnapshot(target.NodeRef)
	if err != nil {
		return err
	}
	sessionID := edgeservice.NormalizeSessionID(target.SessionID)
	fmt.Fprintf(out, "[edge] sent command=status node=%s adapter=%s target=%s session=%s\n", snap.Label, target.Adapter, target.Target, sessionID)

	result, err := edgeSvc.UsageStatus(ctx, edgeservice.UsageStatusRequest{
		NodeRef:    target.NodeRef,
		Adapter:    target.Adapter,
		Target:     target.Target,
		SessionID:  target.SessionID,
		TimeoutSec: target.TimeoutSec,
	})
	if err != nil {
		return err
	}

	FormatUsageStatus(out, result.NodeLabel, result.Adapter, result.Target, result.SessionID, result.UsageStatus)
	return nil
}

// SendCapabilities dispatches a CAPABILITIES command and renders the result.
func SendCapabilities(ctx context.Context, edgeSvc *edgeservice.Service, out io.Writer, target *TargetState) error {
	return sendNodeCommandView(ctx, out, target, "capabilities",
		edgeSvc.Capabilities,
	)
}

// SendSessionList dispatches a SESSION_LIST command and renders the result.
func SendSessionList(ctx context.Context, edgeSvc *edgeservice.Service, out io.Writer, target *TargetState) error {
	return sendNodeCommandView(ctx, out, target, "sessions",
		edgeSvc.SessionList,
	)
}

// SendTransportStatus dispatches a TRANSPORT_STATUS command and renders the result.
func SendTransportStatus(ctx context.Context, edgeSvc *edgeservice.Service, out io.Writer, target *TargetState) error {
	return sendNodeCommandView(ctx, out, target, "transport",
		edgeSvc.TransportStatus,
	)
}

type nodeCommandFunc func(context.Context, edgeservice.NodeCommandRequestSpec) (edgeservice.NodeCommandView, error)

func sendNodeCommandView(ctx context.Context, out io.Writer, target *TargetState, label string, send nodeCommandFunc) error {
	view, err := send(ctx, edgeservice.NodeCommandRequestSpec{
		NodeRef:    target.NodeRef,
		Adapter:    target.Adapter,
		Target:     target.Target,
		SessionID:  target.SessionID,
		TimeoutSec: target.TimeoutSec,
	})
	if err != nil {
		return err
	}
	FormatNodeCommandView(out, label, view)
	return nil
}

// FormatNodeCommandView renders a node command result map. For SESSION_LIST
// results with structured session.N.* keys, a human-readable grouped layout
// is used. All other commands fall back to stable sorted key output.
func FormatNodeCommandView(out io.Writer, label string, view edgeservice.NodeCommandView) {
	fmt.Fprintf(out, "[%s-%s] adapter=%s target=%s session=%s\n", view.NodeLabel, label, view.Adapter, view.Target, view.SessionID)
	if len(view.Result) == 0 {
		fmt.Fprintf(out, "no %s payload returned\n", label)
		return
	}
	if label == "sessions" && hasSessionKeys(view.Result) {
		renderStructuredSessions(out, view.Result)
		return
	}
	keys := make([]string, 0, len(view.Result))
	for k := range view.Result {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	for _, k := range keys {
		fmt.Fprintf(out, "  %s = %s\n", k, view.Result[k])
	}
}

func hasSessionKeys(result map[string]string) bool {
	_, ok := result["session.0.mode"]
	return ok
}

func renderStructuredSessions(out io.Writer, result map[string]string) {
	countStr := result["count"]
	count := 0
	if n, err := strconv.Atoi(countStr); err == nil {
		count = n
	}
	fmt.Fprintf(out, "sessions: %s\n", countStr)
	for i := 0; i < count; i++ {
		prefix := fmt.Sprintf("session.%d.", i)
		fmt.Fprintf(out, "  [%d] mode=%s target=%s session=%s\n",
			i, result[prefix+"mode"], result[prefix+"target"], result[prefix+"session_id"])
	}
}

func FormatUsageStatus(out io.Writer, nodeAlias, adapterName, targetName, sessionID string, status *iop.AgentUsageStatus) {
	fmt.Fprintf(out, "[%s-status] adapter=%s target=%s session=%s\n", nodeAlias, adapterName, targetName, sessionID)

	if status == nil {
		fmt.Fprintln(out, "no usage status provided")
		return
	}

	hasParsedLimits := false
	dailyLabel := usageStatusLabel(status.GetMetadata(), "daily_label", "Daily limit")
	if status.GetDailyLimit() != "" {
		fmt.Fprintf(out, "%s: %s remaining%s\n", dailyLabel, status.GetDailyLimit(), resetSuffix(status.GetDailyResetTime()))
		hasParsedLimits = true
	}
	weeklyLabel := usageStatusLabel(status.GetMetadata(), "weekly_label", "Weekly limit")
	if status.GetWeeklyLimit() != "" {
		fmt.Fprintf(out, "%s: %s remaining%s\n", weeklyLabel, status.GetWeeklyLimit(), resetSuffix(status.GetWeeklyResetTime()))
		hasParsedLimits = true
	}
	metadata := status.GetMetadata()
	if used := usageStatusMetadata(metadata, "used_percent"); used != "" {
		fmt.Fprintf(out, "Used: %s\n", used)
	}
	if limit := usageStatusMetadata(metadata, "usage_limit"); limit != "" {
		fmt.Fprintf(out, "Usage limit: %s\n", limit)
	}
	if formatModelUsage(out, metadata) {
		hasParsedLimits = true
	}
	if formatClaudeStatus(out, metadata) {
		hasParsedLimits = true
	}

	if !hasParsedLimits {
		parseStatus := usageStatusMetadata(metadata, "parse_status")
		if status.GetRawOutput() == "" {
			if parseStatus != "" {
				fmt.Fprintf(out, "(parse_status=%s) raw output did not include parsed limits and was empty\n", parseStatus)
			} else {
				fmt.Fprintln(out, "raw output did not include parsed limits and was empty")
			}
		} else {
			if parseStatus != "" {
				fmt.Fprintf(out, "(parse_status=%s) raw output did not include parsed limits:\n", parseStatus)
			} else {
				fmt.Fprintln(out, "raw output did not include parsed limits:")
			}
			lines := strings.Split(status.GetRawOutput(), "\n")
			for i, line := range lines {
				if i >= 5 {
					fmt.Fprintln(out, "...")
					break
				}
				fmt.Fprintln(out, line)
			}
		}
	}
}

func formatClaudeStatus(out io.Writer, metadata map[string]string) bool {
	if usageStatusMetadata(metadata, "status_kind") != "claude" {
		return false
	}
	rows := []struct {
		label string
		key   string
	}{
		{"version", "claude_status_version"},
		{"login_method", "claude_status_login_method"},
		{"model", "claude_status_model"},
		{"cwd", "claude_status_cwd"},
		{"session_id", "claude_status_session_id"},
	}
	printed := false
	fmt.Fprintln(out, "Claude status:")
	for _, row := range rows {
		if value := usageStatusMetadata(metadata, row.key); value != "" {
			fmt.Fprintf(out, "  %s = %s\n", row.label, value)
			printed = true
		}
	}
	return printed
}

func usageStatusLabel(metadata map[string]string, key, fallback string) string {
	if metadata != nil {
		if label, ok := metadata[key]; ok {
			return label
		}
	}
	return fallback
}

func usageStatusMetadata(metadata map[string]string, key string) string {
	if metadata == nil {
		return ""
	}
	return metadata[key]
}

func formatModelUsage(out io.Writer, metadata map[string]string) bool {
	countRaw := usageStatusMetadata(metadata, "model_usage_count")
	if countRaw == "" {
		return false
	}
	count, err := strconv.Atoi(countRaw)
	if err != nil || count <= 0 {
		return false
	}

	printed := false
	for i := 0; i < count; i++ {
		prefix := fmt.Sprintf("model_usage_%d", i)
		name := usageStatusMetadata(metadata, prefix+"_name")
		if name == "" {
			continue
		}
		used := usageStatusMetadata(metadata, prefix+"_used_percent")
		reset := usageStatusMetadata(metadata, prefix+"_reset")
		if !printed {
			fmt.Fprintln(out, "Model usage:")
			printed = true
		}
		switch {
		case used != "" && reset != "":
			fmt.Fprintf(out, "  %s: %s used (resets %s)\n", name, used, reset)
		case used != "":
			fmt.Fprintf(out, "  %s: %s used\n", name, used)
		case reset != "":
			fmt.Fprintf(out, "  %s: resets %s\n", name, reset)
		default:
			fmt.Fprintf(out, "  %s\n", name)
		}
	}
	return printed
}

func resetSuffix(resetTime string) string {
	if resetTime == "" {
		return ""
	}
	return " (resets " + resetTime + ")"
}