#! /usr/bin/env python
"""Clean files generated by unittests"""
import re
import os
import argparse
from typing import List


class FileSet:
    def __init__(self, prefix: List[str] = None,
                 suffix: List[str] = None,
                 head: List[str] = None,
                 tail: List[str] = None,
                 body: List[str] = None,
                 word: List[str] = None) -> None:
        self.patterns = []
        if prefix is not None:  # prefix.*
            self.patterns.extend([re.compile(fr"^{_}\.\w+$") for _ in prefix])
        if suffix is not None:  # *.suffix
            self.patterns.extend([re.compile(fr"^\w+\.{_}$") for _ in suffix])
        if head is not None:  # head*
            self.patterns.extend([re.compile(fr"^{_}.*$") for _ in head])
        if tail is not None:  # *tail
            self.patterns.extend([re.compile(fr"^.*{_}$") for _ in tail])
        if body is not None:  # *body*
            self.patterns.extend([re.compile(fr"^.*{_}.*$") for _ in body])
        if word is not None:  # exact match word
            self.patterns.extend([re.compile(fr"^{_}$") for _ in word])

    def __add__(self, other):
        result = FileSet()
        result.patterns = self.patterns + other.patterns
        return result

    def include(self, file_name: str) -> bool:
        included = False
        for p in self.patterns:
            if re.search(p, file_name) is not None:
                included = True
                break
        return included


def main():
    files_to_remove = []

    # Wannier90 dat file should be reserved
    cmd = "find ./examples ./tests -name *.dat"
    all_files = os.popen(cmd).readlines()
    dat_set = FileSet(body=["hr", "wsvec"])
    for f in all_files:
        if not dat_set.include(f):
            files_to_remove.append(f)

    # Other files can be removed directly
    suffixes = ["pkl", "png", "gif", "cube", "npy", "h5"]
    for s in suffixes:
        cmd = f"find ./examples ./tests -name *.{s}"
        all_files = os.popen(cmd).readlines()
        for f in all_files:
            files_to_remove.append(f)

    # Parse command-line parameters
    parser = argparse.ArgumentParser()
    parser.add_argument("-d", "--delete", action="store_true",
                        default=False)
    args = parser.parse_args()
    if args.delete:
        for f in files_to_remove:
            os.system(f"rm -rf {f.lstrip('\n').rstrip('\n')}")
    else:
        print("\nFiles to remove, but not actually removed:\n")
        for f in files_to_remove:
            print(f)


if __name__ == "__main__":
    main()
