#!/usr/bin/env python3
import os
import sys

def main():
    # ----- Parse arguments -----
    input_cat = ""
    neorg_dir = ""
    for (i, arg) in enumerate(sys.argv):
        if sys.argv[i-1] == "--category":
            continue
        if arg == "--category":
            if len(sys.argv) <= i + 1:
                print("Category not valid")
                help()
                return
            input_cat = sys.argv[i + 1]
            continue
        elif len(sys.argv) - 1 < i:
            help()
            return
        neorg_dir = arg

    while not os.path.isdir(neorg_dir):
        print("Input target directory")
        print("> ", end="")
        neorg_dir = os.path.abspath(os.path.expanduser(input()))
        if not os.path.isdir(neorg_dir):
            print("Directory not found")

    # ----- Main Program -----

    categories = []
    category_dict = {}
    no_categories = []
    fill_category_dict(neorg_dir, neorg_dir, categories, no_categories, category_dict)

    # ----- Output Result -----

    if len(categories) == 0:
        print("No Neorg files found in", neorg_dir)
        return

    # Category selection
    if input_cat == "":
        # Print all categories
        print("Found", len(categories), "categories:")
        for cat in categories:
            print("\t", cat, sep="")

        # Print empty warning
        if len(no_categories) != 0:
            print("Found", len(no_categories), "files with no categories, enter \"None\" to list them")

        # Get user input
        print("Please enter a category to search for:")
        print("> ", end="")
        input_cat = input()

    # special case for no category files
    if input_cat == "None":
        print("Found", len(no_categories), "files with no categories:")
        for file in no_categories:
            print("\t", file)
        return

    # check if input is valid
    if input_cat not in categories:
        print("Category not valid")
        return

    # print found files
    print("Found ", len(category_dict[input_cat]), " files with category ", input_cat, " in ", neorg_dir, ":", sep="")
    for file in category_dict[input_cat]:
        print("\t", file)


def help():
    print("Usage: neorg-tags.py <path to neorg directory>")
    print("\tOptionally add --category <category> to search for a specific category")


def fill_category_dict(curr_dir, neorg_dir, categories, no_categories, category_dict):
    for file in os.listdir(curr_dir):

        # if file is a not hidden directory, recurse 
        if os.path.isdir(os.path.join(curr_dir, file)) and not file.startswith("."):
            fill_category_dict(os.path.join(curr_dir, file), neorg_dir, categories, no_categories, category_dict)

        if not file.endswith(".norg"):
            continue

        with open(os.path.join(curr_dir, file)) as f:

            # get all the categories
            curr_categories = []
            for line in f:
                if not line.startswith("categories: "):
                    continue

                curr_categories = line.split("categories: ")[1].split(" ")
                curr_categories[-1] = curr_categories[-1].strip() # remove trailing newline

            # add all categories to categories and category_dict
            for category in curr_categories:
                filepath = os.path.join(curr_dir, file).split(neorg_dir)[1]
                if category == "":
                    no_categories.append(filepath)
                    continue
                if category not in categories:
                    categories.append(category)
                if category not in category_dict:
                    category_dict[category] = []
                category_dict[category].append(filepath)


if __name__ == "__main__":
    main()