#!/usr/bin/env python
"""
Command-line interface to access the documentation of ABINIT input variables.

This script allows users to search for variables, view their documentation,
browse them in an external web browser, and visualize variable dependencies.
It also provides documentation for AbiPy's TaskManager and Scheduler options.

Examples:
    Show documentation for the `ecut` variable:
        $ abidoc.py man ecut

    Search for variables related to `paw`:
        $ abidoc.py find paw

    Open the official ABINIT documentation for `acell` in a web browser:
        $ abidoc.py browse acell

    Visualize dependencies for `tolvrs` using Graphviz:
        $ abidoc.py graphviz tolvrs

    List all variables in the `base` varset:
        $ abidoc.py list --mode s
"""

from __future__ import annotations

import argparse
import sys

from monty.termcolor import cprint

import abipy.tools.cli_parsers as cli
from abipy import abilab, flowtk
from abipy.core.release import __version__


def print_vlist(vlist, options):
    """
    Print a list of variables to the terminal.

    Args:
        vlist: List of AbinitVariable objects.
        options: Namespace object containing command-line options.
    """
    for v in vlist:
        print(repr(v))

    if options.verbose:
        for v in vlist:
            abilab.abinit_help(v)
    else:
        print("\nUse -v for more info")


def get_epilog() -> str:
    """
    Return the epilog string for the command-line parser.
    """
    return """\
Usage example:

    abidoc.py man ecut        --> Show documentation for ecut input variable.
    abidoc.py graphviz acell  --> Draw parents and children of variable with graphviz package.
    abidoc.py browse acell    --> Open url in external browser.
    abidoc.py apropos ecut    --> To search in the database for the variables related to ecut.
    abidoc.py find paw        --> To search in the database for the variables whose name contains paw.
    abidoc.py list            --> Print full list of variables.
    abidoc.py withdim natom   --> Print arrays depending on natom.
    abidoc.py scheduler       --> Document scheduler options.
    abidoc.py manager         --> Document manager options.

Use `abidoc.py --help` for help and `abidoc.py COMMAND --help` to get the documentation for `COMMAND`.
Use `-v` to increase verbosity level (can be supplied multiple times e.g -vv).
"""


