fj/cmd/json.go

package cmd

import (
	"encoding/json"
	"fmt"
	"strings"

	"github.com/itchyny/gojq"
	"github.com/spf13/cobra"
)

// addJSONFlags adds --json, --json-fields, and --jq flags to a command.
// --json is an optional-value string flag:
//   - --json (no value) → output all fields as JSON
//   - --json title,state → output only those fields (gh-compatible)
//
// --json-fields is kept as a backwards-compatible alias.
func addJSONFlags(cmd *cobra.Command, jsonDesc string) {
	f := cmd.Flags()
	f.String("json", "", jsonDesc)
	f.Lookup("json").NoOptDefVal = " " // space sentinel: flag present with no value
	f.String("json-fields", "", "Comma-separated list of JSON fields to include")
	f.String("jq", "", "Filter JSON output using a jq expression")
}

// wantJSON returns true if the user requested JSON output via --json, --json-fields, or --jq.
func wantJSON(cmd *cobra.Command) bool {
	if j, _ := cmd.Flags().GetString("json"); j != "" {
		return true
	}
	if jq, _ := cmd.Flags().GetString("jq"); jq != "" {
		return true
	}
	if f, _ := cmd.Flags().GetString("json-fields"); f != "" {
		return true
	}
	return false
}

// outputJSON writes a value as JSON, respecting --json, --json-fields, and --jq flags.
func outputJSON(cmd *cobra.Command, value any) error {
	jsonVal, _ := cmd.Flags().GetString("json")
	jsonFields, _ := cmd.Flags().GetString("json-fields")
	jqExpr, _ := cmd.Flags().GetString("jq")

	fields := ""
	jsonVal = strings.TrimSpace(jsonVal)
	if jsonVal != "" {
		fields = jsonVal
	} else if jsonFields != "" {
		fields = jsonFields
	}

	return writeJSONFiltered(value, fields, jqExpr)
}

// writeJSON writes a value as pretty-printed JSON to ios.Out.
func writeJSON(value any) error {
	enc := json.NewEncoder(ios.Out)
	enc.SetIndent("", "  ")
	return enc.Encode(value)
}

// writeJSONFiltered writes a value as JSON, optionally selecting specific fields
// and/or applying a jq expression. If fields is empty and jqExpr is empty, it
// writes the full value.
func writeJSONFiltered(value any, fields string, jqExpr string) error {
	// If no filtering, just write the full JSON.
	if fields == "" && jqExpr == "" {
		return writeJSON(value)
	}

	// Convert value to a generic interface via JSON round-trip so we can
	// manipulate it with maps/slices.
	raw, err := json.Marshal(value)
	if err != nil {
		return fmt.Errorf("marshaling JSON: %w", err)
	}

	var data any
	if err := json.Unmarshal(raw, &data); err != nil {
		return fmt.Errorf("unmarshaling JSON: %w", err)
	}

	// Apply field selection if specified.
	if fields != "" {
		fieldList := strings.Split(fields, ",")
		for i, f := range fieldList {
			fieldList[i] = strings.TrimSpace(f)
		}
		data = selectFields(data, fieldList)
	}

	// Apply jq expression if specified.
	if jqExpr != "" {
		return applyJQ(data, jqExpr)
	}

	return writeJSON(data)
}

// selectFields filters a JSON value to only include the specified fields.
// Works on both single objects and arrays of objects.
func selectFields(data any, fields []string) any {
	switch v := data.(type) {
	case []any:
		result := make([]any, len(v))
		for i, item := range v {
			result[i] = selectFields(item, fields)
		}
		return result
	case map[string]any:
		result := make(map[string]any)
		for _, field := range fields {
			if val, ok := v[field]; ok {
				result[field] = val
			}
		}
		return result
	default:
		return data
	}
}

// applyJQ applies a jq expression to data and writes each output value.
func applyJQ(data any, expr string) error {
	query, err := gojq.Parse(expr)
	if err != nil {
		return fmt.Errorf("invalid jq expression: %w", err)
	}

	iter := query.Run(data)
	enc := json.NewEncoder(ios.Out)
	enc.SetIndent("", "  ")

	for {
		v, ok := iter.Next()
		if !ok {
			break
		}
		if err, isErr := v.(error); isErr {
			return fmt.Errorf("jq error: %w", err)
		}
		// For string values, print raw (no JSON encoding) to match jq behavior.
		if s, ok := v.(string); ok {
			fmt.Fprintln(ios.Out, s)
		} else {
			if err := enc.Encode(v); err != nil {
				return err
			}
		}
	}

	return nil
}