#! /usr/bin/env python3
# -*- coding: utf-8 -`-
"""
Code generation script for class methods
to be exported as public API
"""
import argparse
import ast
import astor
import os
from pathlib import Path
import sys
import yapf.yapflib.yapf_api as formatter

from textwrap import indent

PREFIX = '_generated'

HEADER = """# ***********************************************************
# ******* WARNING: AUTOGENERATED! ALL EDITS WILL BE LOST ******
# *************************************************************
from ._run import GLOBAL_RUN_CONTEXT, _NO_SEND
from ._ki import LOCALS_KEY_KI_PROTECTION_ENABLED

    
"""

TEMPLATE = """locals()[LOCALS_KEY_KI_PROTECTION_ENABLED] = True
try:
    return {} GLOBAL_RUN_CONTEXT.{}.{}
except AttributeError:
    raise RuntimeError('must be called from async context')
"""


def is_function(node):
    """Check if the AST node is either a function
    or an async function
    """
    if (
        isinstance(node, ast.FunctionDef)
        or isinstance(node, ast.AsyncFunctionDef)
    ):
        return True
    return False


def is_public(node):
    """Check if the AST node has a _public decorator
    """
    if not is_function(node):
        return False
    for decorator in node.decorator_list:
        if isinstance(decorator, ast.Name) and decorator.id == "_public":
            return True
    return False


def get_public_methods(tree):
    """ Return a list of methods marked as public.
    The function walks the given tree and extracts
    all objects that are functions which are marked
    public.
    """
    for node in ast.walk(tree):
        if is_public(node):
            yield node


def create_passthrough_args(funcdef):
    """Given a function definition, create a string that represents taking all
    the arguments from the function, and passing them through to another
    invocation of the same function.

    Example input: ast.parse("def f(a, *, b): ...")
    Example output: "(a, b=b)"
    """
    call_args = []
    for arg in funcdef.args.args:
        call_args.append(arg.arg)
    if funcdef.args.vararg:
        call_args.append("*" + funcdef.args.vararg.arg)
    for arg in funcdef.args.kwonlyargs:
        call_args.append(arg.arg + "=" + arg.arg)
    if funcdef.args.kwarg:
        call_args.append("**" + funcdef.args.kwarg.arg)
    return "({})".format(", ".join(call_args))


def gen_public_wrappers_source(source_path: Path, lookup_path: str) -> str:
    """Scan the given .py file for @_public decorators, and generate wrapper
    functions.

    """
    generated = [HEADER]
    source = astor.code_to_ast.parse_file(source_path)
    for method in get_public_methods(source):
        # Remove self from arguments
        assert method.args.args[0].arg == "self"
        del method.args.args[0]

        # Remove decorators
        method.decorator_list = []

        # Create pass through arguments
        new_args = create_passthrough_args(method)

        # Remove method body without the docstring
        if ast.get_docstring(method) is None:
            del method.body[:]
        else:
            # The first entry is always the docstring
            del method.body[1:]

        # Create the function definition including the body
        func = astor.to_source(method, indent_with=' ' * 4)

        # Create export function body
        template = TEMPLATE.format(
            'await' if isinstance(method, ast.AsyncFunctionDef) else '',
            lookup_path,
            method.name + new_args,
        )

        # Assemble function definition arguments and body
        snippet = func + indent(template, ' ' * 4)

        # Append the snippet to the corresponding module
        generated.append(snippet)
    return "\n".join(generated)


def matches_disk_files(new_files):
    for new_path, new_source in new_files.items():
        if not os.path.exists(new_path):
            return False
        with open(new_path, "r", encoding="utf-8") as old_file:
            old_source = old_file.read()
        if old_source != new_source:
            return False
    return True


def process(sources_and_lookups, *, do_test):
    new_files = {}
    for source_path, lookup_path in sources_and_lookups:
        print("Scanning:", source_path)
        new_source = gen_public_wrappers_source(source_path, lookup_path)
        dirname, basename = os.path.split(source_path)
        new_path = os.path.join(dirname, PREFIX + basename)
        new_files[new_path] = new_source
    if do_test:
        if not matches_disk_files(new_files):
            print("Generated sources are outdated. Please regenerate.")
            sys.exit(1)
        else:
            print("Generated sources are up to date.")
    else:
        for new_path, new_source in new_files.items():
            with open(new_path, "w", encoding="utf-8") as f:
                f.write(new_source)
        print("Regenerated sources successfully.")


# This is in fact run in CI, but only in the formatting check job, which
# doesn't collect coverage.
def main():  # pragma: no cover
    parser = argparse.ArgumentParser(
        description='Generate python code for public api wrappers'
    )
    parser.add_argument(
        '--test',
        '-t',
        action='store_true',
        help='test if code is still up to date'
    )
    parsed_args = parser.parse_args()

    source_root = Path.cwd()
    # Double-check we found the right directory
    assert (source_root / "LICENSE").exists()
    core = source_root / 'trio/_core'
    to_wrap = [
        (core / "_run.py", "runner"),
        (core / "_io_windows.py", "runner.io_manager"),
        (core / "_io_epoll.py", "runner.io_manager"),
        (core / "_io_kqueue.py", "runner.io_manager"),
    ]

    process(to_wrap, do_test=parsed_args.test)


if __name__ == '__main__':  # pragma: no cover
    main()
