#! /usr/bin/env python
# _*_ coding: utf-8 _*_
"""
author:     Chen XuHui
date:       2024/7/29
file:       path.py
version:
doc:        路径相关.
"""
import os
from typing import List

CORE_DIR = os.path.abspath(os.path.join(__file__, '..', '..'))  # 本核心kit的路径
WORK_DIR = os.path.abspath(os.path.join(__file__, '..', '..', '..'))  # 代码工程路径
if os.environ.get('USERPROFILE'):
    DESKTOP = os.path.join(os.environ['USERPROFILE'], 'Desktop')  # 桌面路径
else:
    DESKTOP = r'/home'


def get_file_paths(
    root_path, fileformat: str = '', include='', uninclude='',
    relative: bool = False, depth: int = None, verbose: bool = False
) -> List[str]:
    """获取目标目录下所有特定格式的文件路径列表
    :param root_path: str, 根目录
    :param fileformat: str, 文件格式扩展名, 例如 'json'
    :param include: str, 文件名应包含该字符串
    :param uninclude: str, 文件路径不包含该字符串, 优先级高于 include 参数
    :param relative: bool, 返回结果是否采用以 root_path 为参考的相对路径, False时为绝对路径
    :param depth: int, 查找深度, None代表无限深度
    :param verbose: bool, 是否打印查找结果
    :return: file_paths: list[str], 文件路径的列表
    """
    root_path = os.path.normpath(os.path.abspath(root_path))
    if not os.path.exists(root_path):
        raise OSError(f'目标路径{root_path}不存在')
    if not os.path.isdir(root_path):
        raise OSError(f'请输入一个文件夹路径, 当前输入的是文件路径 {root_path}')
    if isinstance(depth, int) and depth <= 0:
        return []  # 已搜索到目标深度, 直接退出
    file_paths = []

    files = os.listdir(root_path)
    new_depth = depth - 1 if isinstance(depth, int) else None
    for obj in files:
        new_path = os.path.join(root_path, obj)
        if uninclude and uninclude in new_path:
            continue  # 包含排除字符, 跳过

        if os.path.isdir(new_path):
            file_paths.extend(
                get_file_paths(new_path, fileformat=fileformat, include=include, uninclude=uninclude, depth=new_depth)
            )
        elif os.path.isfile(new_path):
            if include and include not in obj:
                continue  # 文件名不包含挑选字符, 跳过
            if fileformat and not obj.endswith(fileformat):
                continue  # 文件格式不符合要求
            file_paths.append(new_path)

    if relative:
        file_paths = [os.path.relpath(path_i, start=root_path) for path_i in file_paths]
    if verbose:
        for i in file_paths:
            try:
                print(i)
            except UnicodeDecodeError:
                pass
    return file_paths


def get_all_imports(path=WORK_DIR, encoding='utf-8', extra_ignore_dirs=None):
    import ast
    import logging
    import traceback
    _bulitin_module = {'os', 'sys', 'math', 'random', 'time', 'itertools'}
    _standard_module = os.listdir(os.path.dirname(os.__file__))
    _standard_module = {i.replace('.py', '') for i in _standard_module}
    myself_module = [module for module in os.listdir(path) if os.path.isdir(module)]

    imports = set()
    raw_imports = set()
    candidates = []
    ignore_errors = False
    ignore_dirs = [".hg", ".svn", ".git", ".tox", "__pycache__", "env", "venv"]

    if extra_ignore_dirs:
        ignore_dirs_parsed = []
        for e in extra_ignore_dirs:
            ignore_dirs_parsed.append(os.path.basename(os.path.realpath(e)))
        ignore_dirs.extend(ignore_dirs_parsed)

    for root, dirs, files in os.walk(path):
        dirs[:] = [d for d in dirs if d not in ignore_dirs]

        candidates.append(os.path.basename(root))
        files = [fn for fn in files if os.path.splitext(fn)[1] == ".py"]

        candidates += [os.path.splitext(fn)[0] for fn in files]
        for file_name in files:
            try:
                with open(os.path.join(root, file_name), "r", encoding=encoding) as f:
                    contents = f.read()
            except UnicodeDecodeError:
                with open(os.path.join(root, file_name), "r", encoding='gbk') as f:
                    contents = f.read()
            try:
                tree = ast.parse(contents)
                for node in ast.walk(tree):
                    if isinstance(node, ast.Import):
                        for subnode in node.names:
                            raw_imports.add(subnode.name)
                    elif isinstance(node, ast.ImportFrom):
                        raw_imports.add(node.module)
            except Exception as exc:
                if ignore_errors:
                    traceback.print_exc()
                    logging.warning("Failed on file: %s" % os.path.join(root, file_name))
                    continue
                else:
                    logging.error("Failed on file: %s" % os.path.join(root, file_name))
                    raise exc

    # Clean up imports
    for name in [n for n in raw_imports if n]:
        # Sanity check: Name could have been None if the import statement was as from . import X
        # Cleanup: We only want to first part of the import.
        # Ex: from django.conf --> django.conf. But we only want django as an import
        cleaned_name, _, _ = name.partition('.')
        imports.add(cleaned_name)

    packages = set(imports) - set(set(candidates) & set(imports))
    logging.debug('Found packages: {0}'.format(packages))

    third_module = sorted(
        func for func in imports
        if func not in _standard_module and func not in _bulitin_module and func not in myself_module
    )
    return third_module


if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--path", "--root_path", dest='root_path', type=str, help='查找根目录', default=os.getcwd())
    parser.add_argument("--f", "--fileformat", dest="fileformat", type=str, help='文件格式扩展名约束', default='')
    parser.add_argument("--include", type=str, help='文件名应包含该字符串', default='')
    parser.add_argument("--uninclude", type=str, help='文件路径不包含该字符串, 优先级高于 include 参数', default='')
    parser.add_argument("--d", "--depth", dest="depth", type=int, help='查找深度, None代表无限深度', default=None)
    parser.add_argument("--r", "--relative", dest="relative",
                        help='返回结果是否采用以 root_path 为参考的相对路径, False时为绝对路径',
                        default=False, action='store_true')  # 训练的worker数量
    parser.add_argument("--v", "--verbose", dest="verbose", help='是否打印查找结果', default=True, action='store_true')
    args = parser.parse_args()
    args = vars(args)
    get_file_paths(**args)
