from typing import Dict
import copy
import json
import shlex


def str_wrap_fbnet_arch_def(d: Dict, inplace=False) -> Dict:
    """Replaces MODEL.FBNET_V2.ARCH_DEF with wrapped json string

    Searches the input dict to see if it contains MODEL.FBNET_V2.ARCH_DEF
    and replaces the value with a wrapped json string

    The json string is created because FBNet builder runs json.loads on the
    archdef. The json string needs to be wrapped in another string because
    CfgNode runs literal_eval in order to check whether it should continue
    to create CfgNodes if the value is a dict.

        arch_def = {...}                               # {...}
        arch_def = json.dumps(arch_def)                # '{...}'
        arch_def = strwrap(arch_def)                   # '"{...}"'

        CfgNode(arch_def) => literal_eval(arch_def)    # '{...}'
        FBNetBuilder(arch_def) => json.loads(arch_def) # {...}

    Example:
        config = {"MODEL": {"FBNET_V2": {"ARCH_DEF": [1, 1, 1]}}}
        str_wrap_fbnet_arch_def(config)
          => {"MODEL": {"FBNET_V2": {"ARCH_DEF": '''"[1, 1, 1]"'''}}}
    """
    if not inplace:
        d = copy.deepcopy(d)

    try:
        archdef = d["MODEL"]["FBNET_V2"]["ARCH_DEF"]
        # MODEL.FBNET_V2.ARCH_DEF needs to be json str
        archdef = json.dumps(archdef)
        # CfgNode runs literal_eval when merging so wrap around str
        archdef = shlex.quote(archdef)
        d["MODEL"]["FBNET_V2"]["ARCH_DEF"] = archdef
    except KeyError:
        pass

    return d