from google_sheet import conf_repo, secret_repo
import json

from kaggle_service import KernelRerunService, NbJob
from logger import sheet_logger, get_now
from run_stt import run_stt_service


def get_secret_dict():
    # load all account and secret into a dict
    secret_dict = {}
    try:
        for i in range(2, 50):  # note: read row 2 - 50
            rs = secret_repo.read(i)
            if not rs:
                break
            assert 'username' in rs
            assert 'secret' in rs

            username = rs['username'].rstrip()
            secret = rs['secret'].rstrip()

            secret_dict[username] = secret
    except Exception as e:
        sheet_logger.log(log="Get secret table failed!!" + str(e))
    return secret_dict


def get_config_list(secret_dict):
    configs = []
    ids = []
    try:
        for i in range(2, 50):     # note: read row 2 - 50
            rs = conf_repo.read(i)
            if not rs:
                break
            if not rs['config']:
                reset_keys = []
                for key in rs:
                    if rs[key]:
                        reset_keys.append(key)
                if len(reset_keys) > 0:
                    conf_repo.update(row_index=i, data={k: '' for k in reset_keys})
                break
            print(i, rs)
            try:
                # validation
                # assert 'config' in rs, 'require config column'
                assert 'usernames' in rs, 'require usernames column'
                assert 'enable' in rs, 'require enable column'

                assert rs['enable'].rstrip() in ['enable', 'disable'], rs['enable']
                assert 'slug' in rs['config'], 'require config.slug'

                cfg = json.loads(rs['config'])
                usernames = set(rs['usernames']
                             .rstrip().replace('\n', ' ').replace(',', ' ').replace(';', ' ').replace('|', ' ')
                             .split(' '))
                usernames = [u for u in usernames if u]
                is_enabled = rs['enable'].rstrip() == 'enable'
                if not is_enabled:
                    sheet_logger.log(log="Disabled, skip!", nb=cfg['slug'])
                    print("skip ", cfg['slug'])
                    continue

                assert len(usernames) > 0, 'len usernames == 0'

                # process config
                accounts = {u: secret_dict[u] for u in secret_dict if u in usernames}
                assert not set(usernames).difference(set(accounts.keys())), set(usernames).difference(set(accounts.keys()))
                cfg = {**cfg, "accounts": accounts}

                # save
                configs.append(cfg)
                ids.append(i)
                print(cfg)
            except AssertionError:
                import traceback
                sheet_logger.update_job_status(i, validate_status=str(traceback.format_exc()))
    except Exception:
        import traceback
        sheet_logger.log(log="Get config failed!!" + str(traceback.format_exc()))
    return configs, ids


def create_service(configs, ids):
    service = KernelRerunService()
    for config, idx in zip(configs, ids):
        try:
            service.add_job(NbJob.from_dict(config, id=idx))
        except Exception as e:
            sheet_logger.update_job_status(idx, validate_status=str(e))
    return service


def main():
    sheet_logger.log("========= start ==========")

    obj = run_stt_service.get_obj()
    if run_stt_service.is_need_to_run(obj):
        # start
        run_stt_service.set_is_running({})

        secret_dict = get_secret_dict()

        # load config from gg sheet file
        configs, ids = get_config_list(secret_dict)

        # add config to service
        service = create_service(configs, ids)

        # run service
        try:
            service.validate_all()
            # service.status_all()
            service.run_all()
        except Exception as e:
            sheet_logger.log(log=str(e))

        # done
        run_stt_service.set_run_done()

    else:
        sheet_logger.log(f"Current time is {get_now()} is not after {obj.get('last_run')} {obj.get('auto_run_after_last_run')} or not require new run")

    sheet_logger.log("========= end ==========")


if __name__ == "__main__":
    main()