#!python3
from argparse import ArgumentParser, BooleanOptionalAction
import sys
from rdrag import util
from rdrag import constants


# PARSE_ARGS = 'answer_only temperature top_p max_tokens limit'.split()
# def extract_options(argv=sys.argv):
#     args = sys.argv[1:]
#     options = {}
#     remove_indexes = []
#     for i, a in enumerate(args):
#         k = a.split('=')[0]
#         if not k.startswith('--') or k[2:] not in PARSE_ARGS:
#             continue
#         remove_indices.append(i)
#         k = k[2:]
#         if i == len(args) and len(k.split('=') < 2):
#             options[k] = True
#             continue
# options[i] = args[i+1]


ask_parser = ArgumentParser(
    prog='privatera.ask',
    description='Ask the Privera RA to research a question for you.',
    epilog='SEE: gitlab.com/tangibleai/privatera'
)
ask_parser.add_argument('words', nargs='*')           # positional argument
ask_parser.add_argument('-a', '-o', '--answer_only', action='store_true')
ask_parser.add_argument('--json', action=BooleanOptionalAction)
ask_parser.add_argument('-v', '--verbose', action='store_true')

ask_parser.add_argument('-c', '--cache_prompt', action=BooleanOptionalAction)      # option that takes a value
ask_parser.add_argument('-l', '--limit', default='3', type=int)      # option that takes a value
ask_parser.add_argument('-m', '--max_tokens', default='256', type=int)      # option that takes a value
ask_parser.add_argument('-p', '--top_p', default='0.5', type=float)      # option that takes a value
ask_parser.add_argument('-t', '--temperature', default='0.3', type=float)      # option that takes a value
ask_parser.add_argument('-k', '--top_k', default='10', type=int)      # option that takes a value
ask_parser.add_argument('-b', '--base_url', default=constants.LLM_BASE_URL, type=str)      # option that takes a value
ask_parser.add_argument('-n', '--llm_name', default=constants.LLM_NAME, type=str)      # option that takes a value
# limit: int = 3,
# max_tokens: int = 128,
# top_p: float = 0.5,
# temperature: float = 0.3,
# base_url: str = LLM_BASE_URL,
# llm_name: str = LLM_NAME,
# answer_only: bool = True,
# automatic_routing: bool = True

search_parser = ArgumentParser(
    prog='privatera.search',
    description='Search the Privera RA meilisearch database.',
    epilog='SEE: gitlab.com/tangibleai/privatera'
)

llm_parser = ArgumentParser(
    prog='privatera.llm',
    description='Chat with one of the Private RA LLMs.',
    epilog='SEE: gitlab.com/tangibleai/privatera',
)


def search():
    print(util.search(' '.join(sys.argv[1:])))


# def ask():
#     print(util.ask(' '.join(sys.argv[1:])))


def llm():
    print(util.complete_with_llm(' '.join(sys.argv[1:])))


def ask(*args, **kwargs):
    p = ask_parser.parse_args(args=sys.argv[1:] + list(args))
    options = {'answer_only': True}
    options.update(vars(p))
    options.update(kwargs)
    words = options.pop('words')
    json = options.pop('json')
    verbose = json or options.pop('verbose')
    verbose = verbose and not options.pop('answer_only')
    print(util.ask(' '.join(words), **options))
