#!/usr/bin/env python3
# SPDX-License-Identifier: WTFPL

# /// script
# dependencies = ["prettytable"]
# ///

"""
Usage:

Print a table of file.json:
    json2table < file.json

Same, but put 'id' and 'name' as first columns:
    json2table id name < file.json
"""

import argparse
from collections import OrderedDict
import prettytable
import json
import sys


def read(fd):
    with fd:
        text = fd.read()

    try:
        return json.loads(text, object_pairs_hook=OrderedDict)
    except json.JSONDecodeError as initial:
        try:
            return [json.loads(line, object_pairs_hook=OrderedDict) for line in text.split("\n") if line]
        except json.JSONDecodeError:
            raise initial


parser = argparse.ArgumentParser()
parser.add_argument('-f', '--file', default='-', help='JSON file to read (standard input by default)')
parser.add_argument('-k', '--keep-only', action='store_const', const=True, help='Keep only specified keys')
parser.add_argument('-b', '--box', action='store_true', help='Use Unicode pretty characters instead of plain ASCII table')
parser.add_argument('--markdown', action='store_true', help='Output markdown table')
parser.add_argument('field', nargs='*', help='display field in the leftmost columns')
args = parser.parse_args()

if args.file == '-':
    fd = sys.stdin
else:
    fd = open(args.file, 'r')

data = read(fd)

okdata = False
if isinstance(data, list):
    if all(isinstance(obj, dict) for obj in data):
        okdata = True
    elif all(isinstance(obj, list) for obj in data):
        data = [{col: obj[n] for n, col in enumerate(data[0])} for obj in data[1:]]
        okdata = True
elif isinstance(data, dict) and all(isinstance(obj, list) for obj in data.values()):
    def get(l, n):
        try:
            return l[n]
        except IndexError:
            return None

    nrows = max(len(v) for v in data.values())
    data = [OrderedDict((k, get(data[k], n)) for k in data) for n in range(nrows)]
    okdata = True

if not okdata:
    raise ValueError('input must be a JSON list of dictionaries or a dictionary with lists')
order = args.field or []
keys = dict((k, None) for k in order)

if args.keep_only:
    def keep(d):
        return dict((k, v) for k, v in d.items() if k in keys)
    data = [keep(r) for r in data]
else:
    keys.update((k, None) for d in data for k in d.keys())

keys = list(keys)

table = prettytable.PrettyTable(field_names=keys)
if args.box:
    table.set_style(prettytable.TableStyle.SINGLE_BORDER)
elif args.markdown:
    table.set_style(prettytable.TableStyle.MARKDOWN)

for d in data:
    table.add_row([d.get(k) for k in keys])
table.align = "l"
print(table)
