#!/usr/bin/env python3
from __future__ import annotations

import argparse
import io
import json
import sys
import unittest
from pathlib import Path
from typing import Any, TextIO


def pretty_print_json(source: TextIO, indent: int = 2, sort_keys: bool = False) -> str:
    """Read JSON from a text stream and return formatted JSON."""
    data: Any = json.load(source)
    return json.dumps(data, indent=indent, sort_keys=sort_keys) + "\n"


def open_input(path: str | None) -> TextIO:
    # Keep stdin as the default so the script works well in pipelines.
    if path is None or path == "-":
        return sys.stdin
    return Path(path).open("r", encoding="utf-8")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Tiny JSON pretty-printer.")
    parser.add_argument("path", nargs="?", help="JSON file to print, or '-' / omitted for stdin")
    parser.add_argument("-i", "--indent", type=int, default=2, help="spaces per indent level")
    parser.add_argument("--sort-keys", action="store_true", help="sort object keys alphabetically")
    parser.add_argument("--test", action="store_true", help="run tests and exit")
    return parser


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)

    if args.test:
        suite = unittest.defaultTestLoader.loadTestsFromTestCase(PrettyPrintJsonTests)
        result = unittest.TextTestRunner(verbosity=2).run(suite)
        return 0 if result.wasSuccessful() else 1

    try:
        with open_input(args.path) as source:
            sys.stdout.write(pretty_print_json(source, indent=args.indent, sort_keys=args.sort_keys))
    except (OSError, json.JSONDecodeError) as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 1

    return 0


class PrettyPrintJsonTests(unittest.TestCase):
    def test_pretty_prints_nested_json(self) -> None:
        source = io.StringIO('{"name":"Ada","items":[1,true,null]}')
        expected = '{\n  "name": "Ada",\n  "items": [\n    1,\n    true,\n    null\n  ]\n}\n'
        self.assertEqual(pretty_print_json(source), expected)

    def test_sorts_keys_when_requested(self) -> None:
        source = io.StringIO('{"b":2,"a":1}')
        expected = '{\n  "a": 1,\n  "b": 2\n}\n'
        self.assertEqual(pretty_print_json(source, sort_keys=True), expected)


if __name__ == "__main__":
    raise SystemExit(main())