from argparse import ArgumentParser
from collections import OrderedDict
from datetime import datetime

from . import hookspy
from .body import create_hook
from .signatures import get_method_signature, fastcall_typedef, fastcall_kwargs_typedef


STREAM_TYPES = ["bytesio", "stringio", "iobase"]
TYPES = STREAM_TYPES

stream_methods = [
    "read",
    "readline",
    "readlines",
    "read1",
]


stream_hook_template = """    ADD_STREAM_HOOK(StreamType, {1}, {2});"""
stream_reverse_hook_template = """    REVERSE_STREAM_HOOK(StreamType, {1}, {2});"""


stream_patch_template = """
int patch_{}_methods(PyTypeObject *StreamType) {{

{}

    /* Indicates success */
    return 0;
}}
"""

stream_reverse_template = """
void reverse_{}_methods(PyTypeObject *StreamType) {{

{}

}}
"""

YEAR = datetime.now().year

C_HEADER = f"""/*
* Copyright © {YEAR} Contrast Security, Inc.
* See https://www.contrastsecurity.com/enduser-terms-0317a for more details.
*/
"""

source_file_header = f"""{C_HEADER}/* THIS FILE WAS AUTOMATICALLY GENERATED BY HOOKSPY */
/* Python requires its own header to always be included first */
#define PY_SSIZE_T_CLEAN
#include <Python.h>

#include <contrast/assess/patches.h>
#include <contrast/assess/propagate.h>
#include <contrast/assess/scope.h>
#include <contrast/assess/utils.h>

{fastcall_typedef};
{fastcall_kwargs_typedef};
"""


pytype_map = {
    "bytesio": hookspy.find_bytesio_hook,
    "stringio": hookspy.find_stringio_hook,
    "iobase": hookspy.find_iobase_hook,
}


def transform_name(name):
    # This is kind of annoying, but the name of this method in the C-level
    # method table has an underscore.
    return name.replace("formatmap", "format_map")


def find_hooks(spyfunc, methods):
    hook_methods = OrderedDict()
    for name in methods:
        result = spyfunc(transform_name(name))
        if result is None:
            continue
        hook_methods[name] = result
    return hook_methods


def raise_unsupported_type(typename):
    raise ValueError(f"Unsupported type: {typename}")


def autogen(typename):
    spyfunc = pytype_map[typename]
    stream_type = typename in STREAM_TYPES
    methods = stream_methods if stream_type else []
    hook_template = (
        stream_hook_template if stream_type else raise_unsupported_type(typename)
    )
    patch_template = (
        stream_patch_template if stream_type else raise_unsupported_type(typename)
    )

    macros = []
    declarations = []
    hook_funcs = []
    hook_points = []
    reverse_points = []
    hook_methods = find_hooks(spyfunc, methods)

    for name, (offset, argtype) in hook_methods.items():
        signature = get_method_signature(argtype)

        hook_definition, hook_name = create_hook(typename, name, signature, stream_type)
        declarations.append(f"{signature} {hook_name}_orig;")
        hook_funcs.append(hook_definition)

        macro_name = f"{typename}_{name}_OFFSET".upper()
        macros.append(f"#define {macro_name} {offset}")

        hook_points.append(hook_template.format("", hook_name, macro_name))
        reverse_points.append(
            stream_reverse_hook_template.format("", hook_name, macro_name)
        )

    print(source_file_header)
    for m in macros:
        print(m)
    print("\n")
    for d in declarations:
        print(d)
    print("\n")
    for f in hook_funcs:
        print(f)
    print("\n")
    print(patch_template.format(typename, "\n".join(hook_points)))

    if stream_type:
        print(stream_reverse_template.format(typename, "\n".join(reverse_points)))


def main():
    p = ArgumentParser("Automatically generate Python type hooks")
    p.add_argument("type", type=str, help="Python type", choices=TYPES)
    args = p.parse_args()

    autogen(args.type)
