# -*- coding: utf-8 -*-
from argparse import ArgumentParser
from collections import OrderedDict
from datetime import datetime
from sys import version_info

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


PY2 = version_info[0] == 2


STRTYPES = ["unicode", "bytes", "bytearray"]
STREAM_TYPES = ["bytesio", "stringio", "iobase"]
TYPES = STRTYPES + STREAM_TYPES


string_methods = [
    "encode",
    "decode",
    "replace",
    "split",
    "rsplit",
    "join",
    "capitalize",
    "casefold",
    "title",
    "center",
    "expandtabs",
    "partition",
    "ljust",
    "lower",
    "lstrip",
    "rjust",
    "rstrip",
    "rpartition",
    "splitlines",
    "strip",
    "swapcase",
    "translate",
    "upper",
    "zfill",
    "formatmap",
    "removeprefix",
    "removesuffix",
]


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


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


if PY2:
    string_patch_template = """
PyObject *create_{0}_hook_module(PyObject *self, PyObject *arg) {{
    return Py_InitModule("{0}_hooks", methods);
}}
"""
else:
    string_patch_template = """
PyObject *create_{0}_hook_module(PyObject *self, PyObject *arg) {{
    return PyModule_Create(&{0}_module_def);
}}
"""

method_entry_template = (
    '{{"apply_{1}_hook", apply_{0}_{1}_hook, METH_O, "Enable {0}.{1} hook"}}'
)

method_table_template = """
static PyMethodDef methods[] = {{
    {},
    {{NULL, NULL, 0, NULL}},
}};
"""

method_def_template = """
static struct PyModuleDef {0}_module_def = {{
    PyModuleDef_HEAD_INIT,
    "{0}_hooks",
    "methods for hooking {0} methods",
    -1,
    methods,
    NULL,
    NULL,
    NULL,
    NULL
}};
"""

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 = """/*
* Copyright © {} Contrast Security, Inc.
* See https://www.contrastsecurity.com/enduser-terms-0317a for more details.
*/
""".format(
    YEAR
)

source_file_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>

{};
{};
""".format(
    C_HEADER, fastcall_typedef, fastcall_kwargs_typedef
)


pytype_map = {
    "unicode": ("PyUnicode_Type", hookspy.find_unicode_hook),
    "bytearray": ("PyByteArray_Type", hookspy.find_bytearray_hook),
    "bytes": ("PyBytes_Type", hookspy.find_bytes_hook),
    "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 autogen(typename):
    pytype, spyfunc = pytype_map[typename]
    stream_type = typename in STREAM_TYPES
    methods = stream_methods if stream_type else string_methods
    hook_template = stream_hook_template if stream_type else string_hook_template
    patch_template = stream_patch_template if stream_type else string_patch_template

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

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

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

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

        if not stream_type:
            hook_method = "CREATE_HOOK_METHOD({}, {}, {})".format(
                pytype, hook_name, offset
            )
            create_hook_methods.append(hook_method)

            hook_method_entries.append(method_entry_template.format(typename, name))

        hook_points.append(hook_template.format(pytype, hook_name, macro_name))
        reverse_points.append(
            stream_reverse_hook_template.format(pytype, 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")
    for m in create_hook_methods:
        print(m)
    print("\n")
    if not stream_type:
        print(method_table_template.format(",\n    ".join(hook_method_entries)))
        print("\n")
        if not PY2:
            print(method_def_template.format(typename))
        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)
