import hashlib
import os
import sys
import shutil
from functools import wraps, partial

import pytest

if os.path.dirname('src') not in sys.path:
    sys.path.append('src')

os.environ['HARD_ASSERTS'] = "1"

from src.utils import call_subprocess_onetask, makedirs, FakeTokenizer, download_simple, sanitize_filename


def get_inf_port():
    if os.getenv('HOST') is not None:
        inf_port = os.environ['HOST'].split(':')[-1]
    elif os.getenv('GRADIO_SERVER_PORT') is not None:
        inf_port = os.environ['GRADIO_SERVER_PORT']
    else:
        inf_port = str(7860)
    return int(inf_port)


def get_inf_server():
    if os.getenv('HOST') is not None:
        inf_server = os.environ['HOST']
    elif os.getenv('GRADIO_SERVER_PORT') is not None:
        inf_server = "http://localhost:%s" % os.environ['GRADIO_SERVER_PORT']
    else:
        raise ValueError("Expect tests to set HOST or GRADIO_SERVER_PORT")
    return inf_server


def get_mods():
    testtotalmod = int(os.getenv('TESTMODULOTOTAL', '1'))
    testmod = int(os.getenv('TESTMODULO', '0'))
    return testtotalmod, testmod


def do_skip_test(name):
    """
    Control if skip test.  note that skipping all tests does not fail, doing no tests is what fails
    :param name:
    :return:
    """
    testtotalmod, testmod = get_mods()
    return int(get_sha(name), 16) % testtotalmod != testmod


def wrap_test_forked(func):
    """Decorate a function to test, call in subprocess"""

    @wraps(func)
    def f(*args, **kwargs):
        # automatically list or set, so can globally control server ports or host for all tests
        gradio_port = os.environ['GRADIO_SERVER_PORT'] = os.getenv('GRADIO_SERVER_PORT', str(7860))
        gradio_port = int(gradio_port)
        # testtotalmod, testmod = get_mods()
        # gradio_port += testmod
        os.environ['HOST'] = os.getenv('HOST', "http://localhost:%s" % gradio_port)

        pytest_name = get_test_name()
        if do_skip_test(pytest_name):
            # Skipping is based on raw name, so deterministic
            pytest.skip("[%s] TEST SKIPPED due to TESTMODULO" % pytest_name)
        func_new = partial(call_subprocess_onetask, func, args, kwargs)
        return run_test(func_new)

    return f


def run_test(func, *args, **kwargs):
    return func(*args, **kwargs)


def get_sha(value):
    return hashlib.md5(str(value).encode('utf-8')).hexdigest()


def get_test_name():
    tn = os.environ['PYTEST_CURRENT_TEST'].split(':')[-1]
    tn = "_".join(tn.split(' ')[:-1])  # skip (call) at end
    return sanitize_filename(tn)


def make_user_path_test():
    import os
    import shutil
    user_path = makedirs('user_path_test', use_base=True)
    if os.path.isdir(user_path):
        shutil.rmtree(user_path)
    user_path = makedirs('user_path_test', use_base=True)
    db_dir = "db_dir_UserData"
    db_dir = makedirs(db_dir, use_base=True)
    if os.path.isdir(db_dir):
        shutil.rmtree(db_dir)
    db_dir = makedirs(db_dir, use_base=True)
    shutil.copy('data/pexels-evg-kowalievska-1170986_small.jpg', user_path)
    shutil.copy('README.md', user_path)
    shutil.copy('docs/FAQ.md', user_path)
    return user_path


def get_llama(llama_type=3):
    from huggingface_hub import hf_hub_download

    # FIXME: Pass into main()
    if llama_type == 1:
        file = 'ggml-model-q4_0_7b.bin'
        dest = 'models/7B/'
        prompt_type = 'plain'
    elif llama_type == 2:
        file = 'WizardLM-7B-uncensored.ggmlv3.q8_0.bin'
        dest = './'
        prompt_type = 'wizard2'
    elif llama_type == 3:
        file = download_simple('https://huggingface.co/TheBloke/Llama-2-7b-Chat-GGUF/resolve/main/llama-2-7b-chat.Q6_K.gguf?download=true')
        dest = './'
        prompt_type = 'llama2'
    else:
        raise ValueError("unknown llama_type=%s" % llama_type)

    makedirs(dest, exist_ok=True)
    full_path = os.path.join(dest, file)

    if not os.path.isfile(full_path):
        # True for case when locally already logged in with correct token, so don't have to set key
        token = os.getenv('HUGGING_FACE_HUB_TOKEN', True)
        out_path = hf_hub_download('h2oai/ggml', file, token=token, repo_type='model')
        # out_path will look like '/home/jon/.cache/huggingface/hub/models--h2oai--ggml/snapshots/57e79c71bb0cee07e3e3ffdea507105cd669fa96/ggml-model-q4_0_7b.bin'
        shutil.copy(out_path, dest)
    return prompt_type, full_path


def kill_weaviate(db_type):
    """
    weaviate launches detatched server, which accumulates entries in db, but we want to start freshly
    """
    if db_type == 'weaviate':
        os.system('pkill --signal 9 -f weaviate-embedded/weaviate')


def count_tokens_llm(prompt, base_model='h2oai/h2ogpt-oig-oasst1-512-6_9b', tokenizer=None):
    import time
    if tokenizer is None:
        assert base_model is not None
        from transformers import AutoTokenizer
        tokenizer = AutoTokenizer.from_pretrained(base_model)
    t0 = time.time()
    a = len(tokenizer(prompt)['input_ids'])
    print('llm: ', a, time.time() - t0)
    return dict(llm=a)


def count_tokens(prompt, base_model='h2oai/h2ogpt-oig-oasst1-512-6_9b'):
    tokenizer = FakeTokenizer()
    num_tokens = tokenizer.num_tokens_from_string(prompt)
    print(num_tokens)

    from transformers import AutoTokenizer

    t = AutoTokenizer.from_pretrained("distilgpt2")
    llm_tokenizer = AutoTokenizer.from_pretrained(base_model)

    from InstructorEmbedding import INSTRUCTOR
    emb = INSTRUCTOR('hkunlp/instructor-large')

    import nltk


    def nltkTokenize(text):
        words = nltk.word_tokenize(text)
        return words


    import re

    WORD = re.compile(r'\w+')


    def regTokenize(text):
        words = WORD.findall(text)
        return words

    counts = {}
    import time
    t0 = time.time()
    a = len(regTokenize(prompt))
    print('reg: ', a, time.time() - t0)
    counts.update(dict(reg=a))

    t0 = time.time()
    a = len(nltkTokenize(prompt))
    print('nltk: ', a, time.time() - t0)
    counts.update(dict(nltk=a))

    t0 = time.time()
    a = len(t(prompt)['input_ids'])
    print('tiktoken: ', a, time.time() - t0)
    counts.update(dict(tiktoken=a))

    t0 = time.time()
    a = len(llm_tokenizer(prompt)['input_ids'])
    print('llm: ', a, time.time() - t0)
    counts.update(dict(llm=a))

    t0 = time.time()
    a = emb.tokenize([prompt])['input_ids'].shape[1]
    print('instructor-large: ', a, time.time() - t0)
    counts.update(dict(instructor=a))

    return counts