import os
import shutil
import sys
import tempfile
from contextlib import contextmanager
from distutils import log
from distutils.errors import CompileError, LinkError


__all__ = [
    'CCompilerCapabilities'
]


@contextmanager
def chdir(path):
    cwd = os.getcwd()
    os.chdir(path)
    try:
        yield
    finally:
        os.chdir(cwd)


@contextmanager
def mkdtemp(suffix):
    path = tempfile.mkdtemp(suffix)
    try:
        with chdir(path):
            yield path
    finally:
        shutil.rmtree(path)


@contextmanager
def output(is_quiet):
    if is_quiet:  # pragma: no branch
        devnull = open(os.devnull, 'w')
        oldstderr = os.dup(sys.stderr.fileno())
        oldstdout = os.dup(sys.stdout.fileno())
        os.dup2(devnull.fileno(), sys.stderr.fileno())
        os.dup2(devnull.fileno(), sys.stdout.fileno())
    try:
        yield
    finally:
        if is_quiet:  # pragma: no branch
            os.dup2(oldstderr, sys.stderr.fileno())
            os.dup2(oldstdout, sys.stdout.fileno())
            devnull.close()


class CCompilerCapabilities:
    SIMD_SSSE3 = 0
    SIMD_SSE41 = 1
    SIMD_SSE42 = 2
    SIMD_AVX = 3
    SIMD_AVX2 = 4
    SIMD_NEON32 = 5
    SIMD_NEON64 = 6

    def __init__(self, compiler):
        self.__capabilities = dict()
        self.__cflags = []
        if compiler.compiler_type == 'msvc':
            self.__cflags = ['/WX', '/Od']  # pragma: no cover
        else:
            self.__cflags = ['-O0']
        self.__get_capabilities(compiler)

    def __has_simd_support(self, compiler, flags, define, include, content):
        quiet = True

        with mkdtemp("pybase64simdtest") as dname:
            fname = os.path.join(dname, 'simd.c')
            with open(fname, 'w') as f:
                f.write("""#include <%s>\n""" % include)
                f.write("""\
int main (int argc, char **argv) {
    %s
}
""" % content)
            with output(quiet):
                for flag in flags:
                    lflags = []
                    if not len(flag) == 0:
                        lflags = [flag]
                    try:
                        objects = compiler.compile(
                            ['simd.c'],
                            output_dir=dname,
                            extra_postargs=lflags + self.__cflags)
                    except CompileError:
                        continue
                    try:
                        compiler.link_shared_lib(objects,
                                                 'a.out',
                                                 output_dir=dname)
                    except (LinkError, TypeError):  # pragma: no cover
                        continue  # pragma: no cover
                    return {'support': True, 'flags': lflags + define}
                return {'support': False, 'flags': []}

    def __get_capabilities(self, compiler):
        log.info("getting compiler simd support")
        self.__capabilities[CCompilerCapabilities.SIMD_SSSE3] = \
            self.__has_simd_support(
                compiler,
                ['', '-mssse3'],
                ['-D__SSSE3__'],
                'tmmintrin.h',
                '__m128i t = _mm_loadu_si128((const __m128i*)argv[0]);'
                't = _mm_shuffle_epi8(t, t);'
                'return _mm_cvtsi128_si32(t);'
        )
        log.info(
            "SSSE3: %s" %
            str(self.__capabilities[
                CCompilerCapabilities.SIMD_SSSE3
            ]['support'])
        )
        self.__capabilities[CCompilerCapabilities.SIMD_SSE41] = \
            self.__has_simd_support(
                compiler,
                ['', '-msse4.1'],
                ['-D__SSE4_1__'],
                'smmintrin.h',
                '__m128i t = _mm_loadu_si128((const __m128i*)argv[0]);'
                't = _mm_mpsadbw_epu8(t, t, 1);'
                'return _mm_cvtsi128_si32(t);'
        )
        log.info(
            "SSE41: %s" %
            str(self.__capabilities[
                CCompilerCapabilities.SIMD_SSE41
            ]['support'])
        )
        self.__capabilities[CCompilerCapabilities.SIMD_SSE42] = \
            self.__has_simd_support(
                compiler,
                ['', '-msse4.2'],
                ['-D__SSE4_2__'],
                'nmmintrin.h',
                '__m128i t = _mm_loadu_si128((const __m128i*)argv[0]);'
                'return _mm_cmpistra(t, t, 0);'
        )
        log.info(
            "SSE42: %s" %
            str(self.__capabilities[
                CCompilerCapabilities.SIMD_SSE42
            ]['support'])
        )
        self.__capabilities[CCompilerCapabilities.SIMD_AVX] = \
            self.__has_simd_support(
                compiler,
                ['', '-mavx', '/arch:AVX'],
                ['-D__AVX__'],
                'immintrin.h',
                '__m256i y = _mm256_loadu_si256((const __m256i*)argv[0]);'
                'return _mm256_testz_si256(y, y);'
        )
        log.info(
            "AVX:   %s" %
            str(self.__capabilities[
                CCompilerCapabilities.SIMD_AVX
            ]['support'])
        )
        self.__capabilities[CCompilerCapabilities.SIMD_AVX2] = \
            self.__has_simd_support(
                compiler,
                ['', '-mavx2', '/arch:AVX2'],
                ['-D__AVX2__'],
                'immintrin.h',
                '__m256i y = _mm256_loadu_si256((const __m256i*)argv[0]);'
                'y = _mm256_i32gather_epi32((int const*)argv[1], y, 2);'
                'return _mm_cvtsi128_si32(_mm256_castsi256_si128(y));'
        )
        log.info(
            "AVX2:  %s" %
            str(self.__capabilities[
                CCompilerCapabilities.SIMD_AVX2
            ]['support'])
        )
        self.__capabilities[CCompilerCapabilities.SIMD_NEON32] = \
            self.__has_simd_support(
                compiler,
                [''],
                [],
                'arm_neon.h',
                'uint8x16_t t = vdupq_n_u8(1);'
                'return vgetq_lane_s32(vreinterpretq_s32_u8(t));'
        )
        log.info(
            "NEON32:  %s" %
            str(self.__capabilities[
                CCompilerCapabilities.SIMD_NEON32
            ]['support'])
        )

    def has(self, what):
        if what not in self.__capabilities:
            return False
        return self.__capabilities[what]['support']

    def flags(self, what):
        if not self.has(what):  # pragma: no branch
            return self.__capabilities[666]['flags']  # pragma: no cover
        return self.__capabilities[what]['flags']
