"""Define constants used throughout the project.
"""

from pathlib import Path

from htools.core import shell
from htools.structures import IndexedDict


# Have to define this at the top so we can reference it in functions. We later
# make it available as part of C.
git_clone_url = 'https://github.com/hdmamin/jabberwocky.git'


def resolve_root():
    """Get the project root directory. This directory should contain the 'data'
    subdir. If working from the main git repo (i.e. what I've been doing for
    most of this project), this is ~/jabberwocky. When installed as a pip
    package, this is ~/.jabberwocky.

    Returns
    -------
    pathlib.Path
    """
    root = Path('~/jabberwocky').expanduser()
    if root.is_dir() and (root/'data').is_dir():
        return root
    return clone_or_pull_prompts('~/.jabberwocky', update=False)


def clone_or_pull_prompts(path='~/.jabberwocky', update=False):
    """Download prompt data by cloning/pulling from github.

    Parameters
    ----------
    path: str
    update: bool
        If True and path is an existing directory, we try to pull changes from
        github. Note that this does NOT affect anything if the dir does not
        exist (in which case we will always clone the repo).
    """
    path = Path(path).expanduser()
    if path.is_dir():
        if update:
            print('Pulling prompts from github because update=True...')
            code, *_ = shell(f'cd {path} && git pull', return_output=False)
        else:
            print('Prompts already present, not pulling changes because '
                  'update=False.')
            code = 0
    else:
        print(f'{path} is not an existing directory. Cloning from github...')
        # Must reference the global git clone var var here.
        code, *_ = shell(f'git clone {git_clone_url} {path}',
                         return_output=False)
    if code:
        raise RuntimeError('Cloning/pulling prompts failed.')
    return path


class C:
    # Formatting (note: later realized bold_end actually also ends colors too.)
    bold_start = '\033[1m'
    underline_start = '\033[4m'
    bold_end = '\033[0m'

    # OpenAI constants
    # Parameter Counts: 2.7B, 6.7B, 13B, 175B
    # Empty strings mean the backend doesn't provide an equivalent engine.
    engines = ['text-ada-001',
               'text-babbage-001',
               'text-curie-001',
               'text-davinci-002']
    backend_engines = {
        'openai': engines,
        'gooseai': [
            'gpt-neo-2-7b',
            'gpt-j-6b',
            'fairseq-13b',
            'gpt-neo-20b'
        ],
        'huggingface': [
            'gpt-neo-2.7B',
            'gpt-j-6B',
            '',
            ''
        ],
        # These backends only provide 1 model so these are just informational.
        # (They give GPTBackend.engine() something to return but they don't
        # actually affect the query results.)
        'hobby': ['gpt-j-6B' for _ in range(2)] + ['', ''],
        'banana': ['gpt-j-6B' for _ in range(2)] + ['', '']
    }

    # Dollars per thousand tokens with openai backend.
    prices = IndexedDict(zip(engines, [.0008, .0012, .006, .06]))

    # Data/Paths
    root = resolve_root()
    data_dir = root/'data'
    prompt_dir = data_dir/'prompts'
    # Make this available as an attribute now. Need to keep global var to use
    # in functions in this module which have to execute before this line.
    git_clone_url = git_clone_url
    # This contains a dict generated by scripts/s01_fetch_sample_responses.py.
    # Contains more variants of responses (see script docstring for details).
    all_mocks_path = root/'data/misc/gooseai_sample_responses.pkl'

    # Use with datetime.strptime() to convert from format like
    # 'Sun Apr 17 13:18:23 2022' back to datetime object. This is the format
    # GPTBackend.query logger uses in the meta.datetime field.
    ctime_fmt = '%a %b %d %H:%M:%S %Y'

    # Notice this still has a {}: we must call remote_prompt_url.format(task)
    # to get the fully resolved URL.
    prompt_base_url = 'https://raw.githubusercontent.com/hdmamin/'\
        'jabberwocky/main/data/prompts/{}.yaml'


