import os
import sys
import uuid
import yaml
import json
import glob
import shutil
import getpass
import pandas as pd
import axiom.common
from itertools import chain

from axiom.settings import config
def get_global_values():
    shell_pid = str(os.getppid())
    global_value_file = os.path.join(config['axiom_path'], '.pid'+shell_pid)
    return json.load(open(global_value_file))

def set_cwd(cwd):
    shell_pid = str(os.getppid())
    global_value_file = os.path.join(config['axiom_path'], '.pid'+shell_pid)
    global_value = {'cwd': cwd}
    json.dump(global_value, open(global_value_file,'w'))

def main(sys_argv):
    shell_pid = str(os.getppid())
    global_value_file = os.path.join(config['axiom_path'], '.pid'+shell_pid)
    if not os.path.exists(global_value_file):
        cwd = os.path.join(config['axiom_path'], 'info')
        global_value = {'cwd': cwd}
        json.dump(global_value, open(global_value_file,'w'))
    if len(sys_argv) == 1:
        lines = open(__file__).read().split('\n')
        lines = [line for line in lines if line.find('if sys_argv[1]') > 0][1:]
        cmds = [line.split("'")[1].split("'")[0] for line in lines]
        comments = [line.split("#")[-1] for line in lines]
        usages = [cmds[i]+': '+comments[i] for i in range(len(cmds))]
        print('\n'.join(usages))
        return

    if sys_argv[1] == 'info': # show basic info of the project
        print('my file name', __file__)
        print('config', config)

    def compile(txt):
        from axiom.settings import config
        folder = str(uuid.uuid4())
        os.makedirs(os.path.join(folder, 'obj'))
        me, code_txt = txt.split('##__AXIOM_INFO_END__\n')
        code_txt = code_txt.split('##__AXIOM_FUNCTION_END__')[0]
        print('me'.center(50,'='))
        print(me)
        print('code'.center(50,'='))
        print(code_txt)
        open(os.path.join(folder, 'obj', 'main.py'),'w').write(code_txt)
        info_file = os.path.join(folder, 'info.yaml')

        me = me[me.find('=')+1:]
        yaml.safe_dump(json.loads(me), open(info_file, 'w'))
        print('info_file', info_file)
        axiom.common.add_obj_common(os.path.join(folder, 'obj'))

    if sys_argv[1] == 'compile':
        txt = open(sys_argv[2]).read()
        blocks = txt.split('##__AXIOM_BLOCK_BEGIN__')[1:]
        blocks = [b.split('##__AXIOM_BLOCK_END__')[0] for b in blocks]
        for b in blocks:
            compile(b)

    if sys_argv[1] == 'what': # show what can we show
        fs = glob.glob(os.path.join(config['axiom_path'], 'info', '*'))
        yamls = [yaml.safe_load(open(f).read()) for f in fs]
        if len(sys_argv) == 2:
            keys = [list(y.keys()) for y in yamls]
            print(list(set(list(chain.from_iterable(keys)))))
        if len(sys_argv) == 3:
            obj_properties = []
            for y in yamls:
                x = y[sys_argv[2]]
                if isinstance(x, list):
                    obj_properties += x
                else:
                    obj_properties.append(x)

            obj_properties = sorted(list(set(obj_properties)))
            print('\n'.join(obj_properties))

    if sys_argv[1] == 'template': # axiom template report_matrix, add a function template at the current folder
        func_name = sys_argv[2]
        os.mkdir(func_name)
        open(os.path.join(func_name, '__main.py'), 'w').write(template_main)
        open(os.path.join(func_name, 'info.yaml'), 'w').write(template_info)
        open(os.path.join(func_name, '__INCLUDE__'), 'w').write('external folder')
        open(os.path.join(func_name, 'config.yaml'), 'w').write('{}')

    def list_axiom(fs):
        ys = [yaml.safe_load(open(f).read()) for f in fs]
        md5s = [os.path.split(f)[-1][:-5] for f in fs]
        columns = ['given_name', 'type', 'description']
        d = {}
        for i in range(len(fs)):
            d[md5s[i]] = [ys[i]['given_name'], ys[i]['type'], ys[i]['description']]
        df = pd.DataFrame.from_dict(d, orient='index', columns=columns)
        print(df)

    if sys_argv[1] == 'ls_old': # show the axiom.obj
        if len(sys_argv) == 3:
            axiom_path = sys_argv[2]
        else:
            axiom_path = config['axiom_path']
        list_axiom(os.path.join(axiom_path, 'info'))

    if sys_argv[1] == 'lsb': # list all the objects at huggingface axiom/info
        list_axiom(axiom.common.config['bypy_cache'])

    if sys_argv[1] == 'add': # add the current folder as the obj
        axiom.common.add_obj_common(sys_argv[2], to_move=False)

    if sys_argv[1] == 'rm': # remove a md5 obj
        md5 = sys_argv[2]
        os.remove(os.path.join(config['axiom_path'], 'obj', md5))
        os.remove(os.path.join(config['axiom_path'], 'info', md5+'.yaml'))
        shutil.rmtree(os.path.join(config['axiom_cache'], md5))

    if sys_argv[1] == 'extract': # exting an object to the corrent folder
        md5 = sys_argv[2]
        f = os.path.join(config['axiom_path'], 'obj', md5)
        cmd = 'cp %s ./&& tar zfx %s&&rm %s' % (f, md5, md5)
        print(cmd)
        os.system(cmd)


    if sys_argv[1] == 'commit': # upload an md5 to bypy
        md5 = sys_argv[2]
        axiom.common.upload_to_server(md5)

    def run_python_block(txt):
        file_name = str(uuid.uuid4())
        open(file_name, 'w').write(txt)
        os.system('python '+file_name)
        os.remove(file_name)

    if sys_argv[1] == 't': # test a file with __AXIOM__BLOCK__BEGIN__ block, with test
        txt = open(sys_argv[2]).read()
        blocks = txt.split('##__AXIOM_BLOCK_BEGIN__')[1:]
        blocks = [b.split('##__AXIOM_BLOCK_END__')[0] for b in blocks]
        for b in blocks:
            run_python_block(b)

    if sys_argv[1] == 'aaa':
        import huggingface_hub as hh
        cache_dir = config['huggingface_cache']
        repo_id = 'reaxiom/info'
        hh.snapshot_download(repo_id, repo_type='dataset', cache_dir=cache_dir, token=True)
        sha = hh.dataset_info(repo_id).sha
        info_folder = os.path.join(config['huggingface_cache'], 'datasets--'+repo_id.replace('/','--'), 'snapshots', sha)
        fs = glob.glob(os.path.join(info_folder, '*.yaml'))
        print('info_folder'.center(50, '-'))
        print(info_folder)
        print('\n'.join(fs))

    if sys_argv[1] == 'pwd':
        cwd = get_global_values()['cwd']
        print(cwd)

    if sys_argv[1] == 'ls':
        cwd = get_global_values()['cwd']
        if len(sys_argv) == 2:
            folders = os.listdir(cwd)
            print('\n'.join(folders))
        else:
            folder = os.path.join(config['axiom_path'], 'info', sys_argv[2])
            print('folder', folder)
            fs = axiom.common.get_folder_fs(folder)
            list_axiom(fs)

    if sys_argv[1] == 'cd':
        if len(sys_argv) == 2:
            target_folder = os.path.join(config['axiom_path'], 'info')
            set_cwd(target_folder)
        else:
            sub = sys_argv[2]
            d = get_global_values()
            cwd = d['cwd']
            if sub == '..':
                target_folder = os.path.split(cwd)[0]
            else:
                target_folder = os.path.join(cwd, sub)
            if os.path.isdir(target_folder):
                set_cwd(target_folder)
            else:
                print('not found:', target_folder)
    if sys_argv[1] == 'set.huggingface': # settup the huggingface token and username
        token = getpass.getpass('huggingface token:')
        user_name = input('huggingface private db user:')
        obj_db = input('private huggingface obj db: %s/' % user_name)
        info_db = input('private huggingface info db: %s/' % user_name)
        public_user_name = input('huggingface public db user:')
        public_obj_db = input('public huggingface obj db: %s/' % public_user_name)
        public_info_db = input('public huggingface info db: %s/' % public_user_name)
        


        huggingface_path = os.path.join(config['home'], '.huggingface')
        if not os.path.isdir(huggingface_path):
            os.mkdir(huggingface_path)
        open(os.path.join(huggingface_path, 'token'), 'w').write(token)

        d = {}
        d['user_name'] = user_name
        d['obj_db'] = obj_db
        d['info_db'] = info_db
        d['public_user_name'] = public_user_name
        d['public_obj_db'] = public_obj_db
        d['public_info_db'] = public_info_db
        json.dump(d, open(os.path.join(huggingface_path, 'setting.json'), 'w'))


if __name__ == '__main__':
    main(sys.argv)
