# Copyright 2017 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""fd.py is a fascinating directory changer to save your time in typing.

fd.py is intended to be used through shell function "fd()" in fx-env.sh.

See examples by
$ fd --help

fd stores two helper files, fd.txt and fd.pickle in $FUCHSIA_DIR/out/.
If that directory does not exists, fd will create one.
"""

import argparse
import os
import pickle
import sys
import termios
import tty

SEARCH_BASE = os.environ["FUCHSIA_DIR"]  # or 'HOME'
STORE_DIR = SEARCH_BASE + "/out/"
DIRS_FILE = STORE_DIR + "fd.txt"
PICKLE_FILE = STORE_DIR + "fd.pickle"

EXCLUDE_DIRS = [
    '"*/.git"',
    "./build",
    "./out",
    "./prebuilt",
    "./third_party",
    "./zircon/build",
    "./cmake-build-debug",
    "./zircon/third_party",
]


def eprint(*args, **kwargs):
    print(*args, file=sys.stderr, **kwargs)


class Trie(object):
    """Class Trie."""

    def __init__(self):
        self.name = ""  # != path up to here from the root. Key is a valid
        # complete one.
        self.vals = []
        self.kids = {}

    def __getitem__(self, name, idx=0):
        if self.name == name:
            return self.vals
        if idx == name.__len__() or name[idx] not in self.kids:
            return None
        return self.kids[name[idx]].__getitem__(name, idx + 1)

    def __setitem__(self, name, val, idx=0):
        if idx < name.__len__():
            self.kids.setdefault(name[idx], Trie()).__setitem__(
                name, val, idx + 1
            )
            return
        self.name = name
        self.vals.append(val)

    def __contains__(self, name):
        return self[name] is not None

    def walk(self):
        descendants = []
        if self.name:
            descendants.append(self.name)
        for k in self.kids:
            descendants.extend(self.kids[k].walk())
        return descendants

    def prefixed(self, name, idx=0):
        if idx < name.__len__():
            if name[idx] in self.kids:
                return self.kids[name[idx]].prefixed(name, idx + 1)
            return []

        return self.walk()


def build_trie():
    """build trie.

    Returns:
      Trie
    """

    def build_find_cmd():
        paths = []
        for path in EXCLUDE_DIRS:
            paths.append("{} {}".format("-path", path))
        return (
            r"cd {}; find . \( {} \) -prune -o -type d -print > " "{}"
        ).format(SEARCH_BASE, " -o ".join(paths), DIRS_FILE)

    if not os.path.exists(STORE_DIR):
        os.makedirs(STORE_DIR)

    cmd_str = build_find_cmd()
    os.system(cmd_str)

    t = Trie()
    with open(DIRS_FILE, "r") as f:
        for line in f:
            line = line[2:][:-1]
            tokens = line.split("/")
            if tokens.__len__() == 0:
                continue
            target = tokens[-1]

            t[target] = line

    return t


def get_trie():
    """get_trie.

    Returns:
      trie
    """

    def save_pickle(obj):
        with open(PICKLE_FILE, "wb+") as f:
            pickle.dump(obj, f, protocol=pickle.HIGHEST_PROTOCOL)

    def load_pickle():
        with open(PICKLE_FILE, "rb") as f:
            return pickle.load(f)

    if os.path.exists(PICKLE_FILE):
        return load_pickle()

    t = build_trie()
    save_pickle(t)
    return t


def button(idx):
    """button maps idx to an ascii value."""
    ascii = 0
    if 0 <= idx <= 8:
        ascii = ord("1") + idx
    elif 9 <= idx <= 34:
        ascii = ord("a") + idx - 9
    elif 35 <= idx <= 60:
        ascii = ord("A") + idx - 35
    elif 61 <= idx <= 75:
        ascii = ord("!") + idx - 61
    return str(chr(ascii))


def get_button():  # Unix way
    fd = sys.stdin.fileno()
    old_settings = termios.tcgetattr(fd)
    try:
        tty.setraw(sys.stdin.fileno())
        ch = sys.stdin.read(1)
    finally:
        termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
    return ch


def choose_options(t, key, choice):
    # Build options by the given key
    if key in t:
        options = t[key]
    else:
        prefixed_keys = t.prefixed(key)
        options = []
        for pk in prefixed_keys:
            options.extend(t[pk])

    options = sorted(options)
    if options.__len__() == 0:
        eprint("No such directory: {}".format(key))
        return None
    elif options.__len__() == 1:
        return options[0]
    elif options.__len__() > 75:  # See def button() for the limit.
        eprint(
            'Too many ({}) results for "{}". '
            "Refine your prefix or time to buy 4K "
            "monitor\n".format(options.__len__(), key)
        )
        return None

    def list_choices(l):
        for i in range(l.__len__()):
            eprint("[{}] {}".format(button(i), l[i]))
        eprint()

    choice_dic = {}
    for idx, val in enumerate(options):
        choice_dic[button(idx)] = val

    if choice is not None and choice not in choice_dic:
        # Invalid pre-choice
        eprint('Choice "{}" not available\n'.format(choice))

    if choice not in choice_dic:
        list_choices(options)
        choice = get_button()
        if choice not in choice_dic:
            return None

    return choice_dic[choice]


def main():
    def parse_cmdline():
        example_commands = """

