# Based on: https://github.com/allenai/s2-folks/tree/main/examples/python
#
#

import wormneuroatlas as wa
from pprint import PrettyPrinter

FORMAT_TEXT = "Text"
FORMAT_DICT = "Dict"
FORMAT_JSON = "JSON"

FORMATS = (FORMAT_TEXT, FORMAT_DICT)

DS_WORMNEUROATLAS = "WormNeuroAtlas"


pp = PrettyPrinter(depth=6)


def query_data_source(text, source, format):
    if source == DS_WORMNEUROATLAS:
        atlas = wa.NeuroAtlas()
        ds, ss = atlas.everything_about(
            text, return_values=True, return_text=True, print_text=False
        )

        if format == FORMAT_TEXT:
            info = """Extracting info on neuron %s from %s...
""" % (
                text,
                source,
            )
            info += """```%s```""" % ss
        elif format == FORMAT_DICT:
            info = """```
%s```""" % pp.pformat(ds)
        elif format == FORMAT_JSON:
            import json

            jj = json.dumps(ds)
            info = """```
%s```""" % pp.pformat(jj)

    return info