def get_parser(with_epilog=False):
    """
    Return the ArgumentParser object for the script.

    Args:
        with_epilog: If True, include the epilog in the parser.
    """
    # Build the main parser.
    parser = argparse.ArgumentParser(
        epilog=get_epilog() if with_epilog else "", formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument("-V", "--version", action="version", version=__version__)

    # Parent parser for common options.
    copts_parser = argparse.ArgumentParser(add_help=False)
    copts_parser.add_argument(
        "-v",
        "--verbose",
        default=0,
        action="count",  # -vv --> verbose=2
        help="verbose, can be supplied multiple times to increase verbosity",
    )
    copts_parser.add_argument(
        "--loglevel",
        default="ERROR",
        type=str,
        help="Set the loglevel. Possible values: CRITICAL, ERROR (default), WARNING, INFO, DEBUG",
    )
    copts_parser.add_argument(
        "-c", "--codename", type=str, default="abinit", help="Code name e.g. anaddb, optic... Default: abinit"
    )

    var_parser = argparse.ArgumentParser(add_help=False)
    var_parser.add_argument("varname", help="ABINIT variable")

    # Create the parsers for the sub-commands
    subparsers = parser.add_subparsers(dest="command", help="sub-command help", description="Valid subcommands")

    # Subparser for man.
    p_man = subparsers.add_parser("man", parents=[copts_parser, var_parser], help="Show documentation for varname.")

    # Subparser for graphviz.
    p_graphviz = subparsers.add_parser(
        "graphviz",
        parents=[copts_parser, var_parser],
        help=("Draw variable dependencies with graphviz package.See https://graphviz.readthedocs.io/."),
    )
    p_graphviz.add_argument(
        "-e",
        "--engine",
        type=str,
        default="automatic",
        help=(
            "graphviz engine: ['dot', 'neato', 'twopi', 'circo', 'fdp', 'sfdp', 'patchwork', 'osage']. "
            "Default: automatic i.e. the engine is automatically selected. See http://www.graphviz.org/pdf/dot.1.pdf "
            "Use `conda install python-graphviz` or `pip install graphviz` to install the python package"
        ),
    )

    # Subparser for browse.
    p_browse = subparsers.add_parser(
        "browse", parents=[copts_parser, var_parser], help="Open documentation in browser."
    )

    # Subparser for apropos.
    p_apropos = subparsers.add_parser(
        "apropos", parents=[copts_parser, var_parser], help="Find variables related to varname."
    )
    # Subparser for find.
    p_find = subparsers.add_parser(
        "find", parents=[copts_parser, var_parser], help="Find all variables whose name contains varname."
    )
    # Subparser for require.
    # p_require = subparsers.add_parser('require', parents=[copts_parser], help="Find all variables required by varname.")

    # Subparser for withdim.
    p_withdim = subparsers.add_parser(
        "withdim", parents=[copts_parser], help="Find all arrays depending on the given dimension."
    )
    p_withdim.add_argument("dimname", help="Dimension name")

    # Subparser for list.
    p_list = subparsers.add_parser("list", parents=[copts_parser], help="List all variables.")
    p_list.add_argument(
        "--mode", default="a", help="Sort mode, `a` for alphabethical, `s` for varset, `c` for characteristics."
    )

    # Subparser for manager.
    p_manager = subparsers.add_parser("manager", parents=[copts_parser], help="Document the TaskManager options.")
    p_manager.add_argument(
        "qtype",
        nargs="?",
        default=None,
        help=(
            "Write job script to terminal if qtype='script' else "
            "document the qparams for the given QueueAdapter qtype e.g. slurm."
        ),
    )

    # Subparser for scheduler
    p_docsched = subparsers.add_parser(
        "scheduler", parents=[copts_parser], help="Document the options available in scheduler.yml."
    )

    # Subparser for abibuild
    p_abibuild = subparsers.add_parser(
        "abibuild", parents=[copts_parser], help="Show ABINIT build information and exit."
    )

    # Subparser for links
    p_urls = subparsers.add_parser("urls", parents=[copts_parser], help="Produce Yaml file with varname --> ulrs.")

    return parser


@cli.prof_main
def main():
    """
    Main entry point for the script.
    """

    def show_examples_and_exit(err_msg=None, error_code=1):
        """Display the usage of the script."""
        sys.stderr.write(get_epilog())
        if err_msg:
            sys.stderr.write("Fatal Error\n" + err_msg + "\n")
        sys.exit(error_code)

    parser = get_parser(with_epilog=True)

    try:
        options = parser.parse_args()
    except Exception:
        show_examples_and_exit(error_code=1)

    # loglevel is bound to the string value obtained from the command line argument.
    # Convert to upper case to allow the user to specify --loglevel=DEBUG or --loglevel=debug
    import logging

    numeric_level = getattr(logging, options.loglevel.upper(), None)
    if not isinstance(numeric_level, int):
        raise ValueError("Invalid log level: %s" % options.loglevel)
    logging.basicConfig(level=numeric_level)

    # Get the database of variables for codename.
    from abipy.abio.abivar_database.variables import get_codevars

    codevars = get_codevars()
    vdb = codevars[options.codename]

    if options.command == "man":
        abilab.abinit_help(options.varname)

    elif options.command == "browse":
        return vdb[options.varname].browse()

    elif options.command == "graphviz":
        if options.varname in vdb.my_varset_list:
            graph = vdb.get_graphviz(varset=options.varname, vartype=None, engine=options.engine)
        else:
            graph = vdb.get_graphviz_varname(varname=options.varname, engine=options.engine)

        import tempfile

        directory = tempfile.mkdtemp()
        print("Producing source files in:", directory)
        graph.view(directory=directory, cleanup=False)

    elif options.command == "apropos":
        vlist = vdb.apropos(options.varname)
        print_vlist(vlist, options)

    elif options.command == "find":
        vlist = [v for v in vdb.values() if options.varname in v.name]
        print("Find results:\n")
        print_vlist(vlist, options)

    elif options.command == "list":
        if options.mode == "a":
            # Alphabetical
            for i, var in enumerate(vdb.values()):
                print(i, repr(var))

        elif options.mode == "s":
            # Grouped by varset
            for section in vdb.my_varset_list:
                print(30 * "#" + " Section: " + section + " " + 30 * "#")
                print_vlist(vdb.vars_with_section(section), options)

        elif options.mode == "c":
            # Grouped by characteristics.
            for char in vdb.my_characteristics:
                print(30 * "#" + " Characteristic: " + char + 30 * "#")
                print_vlist(vdb.vars_with_char(char), options)

        else:
            raise ValueError("Wrong mode %s" % options.mode)

    elif options.command == "withdim":
        for var in vdb.values():
            if var.depends_on_dimension(options.dimname):
                cprint(repr(var), "yellow")
                print("dimensions:", str(var.dimensions), "\n")

    elif options.command == "scheduler":
        print("Options that can be specified in scheduler.yml:")
        print(flowtk.PyFlowScheduler.autodoc())
        return 0

    elif options.command == "manager":
        # Document TaskManager options and qparams.
        qtype = options.qtype

        if qtype == "script":
            manager = flowtk.TaskManager.from_user_config()
            script = manager.qadapter.get_script_str(
                job_name="job_name",
                launch_dir="workdir",
                executable="executable",
                qout_path="qout_file.path",
                qerr_path="qerr_file.path",
                stdin="stdin",
                stdout="stdout",
                stderr="stderr",
            )
            print(script)

        else:
            print(flowtk.TaskManager.autodoc())
            print("qtype supported: %s" % flowtk.all_qtypes())
            print("Use:\n\n\tabidoc.py manager slurm\n\nto show the list of qparams for slurm.\n")

            if qtype is not None:
                print("QPARAMS for %s" % qtype)
                flowtk.show_qparams(qtype)

    elif options.command == "abibuild":
        abinit_build = flowtk.AbinitBuild()
        print()
        print(abinit_build)
        print()
        if not options.verbose:
            print("Use --verbose for additional info")
        else:
            print(abinit_build.info)

    elif options.command == "urls":
        d = {}
        for codename in sorted(codevars.keys()):
            variables = codevars[codename]
            for varname, var in variables.items():
                key = var.abivarname
                key.replace("@", "_")
                d[key] = f"[{var.name}]({var.website_url})"
        from ruamel.yaml import YAML

        YAML().dump(d, sys.stdout)

    else:
        raise ValueError("Don't know how to handle command %s" % options.command)

    return 0


if __name__ == "__main__":
    sys.exit(main())
