from typing import Union
from pathlib import Path
from zipfile import ZipFile

import numpy as np
import fasteners
import anymarkup


exp_dir = Path('./exp_utils_logdir').expanduser()
log_dir = Path('./exp_utils_logdir').expanduser()


def git_backup(path):
    try:
        from git import Repo
        from git.exc import InvalidGitRepositoryError
    except ImportError as e:
        print(f"Can't import `git`: {e}")
        return

    try:
        repo = Repo('.')
        pkg = ZipFile(path / 'source.zip', 'w')

        for file_name in repo.git.ls_files().split():
            pkg.write(file_name)

    except InvalidGitRepositoryError as e:
        print(f"Can't use git to backup files: {e}")
    except FileNotFoundError as e:
        print(f"Can't find file {e}. Did you delete a file and forget to `git add .`")


def _read_exp_status(_exp_dir):
    status_path = _exp_dir / '.status.json5'
    if status_path.exists():
        return anymarkup.parse_file(status_path)
    else:
        run_id = 0
        while (exp_dir / str(run_id)).exists():
            run_id += 1
        return {'id': run_id}


def resolve(file_name: str):
    if '$LOGDIR' in file_name:
        file_name = file_name.replace('$LOGDIR', str(log_dir))
    return Path(file_name).expanduser()


def save(file_name: str, array):
    resolved = resolve(file_name)
    if resolved:
        np.save(str(resolved), array)


def load(file_name: str):
    return np.load(resolve(file_name), allow_pickle=True)


def init(_exp_dir: Union[str, Path]):
    global log_dir, exp_dir

    exp_dir = Path(_exp_dir)
    with fasteners.InterProcessLock(exp_dir / '.lock'):
        _exp = _read_exp_status(exp_dir)
        run_id = str(_exp.get('id', 0))
        _exp['id'] = _exp.get('id', 0) + 1

        log_dir = Path(exp_dir).expanduser() / run_id
        log_dir.mkdir(parents=True, exist_ok=True)
        git_backup(log_dir)

        anymarkup.serialize_file(_exp, exp_dir / '.status.json5')

        import exp_utils as PQ
        PQ.log_dir = log_dir
