from astoptimizer import UNSET
from astoptimizer.compatibility import (
    NONE_TYPE, INT_TYPES, STR_TYPES, UNICODE_TYPE,
    is_immutable)
import sys

ANY_TYPE = object()

DEFAULT_FEATURES = ('builtin_types', 'math', 'string')

# Don't optimize operations on non-BMP strings because the
# result is different on narrow and wide builds
def optimize_unicode(config, text):
    if sys.version_info >= (3,3):
        return True
    if 'pythonbin' in config.features:
        return True
    if sys.maxunicode == 0x10ffff:
        return all(ord(char) <= 0xffff for char in text)
    else:
        return all(not 0xd800 <= ord(char) <= 0xdfff for char in text)

class BaseFunction:
    def __init__(self, narg, *arg_types, **kw):
        """
        Options:
         - check_args: callback, see below
         - catch: exception or list of exceptions catched when the function is called

        def check_args(config, args, result):
           return True
        """
        if isinstance(narg, (tuple, list)):
            self.min_narg, self.max_narg = narg
        elif isinstance(narg, INT_TYPES):
            self.min_narg = narg
            self.max_narg = narg
        else:
            raise TypeError("expect tuple or int, got %s" % type(narg).__name__)
        self.arg_types = arg_types
        self.check_result_callback = kw.pop("check_result", None)
        self.catch = kw.pop("catch", None)
        if kw:
            raise ValueError("unknown arguments: %s" % ', '.join(kw.keys()))

    def _check_args(self, config, args):
        if not(self.min_narg <= len(args) <= self.max_narg):
            return False
        for index, arg in enumerate(args):
            types = self.arg_types[index]
            if types is ANY_TYPE:
                continue
            if not isinstance(arg, types):
                return False
        return True

    def check_result(self, config, result):
        if not config.check_result(result):
            return False
        if self.check_result_callback is not None:
            if not self.check_result_callback(config, result):
                return False
        return True

    def _call(self, config, func, args):
        try:
            result = func(*args)
        except OverflowError:
            return UNSET
        except Exception:
            if self.catch is None:
                raise
            err = sys.exc_info()[1]
            if isinstance(err, self.catch):
                return UNSET
            else:
                raise
        if result is UNSET:
            return UNSET
        if not self.check_result(config, result):
            return UNSET
        return result


class Function(BaseFunction):
    def __init__(self, func, narg, *arg_types, **kw):
        # def check_args(config, args):
        #    if ...:
        #        # can be done at compile time
        #        return True
        #    else:
        #        # must be done at runtime
        #        return False
        self.check_args_callback = kw.pop("check_args", None)
        BaseFunction.__init__(self, narg, *arg_types, **kw)
        self.func = func

    def call(self, config, args):
        if not self._check_args(config, args):
            return UNSET
        if self.check_args_callback is not None:
            if not self.check_args_callback(config, args):
                return UNSET
        return self._call(config, self.func, args)


class VarargFunction(Function):
    def __init__(self, func, **kw):
        Function.__init__(self, func, 0, **kw)

    def _check_args(self, config, args):
        return True


class Method(BaseFunction):
    def __init__(self, obj_type, method, narg, *arg_types, **kw):
        # def check_args(config, obj, args):
        #    if ...:
        #        # can be done at compile time
        #        return True
        #    else:
        #        # must be done at runtime
        #        return False
        self.check_args_callback = kw.pop("check_args", None)
        BaseFunction.__init__(self, narg, *arg_types, **kw)
        self.obj_type = obj_type
        self.method = method

    def call(self, config, obj, args):
        if not self._check_args(config, args):
            return UNSET
        if self.check_args_callback is not None:
            if not self.check_args_callback(config, obj, args):
                return UNSET
        func = getattr(obj, self.method)
        return self._call(config, func, args)


def display_info(message):
    sys.stderr.write("%s\n" % message)
    sys.stderr.flush()

def display_warning(message):
    sys.stderr.write("WARNING: %s\n" % message)
    sys.stderr.flush()