[eg] # Use "fd" for autocompletion (See //scripts/fx-env.sh)
  $ fd ral        # change directory to an only option: ralink
  $ fd wlan       # shows all "wlan" directories and ask to choose
  $ fd wlan 3     # change directory matching to option 3 of "fd wlan"
  $ fd [TAB]      # Autocomplete subdirectories from the current directory
  $ fd //[TAB]    # Autocomplete subdirectories from ${FUCHSIA_DIR}
  $ fd --rebuild  # rebuilds the directory structure cache
"""
        p = argparse.ArgumentParser(
            description="A fascinating directory changer",
            epilog=example_commands,
            formatter_class=argparse.RawDescriptionHelpFormatter,
        )

        p.add_argument(
            "--rebuild", action="store_true", help="rebuild the directory DB"
        )
        p.add_argument("--base", type=str, default=None)
        p.add_argument("target", nargs="?", default="")
        p.add_argument("choice", nargs="?", default=None)

        # Redirect help messages to stderr
        if len(sys.argv) == 2:
            if sys.argv[1] in ["-h", "--help"]:
                eprint(p.format_help())
                print(".")  # Stay at the current directory
                sys.exit(0)

        return p.parse_args()

    def get_abs_path(relative_dir):
        if relative_dir is not None:
            return os.path.join(SEARCH_BASE, relative_dir)
        return os.getcwd()

    def derive_dest(target):
        if not target:
            # To test if this command was invoked just to rebuild
            return get_abs_path(".") if args.rebuild is False else os.getcwd()

        if target[:2] == "//":
            target = target[2:]

        candidate = target

        # Do not guess-work when the user specifies an option to intend to use.
        # Do guess work otherwise.
        if not args.choice:
            if os.path.exists(candidate):
                return candidate

            candidate = get_abs_path(target)
            if os.path.exists(candidate):
                return candidate

            candidate = os.path.abspath(target)
            if os.path.exists(candidate):
                return candidate

        t = get_trie()
        return get_abs_path(choose_options(t, target, args.choice))

    args = parse_cmdline()
    if args.base:
        global SEARCH_BASE
        SEARCH_BASE = args.base

    if args.rebuild:
        os.remove(PICKLE_FILE)

    dest = derive_dest(args.target)
    dest = os.path.normpath(dest)
    print(dest)


if __name__ == "__main__":
    try:
        main()
    except Exception as e:  # Catch all
        eprint(e.message, e.args)
        print(".")  # Stay at the current directory upon exception