package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"os"
)

// Pretty formats one JSON value using two-space indentation.
func Pretty(src []byte) ([]byte, error) {
	decoder := json.NewDecoder(bytes.NewReader(src))
	decoder.UseNumber()

	var value any
	if err := decoder.Decode(&value); err != nil {
		return nil, fmt.Errorf("invalid JSON: %w", err)
	}
	if err := ensureEOF(decoder); err != nil {
		return nil, err
	}

	var out bytes.Buffer
	encoder := json.NewEncoder(&out)
	encoder.SetEscapeHTML(false)
	encoder.SetIndent("", "  ")
	if err := encoder.Encode(value); err != nil {
		return nil, err
	}

	return out.Bytes(), nil
}

// ensureEOF rejects input containing more than one JSON value.
func ensureEOF(decoder *json.Decoder) error {
	var extra any
	if err := decoder.Decode(&extra); err != io.EOF {
		if err == nil {
			return fmt.Errorf("invalid JSON: multiple values")
		}
		return fmt.Errorf("invalid JSON: %w", err)
	}
	return nil
}

func main() {
	input := io.Reader(os.Stdin)
	if len(os.Args) > 2 {
		fmt.Fprintln(os.Stderr, "usage: jsonpretty [file]")
		os.Exit(2)
	}
	if len(os.Args) == 2 {
		file, err := os.Open(os.Args[1])
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		defer file.Close()
		input = file
	}

	src, err := io.ReadAll(input)
	if err == nil {
		var formatted []byte
		formatted, err = Pretty(src)
		if err == nil {
			_, err = os.Stdout.Write(formatted)
		}
	}
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}