package main

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"strings"
	"testing"
)

// PrettyJSON formats compact or messy JSON with two-space indentation.
func PrettyJSON(input []byte) ([]byte, error) {
	var out bytes.Buffer
	if err := prettyValue(&out, input); err != nil {
		return nil, err
	}
	return out.Bytes(), nil
}

func prettyValue(out *bytes.Buffer, input []byte) error {
	indent := 0
	inString := false
	escaped := false
	needSpace := false

	for i, b := range input {
		if inString {
			out.WriteByte(b)
			if escaped {
				escaped = false
			} else if b == '\\' {
				escaped = true
			} else if b == '"' {
				inString = false
			}
			continue
		}

		switch b {
		case ' ', '\n', '\r', '\t':
			continue
		case '"':
			if needSpace {
				out.WriteByte(' ')
				needSpace = false
			}
			inString = true
			out.WriteByte(b)
		case '{', '[':
			if needSpace {
				out.WriteByte(' ')
				needSpace = false
			}
			out.WriteByte(b)
			indent++
			writeNewline(out, indent)
		case '}', ']':
			if indent == 0 {
				return fmt.Errorf("unexpected %q at byte %d", b, i)
			}
			indent--
			writeNewline(out, indent)
			out.WriteByte(b)
		case ',':
			out.WriteByte(b)
			writeNewline(out, indent)
			needSpace = false
		case ':':
			out.WriteString(": ")
			needSpace = false
		default:
			if needSpace {
				out.WriteByte(' ')
			}
			out.WriteByte(b)
			needSpace = false
		}
	}

	if inString {
		return fmt.Errorf("unterminated string")
	}
	if indent != 0 {
		return fmt.Errorf("unclosed object or array")
	}
	out.WriteByte('\n')
	return nil
}

func writeNewline(out *bytes.Buffer, indent int) {
	out.WriteByte('\n')
	out.WriteString(strings.Repeat("  ", indent))
}

// main reads JSON from stdin and writes pretty JSON to stdout.
func main() {
	input, err := io.ReadAll(os.Stdin)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	output, err := PrettyJSON(input)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	os.Stdout.Write(output)
}

func TestPrettyJSON(t *testing.T) {
	got, err := PrettyJSON([]byte(`{"name":"Ada","scores":[1,2,true]}`))
	if err != nil {
		t.Fatal(err)
	}

	want := "{\n  \"name\": \"Ada\",\n  \"scores\": [\n    1,\n    2,\n    true\n  ]\n}\n"
	if string(got) != want {
		t.Fatalf("got:\n%s\nwant:\n%s", got, want)
	}
}

func TestPrettyJSONPreservesEscapedQuotes(t *testing.T) {
	got, err := PrettyJSON([]byte(`{"quote":"she said \"hi\""}`))
	if err != nil {
		t.Fatal(err)
	}

	want := "{\n  \"quote\": \"she said \\\"hi\\\"\"\n}\n"
	if string(got) != want {
		t.Fatalf("got:\n%s\nwant:\n%s", got, want)
	}
}