#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @Author: Niccolò Bonacchi
# @Date: Tuesday, January 22nd 2019, 12:32:14 pm
from pathlib import Path
import logging

logger_ = logging.getLogger(__name__)

FLAG_FILE_NAMES = [
    'transfer_me.flag', 'extract_me.flag', 'register_me.flag', 'flatiron.flag',
    'extract_me.error', 'register_me.error', 'create_me.flag', 'compress_video.flag',
    'compress_audio.flag', 'extract_ephys.flag',
]


def read_flag_file(fname):
    """
    Flag files are ``*.flag`` files within a session folder used to schedule some jobs
    If they are empty, should return True

    :param fname: full file path of the flag file
    :type fname: str or pahlib.Path
    :return: None
    """
    # the flag file may contains specific file names for a targeted extraction
    with open(fname) as fid:
        save = list(set(list(filter(None, fid.read().splitlines()))))
    # if empty, extract everything by default
    if len(save) == 0:
        save = True
    return save


def excise_flag_file(fname, removed_files=None):
    """
    Remove one or several specific files if they figure within the file
    If no file is left, deletes the flag.

    :param fname: full file path of the flag file
    :type fname: str or pahlib.Path
    :return: None
    """
    if not removed_files:
        return
    file_names = read_flag_file(fname)
    # if the file is empty, can't remove a specific file and return
    if len(file_names) == 0:
        return
    if isinstance(removed_files, str):
        removed_files = [removed_files]
    new_file_names = list(set(file_names).difference(set(removed_files)))
    # if the resulting file has no files in it, delete
    if len(new_file_names) == 0:
        Path(fname).unlink()
    else:
        write_flag_file(fname, file_list=new_file_names, clobber=True)


def write_flag_file(fname, file_list: list = None, clobber=False):
    """
    Flag files are ``*.flag`` files within a session folder used to schedule some jobs
    Each line references to a file to extract or register

    :param fname: full file path of the flag file
    :type fname: str or pathlib.Path
    :param file_list: None or list of relative paths to write in the file
    :type file_list: list
    :param clobber: (False) overwrites the flag file if any
    :type clobber: bool, optional
    :return: None
    """
    exists = Path(fname).exists()
    if exists:
        has_files = Path(fname).stat().st_size != 0
    else:
        has_files = False
    if isinstance(file_list, str) and file_list:
        file_list = [file_list]
    if isinstance(file_list, bool):
        file_list = None
    if clobber:
        mode = 'w+'
    elif exists and has_files and file_list:
        mode = 'w+'
        file_list_flag = read_flag_file(fname)
        # if the file is empty, can't remove a specific file and return
        if len(file_list_flag) == 0:
            file_list = [''] + file_list
        else:
            file_list = list(set(file_list + file_list_flag))
    else:
        mode = 'w+'
        if exists and not has_files:
            file_list = []
    with open(fname, mode) as fid:
        if file_list:
            fid.write('\n'.join(file_list))


def create_register_flags(root_data_folder, force=False, file_list=None):
    ses_path = Path(root_data_folder).glob('**/raw_behavior_data')
    for p in ses_path:
        flag_file = Path(p).parent.joinpath('register_me.flag')
        if p.parent.joinpath('flatiron.flag').is_file() and not force:
            continue
        if p.parent.joinpath('extract_me.error').is_file() and not force:
            continue
        if p.parent.joinpath('register_me.error').is_file() and not force:
            continue
        if force and p.parent.joinpath('flatiron.flag').is_file():
            p.parent.joinpath('flatiron.flag').unlink()
        write_flag_file(flag_file, file_list)
        logger_.info('created flag: ' + str(flag_file))


def create_extract_flags(root_data_folder, force=False, file_list=None):
    # first part is to create extraction flags
    ses_path = Path(root_data_folder).glob('**/raw_behavior_data')
    for p in ses_path:
        flag_file = Path(p).parent.joinpath('extract_me.flag')
        if p.parent.joinpath('flatiron.flag').is_file() and not force:
            continue
        if p.parent.joinpath('extract_me.error').is_file() and not force:
            continue
        if p.parent.joinpath('register_me.error').is_file() and not force:
            continue
        if force and p.parent.joinpath('flatiron.flag').is_file():
            p.parent.joinpath('flatiron.flag').unlink()
        if force and p.parent.joinpath('register_me.flag').is_file():
            p.parent.joinpath('register_me.flag').unlink()
        write_flag_file(flag_file, file_list)
        logger_.info('created flag: ' + str(flag_file))


