import os

from huggingface_hub import HfApi

API = HfApi(token=os.environ.get("HF_TOKEN"))
REPO_ID = "Limour/llama-python-streamingllm"


def restart_space():
    API.restart_space(repo_id=REPO_ID, token=os.environ.get("HF_TOKEN"))


def init(cfg):
    # ========== 共同 ==========
    model = cfg['model']
    s_info = cfg['s_info']

    def btn_reset(_cache_path):
        try:
            with cfg['session_lock']:
                _tmp = model.load_session(_cache_path)
                print(f'load cache from {_cache_path} {_tmp}')
                cfg['session_active'] = True
                return model.venv_info
        except Exception as e:
            restart_space()
            raise e

    def btn_stop():
        cfg['btn_stop_status'] = True

    cfg['btn_stop'].click(
        fn=btn_stop
    )

    cfg['btn_reset'].click(
        fn=btn_stop
    ).success(
        fn=btn_reset,
        inputs=cfg['setting_cache_path'],
        outputs=s_info,
        **cfg['btn_concurrency']
    ).success(
        **cfg['btn_finish']
    )

    cfg['btn_debug'].click(
        fn=lambda: model.str_detokenize(model._input_ids),
        outputs=cfg['vo']
    )