#
# bmspy
#
import atexit
import argparse
import pprint
import time


def parse_args():
    parser = argparse.ArgumentParser(
                description='Query JBD BMS and report status',
                add_help=True,
    )
    parser.add_argument('--socket', '-s', dest='socket', action='store',
                default='/run/bmspy/bms', help='Socket to communicate with daemon')
    parser.add_argument('--ups', dest='ups', action='store', default=None,
                help='Only show data for this UPS name (default: all)')
    parser.add_argument('--json', '-j', dest='report_json', action='store_true',
                default=False, help='Report data as JSON')
    parser.add_argument('--prometheus', '-p', dest='report_prometheus', action='store_true',
                default=False, help='Daemonize and report data to Prometheus')
    parser.add_argument('--file', '-f', dest='report_textfile', type=str, action='store',
                default=False, help='Report data to Prometheus using textfile <file>')
    parser.add_argument('--influxdb', '-i', dest='report_influxdb', action='store_true',
                default=False, help='Daemonize and report data to InfluxDB using INFLUXDB_V2_URL, INFLUXDB_V2_ORG and INFLUXDB_V2_TOKEN environment variables')
    parser.add_argument('--bucket', '-b', dest='influx_bucket', type=str, action='store',
                default="ups", help='Set the bucket name when sending data to influxdb (defaults to "ups")')
    parser.add_argument('--url', '-u', dest='influx_url', type=str, action='store',
                default=False, help='Set the URL when sending data to influxdb (overrides INFLUXDB environment variables)')
    parser.add_argument('--org', '-o', dest='influx_org', type=str, action='store',
                default=False, help='Set the influx organization when sending data to influxdb (overrides INFLUXDB environment variables)')
    parser.add_argument('--token', '-t', dest='influx_token', type=str, action='store',
                default=False, help='Set the influx token when sending data to influxdb (overrides INFLUXDB environment variables)')
    parser.add_argument('--print', dest='report_print', action='store_true',
                default=True, help='Report data as text')
    parser.add_argument('--verbose', '-v', action='count',
                default=0, help='Print more verbose information (can be specified multiple times)')
    args = parser.parse_args()
    return args


def main():
    try:
        args = parse_args()

        debug = args.verbose

        if args.report_influxdb:
            num_args = 0
            for arg in [args.influx_url, args.influx_org, args.influx_token]:
                if arg is not False:
                    num_args += 1
            if num_args != 0 and num_args != 3:
                raise argparse.ArgumentTypeError('Missing value for --url, --org or --token')

        if args.report_prometheus:
            from bmspy import prometheus
            prometheus.prometheus_export(daemonize=True, debug=debug)

        if args.report_influxdb:
            from bmspy import influxdb as bms_influx

            bms_influx.influxdb_export(
                bucket=args.influx_bucket,
                url=args.influx_url,
                org=args.influx_org,
                token=args.influx_token,
                ups=args.ups,
                debug=debug,
                daemonize=True,
            )

        elif args.report_textfile:
            from bmspy import prometheus
            prometheus.prometheus_export(daemonize=False, filename=args.report_textfile, debug=debug)

        else:
            from bmspy import client
            client.handle_registration(args.socket, 'bmspy', debug)
            atexit.register(client.handle_registration, args.socket, 'bmspy', debug)

            # {ups_name: JBDUPS}
            data = client.read_data(args.socket, 'bmspy', ups=args.ups, debug=debug)

            if args.report_json:
                import json
                print(json.dumps({name: dict(ups.items()) for name, ups in data.items()}, default=str))

            elif args.report_print:
                pp = pprint.PrettyPrinter(indent=4)
                for ups_name, ups_data in data.items():
                    print("=== {} ===".format(ups_name))
                    pp.pprint(ups_data)

    except KeyboardInterrupt as e:
        print(e)


if __name__ == '__main__':
    main()