#!/usr/bin/env python
# (c)2019-2024, karneliuk.com


# Modules
import json
import logging
import os
import sys


# Own modules
from pygnmi.arg_parser import parse_args
from pygnmi.client import gNMIclient
from pygnmi.artefacts.messages import msg


# Variables
path_log = "log/execution.log"


# Body
def main():
    # Setting logger
    if not os.path.exists(path_log.split("/")[0]):
        os.mkdir(path_log.split("/")[0])

    logging.basicConfig(
        filename=path_log,
        level=logging.INFO,
        format="%(asctime)s.%(msecs)03d+01:00,%(levelname)s,%(message)s",
        datefmt="%Y-%m-%dT%H:%M:%S",
    )
    logging.info("Starting application...")

    # Collecting inputs
    args = parse_args(msg)

    # gNMI operation
    with gNMIclient(
        target=args.target,
        username=args.username,
        password=args.password,
        token=args.token,
        path_cert=args.path_cert,
        path_key=args.path_key,
        path_root=args.path_root,
        override=args.override,
        insecure=args.insecure,
        debug=args.debug,
        show_diff=args.compare,
        skip_verify=args.skip_verify,
        gnmi_timeout=args.gnmi_timeout,
        no_qos_marking=args.no_qos_marking,
    ) as GC:
        result = None

        # Collecting supported capabilities (needed to figure out encoding for telemetry)
        GC.capabilities()

        if args.operation == "capabilities":
            print(f"Doing {args.operation} request to {args.target}...")
            result = GC.capabilities()

        elif args.operation == "get":
            print(f"Doing {args.operation} request to {args.target}...")
            result = GC.get(
                prefix=args.gnmi_prefix,
                path=args.gnmi_path,
                datatype=args.datastore,
                encoding=args.encoding,
                target=args.gnmi_path_target,
            )

        elif args.operation.startswith("set"):
            print(f"Doing {args.operation} request to {args.target}...")
            mode = args.operation.split("-")[1]
            kwargs = {}
            if mode == "delete":
                # For a delete request, give kwarg delete=[path]
                kwargs[mode] = args.gnmi_path
            else:
                # For update (or replace) request, give kwarg update=[(path, data)]
                with open(args.file, "r") as f:
                    data = f.read()
                jdata = json.loads(data)
                kwargs[mode] = [(args.gnmi_path[0], jdata)]
            result = GC.set(encoding=args.encoding, target=args.gnmi_path_target, **kwargs)

        elif args.operation.startswith("subscribe"):
            mode = args.operation.split("-")[1]
            subscription_list = [{"path": xpath, "mode": "target_defined"} for xpath in args.gnmi_path]
            subscribe = {
                "subscription": subscription_list,
                "use_aliases": False,
                "mode": mode,
                "encoding": args.encoding,
            }

            # Set up extensions
            if args.ext_history_snapshot_time:
                try:
                    time_to_int = int(args.ext_history_snapshot_time)

                except:
                    time_to_int = args.ext_history_snapshot_time

                EXT = {"history": {"snapshot_time": time_to_int}}

            elif args.ext_history_range_start and args.ext_history_range_end:
                try:
                    time_to_int_1 = int(args.ext_history_range_start)
                    time_to_int_2 = int(args.ext_history_range_end)

                except:
                    time_to_int_1 = args.ext_history_range_start
                    time_to_int_2 = args.ext_history_range_start

                EXT = {"history": {"range": {"start": time_to_int_1, "end": time_to_int_2}}}

            else:
                EXT = None

            # Telemetry
            result = GC.subscribe2(subscribe=subscribe, target=args.gnmi_path_target, extension=EXT)

            if mode == "stream":
                try:
                    for ent in result:
                        _print_result(ent)
                except KeyboardInterrupt:
                    sys.exit("Telemtry collection is temrinated.")

            elif mode == "once":
                for ent in result:
                    _print_result(ent)
                    if "sync_response" in ent:
                        break

            elif mode == "poll":
                while True:
                    try:
                        input("Press enter to poll, ctrl+c to quit")
                        ent = result.get_update(timeout=5)
                        _print_result(ent)
                    except KeyboardInterrupt:
                        sys.exit("Telemtry collection is temrinated.")

        if result and not args.debug and not args.operation.startswith("subscribe"):
            _print_result(result)


def _print_result(value: dict) -> None:
    """This function prints the result in JSON format"""
    print(json.dumps(value, indent=4))


if __name__ == "__main__":
    main()