class Config:
    def __init__(self, *features):
        self.features = set()

        # Hardcoded constants
        self._constants = {
            'True': True,
            'False': False,
            'None': None,
        }

        # List of functions that can be called at compile time (no border
        # effect), dictionnary of name => (check_callback, func) with
        # check_callback returning False if the call cannot be done:
        #
        # def check_callback(config, args):
        #    if ...:
        #        # can be done at compile time
        #        return True
        #    else:
        #        # must be done at runtime
        #        return False
        self.functions = {}

        # List of methods that can be called at compile time (no border
        # effect), dictionnary of type => {attr_name => check_callback}
        # with check_callback returning False if the call cannot be done:
        #
        # def check_callback(config, obj, args):
        #    if ...:
        #        # can be done at compile time
        #        return True
        #    else:
        #        # must be done at runtime
        #        return False
        self.methods = {}

        # List of modules used by the configuration
        self._modules = set()

        # Minimum and maximum value of integers (int and long types).
        # Arbitrary limit to limit memory footprint and disk space
        # (max 128 bits per number).
        self.min_int = -2 ** 128
        self.max_int = 2 ** 128 - 1

        # Minimum and maximum value of C long type
        self.min_c_long = -2 ** 31
        self.max_c_long = 2 ** 31 - 1

        # Maximum length of a string in characters (bytes type) or bytes (str
        # type). CPython peephole optimizer uses a limit of 20.
        self.max_string_length = 4096

        # Maximum number of items of a frozenset, set or tuple
        self.max_tuple_length = 20

        # Remove dead code?
        # Example: "if 0: print('debug')" => "pass"
        self.remove_dead_code = True

        # Remove "almost dead" code?
        # More aggressive removal of dead code. For example, "obj.attr" is
        # considered as dead code, whereas is may be used to check if an
        # attribute exists:
        #     try:
        #         obj.attr
        #     except AttributeError:
        #         pass
        self.remove_almost_dead_code = False

        # Remove documentation strings?
        # False by default, True if the option -O is specified
        # in the command line
        self.remove_docstring = sys.flags.optimize

        # Callback to display a information message, prototype:
        #
        # def info(message):
        #     ...
        self.info_func = display_info

        # Callback to display a warning message, prototype:
        #
        # def warning(message):
        #     ...
        self.warning_func = display_warning

        # The maximum size is at least 2^31-1 according to Python 2.7
        # documentation
        self.max_size = 2**31 - 1

        # Limit of loop unroll
        # Use 0 to disable loop unroll.
        self.unroll_limit = 100

        # Experimental support of assignment (support of variables)
        self.use_experimental_vars = False

        # Enable features
        for feature in DEFAULT_FEATURES:
            self.enable(feature)
        for feature in features:
            self.enable(feature)

    def enable(self, feature):
        if feature in self.features:
            return
        self.features.add(feature)
        modname = "config_%s" % feature
        mod = __import__("astoptimizer.%s" % modname)
        mod = getattr(mod, modname)
        mod.setup_config(self)

    def enable_all_optimizations(self):
        for feature in (
            'builtin_funcs',
            'builtin_types',
            'math',
            'platform',
            'pythonbin',
            'pythonenv',
            'string',
            'struct',
        ):
            self.enable(feature)
        self.remove_dead_code = True
        #self.use_experimental_vars = True
        self.remove_almost_dead_code = True

    def _add_module(self, name):
        if '.' not in name:
            return
        module_name = name.rsplit('.', 1)[0]
        self._modules.add(module_name)

    def is_module_used(self, name):
        return (name in self._modules)

    def add_constant(self, name, value):
        if not is_immutable(value):
            raise TypeError("Unsupport constant type: %r (%s)"
                            % (value, type(value).__name__))
        if name in self._constants:
            raise ValueError("Duplicate constant: %r" % (name,))
        self._add_module(name)
        self._constants[name] = value

    def get_constant(self, name):
        return self._constants.get(name, UNSET)

    def add_func(self, name, func):
        if name in self.functions:
            raise ValueError("Duplicate function: %r" % (name,))
        self._add_module(name)
        self.functions[name] = func

    def add_method(self, method):
        obj_type = method.obj_type
        attr = method.method
        if obj_type in self.methods:
            if attr in self.methods[obj_type]:
                raise ValueError("Duplicate method %s.%r" % (obj_type, attr))
            self.methods[obj_type][attr] = method
        else:
            self.methods[obj_type] = {attr: method}

    def check_result(self, result):
        obj_type = type(result)
        if obj_type in (bool, NONE_TYPE, float, complex):
            return True
        if obj_type in INT_TYPES:
            return (self.min_int <= result <= self.max_int)
        if obj_type in STR_TYPES:
            if len(result) > self.max_string_length:
                return False
            if (obj_type == UNICODE_TYPE
            and not optimize_unicode(self, result)):
                return False
            return True
        if obj_type in (tuple, frozenset):
            if len(result) > self.max_tuple_length:
                return False
            return all(self.check_result(arg) for arg in result)
        return False