def create_transfer_flags(root_data_folder, force=False, file_list=None):
    create_other_flags(root_data_folder, 'transfer_me.flag', force=False, file_list=None)


def create_create_flags(root_data_folder, force=False, file_list=None):
    create_other_flags(root_data_folder, 'create_me.flag', force=False, file_list=None)


def create_other_flags(root_data_folder, name, force=False, file_list=None):
    ses_path = Path(root_data_folder).glob('**/raw_behavior_data')
    for p in ses_path:
        flag_file = Path(p).parent.joinpath(name)
        write_flag_file(flag_file)
        logger_.info('created flag: ' + str(flag_file))


def create_compress_video_flags(root_data_folder, flag_name='compress_video.flag', clobber=False):
    #  only create flags for raw_video_data folders:
    video_paths = Path(root_data_folder).glob('**/raw_video_data')
    for video_path in video_paths:
        ses_path = video_path.parent
        flag_file = ses_path.joinpath(flag_name)
        vfiles = video_path.rglob('*.avi')
        for vfile in vfiles:
            logger_.info(str(vfile.relative_to(ses_path)) + ' added to ' + str(flag_file))
            write_flag_file(flag_file, file_list=str(vfile.relative_to(ses_path)), clobber=clobber)


def create_audio_flags(root_data_folder, flag_name):
    # audio flags could be audio_ephys.flag, audio_training.flag
    if flag_name not in ['audio_ephys.flag', 'audio_training.flag']:
        raise ValueError('Flag name should be audio_ephys.flag or audio_training.flag')
    audio_paths = Path(root_data_folder).glob('**/raw_behavior_data')
    for audio_path in audio_paths:
        ses_path = audio_path.parent
        flag_file = ses_path.joinpath(flag_name)
        afiles = audio_path.rglob('*.wav')
        for afile in afiles:
            logger_.info(str(afile.relative_to(ses_path)) + ' added to ' + str(flag_file))
            write_flag_file(flag_file, file_list=str(afile.relative_to(ses_path)))


def create_dlc_flags(root_path, dry=False, clobber=False, force=False):
    # look for all mp4 raw video files
    root_path = Path(root_path)
    for file_mp4 in root_path.rglob('_iblrig_leftCamera.raw*.mp4'):
        ses_path = file_mp4.parents[1]
        file_label = file_mp4.stem.split('.')[0].split('_')[-1]
        # skip flag creation if there is a file named _ibl_*Camera.dlc.npy
        if (ses_path / 'alf' / f'_ibl_{file_label}.dlc.npy').exists() and not force:
            continue
        if not dry:
            write_flag_file(ses_path / 'dlc_training.flag',
                            file_list=[str(file_mp4.relative_to(ses_path))],
                            clobber=clobber)
        logger_.info(str(ses_path / 'dlc_training.flag'))


def create_flags(root_data_folder: str or Path, flags: list,
                 force: bool = False, file_list: list = None) -> None:
    ses_path = Path(root_data_folder).glob('**/raw_behavior_data')
    for p in ses_path:
        if 'create' in flags:
            create_create_flags(root_data_folder, force=force, file_list=file_list)
        elif 'transfer' in flags:
            create_transfer_flags(root_data_folder, force=force, file_list=file_list)
        elif 'extract' in flags:
            create_extract_flags(root_data_folder, force=force, file_list=file_list)
        elif 'register' in flags:
            create_register_flags(root_data_folder, force=force, file_list=file_list)


def delete_flags(root_data_folder):
    for f in Path(root_data_folder).rglob('*.flag'):
        f.unlink()
        logger_.info('deleted flag: ' + str(f))
    for f in Path(root_data_folder).rglob('*.error'):
        f.unlink()
        logger_.info('deleted flag: ' + str(f))
