"""
Enable optimizations on the builtins types, example: unicode.encode().
"""

from astoptimizer.compatibility import (
    PYTHON3,
    INT_TYPES,
    NATIVE_STR, BYTES_TYPE, UNICODE_TYPE, STR_TYPES,
    IMMUTABLE_ITERABLE_TYPES)
from astoptimizer.config import Method, ANY_TYPE
import codecs

def check_encoding(encoding):
    try:
        encoding = codecs.lookup(encoding).name
    except LookupError:
        return False
    if encoding not in ('ascii', 'iso8859-1', 'utf-8'):
        return False
    return True

def check_str_decode_args(config, obj, args):
    if not check_encoding(args[0]):
        return False
    if len(args) == 2:
        errors = args[1]
        if errors not in ('strict', 'replace', 'ignore'):
            return False
    return True

def check_str_encode_args(config, obj, args):
    if not check_encoding(args[0]):
        return False
    if len(args) == 2:
        errors = args[1]
        if errors not in ('strict', 'replace', 'ignore', 'backslashreplace'):
            return False
    return True

def check_index_args(config, obj, args):
    return (args[0] in obj)

def check_fill_args(config, obj, args):
    if len(args) != 2:
        return True
    fill = args[1]
    return len(fill) == 1

def check_str_join(config, obj, args):
    arg = args[0]
    str_type = type(obj)
    if not isinstance(arg, str_type):
        return all(isinstance(item, str_type) for item in arg)
    else:
        return True

def setup_config(config):
    # bytes, str, unicode
    config.add_method(Method(BYTES_TYPE, 'decode', (1, 2), NATIVE_STR, NATIVE_STR, check_args=check_str_decode_args))
    config.add_method(Method(UNICODE_TYPE, 'encode', (1, 2), NATIVE_STR, NATIVE_STR, check_args=check_str_encode_args))
    for str_type in STR_TYPES:
        config.add_method(Method(str_type, 'count', 1, str_type))
        config.add_method(Method(str_type, 'find', 1, str_type))
        config.add_method(Method(str_type, 'rfind', 1, str_type))
        config.add_method(Method(str_type, 'index', 1, str_type, check_args=check_index_args))
        config.add_method(Method(str_type, 'rindex', 1, str_type, check_args=check_index_args))

        config.add_method(Method(str_type, 'startswith', 1, str_type))
        config.add_method(Method(str_type, 'endswith', 1, str_type))

        config.add_method(Method(str_type, 'ljust', (1, 2), INT_TYPES, str_type, check_args=check_fill_args))
        config.add_method(Method(str_type, 'rjust', (1, 2), INT_TYPES, str_type, check_args=check_fill_args))
        config.add_method(Method(str_type, 'center', (1, 2), INT_TYPES, str_type, check_args=check_fill_args))
        config.add_method(Method(str_type, 'zfill', 1, INT_TYPES))

        config.add_method(Method(str_type, 'rstrip', (0, 1), str_type))
        config.add_method(Method(str_type, 'lstrip', (0, 1), str_type))
        config.add_method(Method(str_type, 'strip', (0, 1), str_type))

        config.add_method(Method(str_type, 'replace', 2, str_type, str_type))

        # don't optimize split, rsplit, splitlines because the result is a list
        config.add_method(Method(str_type, 'partition', 1, str_type))
        config.add_method(Method(str_type, 'rpartition', 1, str_type))

        if PYTHON3 and str_type == BYTES_TYPE:
            config.add_method(Method(str_type, 'join', 1, (tuple, frozenset), check_args=check_str_join))
        else:
            config.add_method(Method(str_type, 'join', 1, (tuple, frozenset) + (str_type,), check_args=check_str_join))

        config.add_method(Method(str_type, 'expandtabs', (0, 1), INT_TYPES))
        # str methods depend on the locale on Python 2
        if PYTHON3 or str_type == UNICODE_TYPE:
            for method in (
                'capitalize', 'lower', 'upper', 'title', 'swapcase',
                'isalnum', 'isalpha', 'isdigit', 'islower', 'isspace',
                'istitle', 'isupper',
            ):
                config.add_method(Method(str_type, method, 0))

        # FIXME: str.format
        # FIXME: str.translate

    # tuple
    config.add_method(Method(tuple, 'count', 1, ANY_TYPE))
    config.add_method(Method(tuple, 'index', 1, ANY_TYPE, check_args=check_index_args))

    # frozenset
    for method in (
        'copy', 'difference', 'intersection', 'union', 'symmetric_difference',
        'isdisjoint', 'issubset', 'issuperset',
    ):
        config.add_method(Method(frozenset, method, 1, IMMUTABLE_ITERABLE_TYPES))

    # int
    for int_type in INT_TYPES:
        config.add_method(Method(int_type, 'bit_length', 0))

    # float
    config.add_method(Method(float, 'is_integer', 0))
    config.add_method(Method(float, 'as_integer_ratio', 0))
    config.add_method(Method(float, 'hex', 0))

