import argparse
import os
import subprocess
import sys
from tempfile import TemporaryDirectory
from typing import Tuple, Optional, List, Union

# since os.PathLike breaks python typechecking
path = bytes
drv_output = Tuple[path, bytes]


def add_root_args(gcroots_dir: Optional[bytes],
                  symlink_name: path) -> List[Union[path, str]]:
    if gcroots_dir is None:
        return []
    else:
        return ['--add-root',
                os.path.join(gcroots_dir, symlink_name),
                '--indirect']


def parse_drv_path(path: path) -> Optional[drv_output]:
    (nix_store, drv_file) = os.path.split(path)

    # drv “filenames” printed by nix-instantiate are of the form
    # ${hash}-${name}.drv[!${output}] where name, hash and output
    # may not contain any '!' among other things
    drv_parts = drv_file.split(b'!')
    if len(drv_parts) > 2:
        return None
    elif len(drv_parts) == 1:
        # default output
        return (path, b'out')
    else:
        return (os.path.join(nix_store, drv_parts[0]), drv_parts[1])


def derivation_outputs(attr: str, gcroots_dir: Optional[path] = None
                       ) -> List[drv_output]:
    cmd: List[Union[str, path]] = ['nix-instantiate', '-Q', '-E']

    # TODO(sterni): make this work with custom stuff like vuizuvi and depot
    expr = """
        let
          pkgs = import <nixpkgs> {{}};
        in
          builtins.map (o: pkgs.{attr}."${{o}}") pkgs.{attr}.outputs
    """.format(attr=attr)

    cmd.append(expr)
    cmd.extend(add_root_args(gcroots_dir, b'instantiation-result'))

    instantiate = subprocess.run(cmd, capture_output=True)

    if instantiate.returncode != 0:
        print(instantiate.stderr, file=sys.stderr)
        return []
    else:
        return [o for o in map(parse_drv_path, instantiate.stdout.splitlines())
                if o is not None]


def main():
    parser = argparse.ArgumentParser(description='nix-shell for man pages')
    parser.add_argument('attribute', metavar='ATTR',
                        help='nixpkgs attribute holding the desired package')
    parser.add_argument('section_or_page', metavar='SECTION|PAGE',
                        nargs='?')
    parser.add_argument('page', metavar='PAGE', nargs='?')

    args = parser.parse_args()

    # first arg is always the derivation attribute
    drv_attr = args.attribute

    # second arg is either the page or the section number
    if args.section_or_page is not None and args.section_or_page.isdigit():
        section = int(args.section_or_page)
        page = None
    else:
        section = 1
        page = args.section_or_page

    # last arg is the page. if it is missing and was not
    # given as the second, use drv_attr
    if args.page is not None:
        page = args.page
    elif page is None:
        page = drv_attr

    print('{}({}) from pkgs.{}'.format(page, section, drv_attr))

    with TemporaryDirectory(prefix='nman-') as gcroots:
        outputs = derivation_outputs(drv_attr,
                                     gcroots_dir=gcroots.encode('utf-8'))

        for (p, o) in outputs:
            print('{}: {}'.format(o, os.path.realpath(p)))


if __name__ == '__main__':
    main()