#!/usr/bin/env python

from notebook.serverextensions import toggle_serverextension_python
from notebook.nbextensions import enable_nbextension, disable_nbextension

import argparse
parser = argparse.ArgumentParser()
parser.add_argument("toggle")
parser.parse_args()

args = parser.parse_args()

if args.toggle == "enable":
    toggle_serverextension_python("comet_ml", enabled=True)
    enable_nbextension("notebook", "comet_ml/main")
elif args.toggle == "disable":
    toggle_serverextension_python("comet_ml", enabled=False)
    disable_nbextension("notebook", "comet_ml/main")
else:
    print("You should either pass `enable` or `disable` to `jupyter comet_ml`")
