package main

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

const defaultIndent = "  "

// PrettyJSON validates src as JSON and returns a consistently indented form.
func PrettyJSON(src []byte, indent string) ([]byte, error) {
	var compact bytes.Buffer
	if err := json.Compact(&compact, src); err != nil {
		return nil, err
	}

	var pretty bytes.Buffer
	if err := json.Indent(&pretty, compact.Bytes(), "", indent); err != nil {
		return nil, err
	}
	pretty.WriteByte('\n')
	return pretty.Bytes(), nil
}

func run(input io.Reader, output io.Writer, indent string) error {
	src, err := io.ReadAll(input)
	if err != nil {
		return err
	}

	pretty, err := PrettyJSON(src, indent)
	if err != nil {
		return fmt.Errorf("invalid JSON: %w", err)
	}

	_, err = output.Write(pretty)
	return err
}

// A couple of tiny self-tests keep this single-file program easy to check.
func selfTest() error {
	tests := []struct {
		name string
		in   string
		want string
	}{
		{
			name: "object",
			in:   `{"name":"Ada","tags":["math","code"]}`,
			want: "{\n  \"name\": \"Ada\",\n  \"tags\": [\n    \"math\",\n    \"code\"\n  ]\n}\n",
		},
		{
			name: "array",
			in:   `[true,false,{"n":3}]`,
			want: "[\n  true,\n  false,\n  {\n    \"n\": 3\n  }\n]\n",
		},
	}

	for _, tt := range tests {
		got, err := PrettyJSON([]byte(tt.in), defaultIndent)
		if err != nil {
			return fmt.Errorf("%s: unexpected error: %w", tt.name, err)
		}
		if string(got) != tt.want {
			return fmt.Errorf("%s: got %q, want %q", tt.name, got, tt.want)
		}
	}

	if _, err := PrettyJSON([]byte(`{"bad":`), defaultIndent); err == nil {
		return fmt.Errorf("invalid JSON: expected error")
	}

	return nil
}

func main() {
	indent := flag.String("indent", defaultIndent, "indent string")
	test := flag.Bool("test", false, "run built-in tests")
	flag.Parse()

	if *test {
		if err := selfTest(); err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		fmt.Fprintln(os.Stderr, "ok")
		return
	}

	if err := run(os.Stdin, os.Stdout, *indent); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}