import os
import sys

import rich
from rich.text import Text
from rich.tree import Tree


def _get_file_size(file):
    size = os.path.getsize(file)
    if size == 0:
        return "空文件"
    num = 0
    while size > 1024:
        size /= 1024
        num += 1
    unit = ["", "KB", "MB", "GB", "TB"]
    return f"{size:.2f}".rstrip(".0") + unit[num]


def _show_dir(path, tree=None):
    if tree is None:
        tree = Tree(f"[bold magenta]{os.path.abspath(path)}")
    for file in os.listdir(path):
        file_path = os.path.join(path, file)
        if (os.path.isdir(file_path)):
            parent = tree.add(f"[bold magenta]{file}")
            _show_dir(file_path, parent)
        else:
            text_filename = Text(file, "green")
            text_filename.highlight_regex(r"\.[^.]+$", "bold red")
            text_filename.append(f" ({_get_file_size(file_path)})", "bold blue")
            tree.add(text_filename)
    return tree


def color_print_dir(path):
    rich.print(_show_dir(path))


if __name__ == '__main__':
    rich.print(_show_dir(sys.argv[1]))