from __future__ import with_statement
from rpython.rtyper.tool import rffi_platform as platform
from rpython.rtyper.lltypesystem import rffi
from rpython.rtyper.lltypesystem import lltype
from pypy.interpreter.error import OperationError, oefmt
from pypy.interpreter.baseobjspace import W_Root
from pypy.interpreter.typedef import TypeDef, interp_attrproperty
from pypy.interpreter.typedef import GetSetProperty
from pypy.interpreter.gateway import interp2app, unwrap_spec
from rpython.translator.tool.cbuild import ExternalCompilationInfo
from rpython.translator.platform import platform as compiler
from rpython.rlib.rarithmetic import intmask, r_longlong
from rpython.rlib import rgil
import sys


if compiler.name == "msvc":
    libname = 'libbz2'
else:
    libname = 'bz2'
eci = ExternalCompilationInfo(
    includes = ['stdio.h', 'sys/types.h', 'bzlib.h'],
    libraries = [libname],
    )
eci = platform.configure_external_library(
    'bz2', eci,
    [dict(prefix='bzip2-')])
if not eci:
    raise ImportError("Could not find bzip2 library")

class CConfig:
    _compilation_info_ = eci
    calling_conv = 'c'

    CHECK_LIBRARY = platform.Has('dump("x", (long)&BZ2_bzCompress)')

    off_t = platform.SimpleType("off_t", rffi.LONGLONG)
    size_t = platform.SimpleType("size_t", rffi.ULONG)
    BUFSIZ = platform.ConstantInteger("BUFSIZ")
    _alloc_type = lltype.FuncType([rffi.VOIDP, rffi.INT, rffi.INT], rffi.VOIDP)
    _free_type = lltype.FuncType([rffi.VOIDP, rffi.VOIDP], lltype.Void)
    SEEK_SET = platform.ConstantInteger("SEEK_SET")
    bz_stream = platform.Struct('bz_stream',
                                [('next_in', rffi.CCHARP),
                                 ('avail_in', rffi.UINT),
                                 ('total_in_lo32', rffi.UINT),
                                 ('total_in_hi32', rffi.UINT),
                                 ('next_out', rffi.CCHARP),
                                 ('avail_out', rffi.UINT),
                                 ('total_out_lo32', rffi.UINT),
                                 ('total_out_hi32', rffi.UINT),
                                 ('state', rffi.VOIDP),
                                 ('bzalloc', lltype.Ptr(_alloc_type)),
                                 ('bzfree', lltype.Ptr(_free_type)),
                                 ('opaque', rffi.VOIDP),
                                 ])

FILE = rffi.COpaquePtr('FILE')
BZFILE = rffi.COpaquePtr('BZFILE')


constants = {}
constant_names = ['BZ_RUN', 'BZ_FLUSH', 'BZ_FINISH', 'BZ_OK',
    'BZ_RUN_OK', 'BZ_FLUSH_OK', 'BZ_FINISH_OK', 'BZ_STREAM_END',
    'BZ_SEQUENCE_ERROR', 'BZ_PARAM_ERROR', 'BZ_MEM_ERROR', 'BZ_DATA_ERROR',
    'BZ_DATA_ERROR_MAGIC', 'BZ_IO_ERROR', 'BZ_UNEXPECTED_EOF',
    'BZ_OUTBUFF_FULL', 'BZ_CONFIG_ERROR']
for name in constant_names:
    setattr(CConfig, name, platform.DefinedConstantInteger(name))

class cConfig(object):
    pass
for k, v in platform.configure(CConfig).items():
    setattr(cConfig, k, v)
if not cConfig.CHECK_LIBRARY:
    raise ImportError("Invalid bz2 library")

for name in constant_names:
    value = getattr(cConfig, name)
    if value is not None:
        constants[name] = value
locals().update(constants)

off_t = cConfig.off_t
bz_stream = lltype.Ptr(cConfig.bz_stream)
BUFSIZ = cConfig.BUFSIZ
SEEK_SET = cConfig.SEEK_SET
BZ_OK = cConfig.BZ_OK
BZ_STREAM_END = cConfig.BZ_STREAM_END
BZ_CONFIG_ERROR = cConfig.BZ_CONFIG_ERROR
BZ_PARAM_ERROR = cConfig.BZ_PARAM_ERROR
BZ_DATA_ERROR = cConfig.BZ_DATA_ERROR
BZ_DATA_ERROR_MAGIC = cConfig.BZ_DATA_ERROR_MAGIC
BZ_IO_ERROR = cConfig.BZ_IO_ERROR
BZ_MEM_ERROR = cConfig.BZ_MEM_ERROR
BZ_UNEXPECTED_EOF = cConfig.BZ_UNEXPECTED_EOF
BZ_SEQUENCE_ERROR = cConfig.BZ_SEQUENCE_ERROR

if BUFSIZ < 8192:
    INITIAL_BUFFER_SIZE = 8192
else:
    INITIAL_BUFFER_SIZE = 8192

UINT_MAX = 2**32-1
MAX_BUFSIZE = min(sys.maxint, UINT_MAX)
assert isinstance(MAX_BUFSIZE, int)

if rffi.sizeof(rffi.INT) > 4:
    BIGCHUNK = 512 * 32
else:
    BIGCHUNK = 512 * 1024

if BZ_CONFIG_ERROR:
    if rffi.sizeof(rffi.LONG) >= 8:
        def _bzs_total_out(bzs):
            return (rffi.getintfield(bzs, 'c_total_out_hi32') << 32) + \
                   rffi.getintfield(bzs, 'c_total_out_lo32')
    else:
        # we can't return a long long value from here, because most
        # callers wouldn't be able to handle it anyway
        def _bzs_total_out(bzs):
            if rffi.getintfield(bzs, 'c_total_out_hi32') != 0 or \
                   rffi.getintfield(bzs, 'c_total_out_lo32') > sys.maxint:
                raise MemoryError
            return rffi.getintfield(bzs, 'c_total_out_lo32')
else:
    XXX    # this case needs fixing (old bz2 library?)
    def _bzs_total_out(bzs):
        return bzs.total_out

def external(name, args, result, **kwds):
    return rffi.llexternal(name, args, result, compilation_info=
                           CConfig._compilation_info_, **kwds)

# the least but one parameter should be rffi.VOIDP but it's not used
# so I trick the compiler to not complain about constanst pointer passed
# to void* arg
BZ2_bzReadOpen = external('BZ2_bzReadOpen', [rffi.INTP, FILE, rffi.INT,
    rffi.INT, rffi.INTP, rffi.INT], BZFILE)
BZ2_bzWriteOpen = external('BZ2_bzWriteOpen', [rffi.INTP, FILE, rffi.INT,
    rffi.INT, rffi.INT], BZFILE)
BZ2_bzReadClose = external('BZ2_bzReadClose', [rffi.INTP, BZFILE], lltype.Void)
BZ2_bzWriteClose = external('BZ2_bzWriteClose', [rffi.INTP, BZFILE,
    rffi.INT, rffi.UINTP, rffi.UINTP], lltype.Void)
BZ2_bzRead = external('BZ2_bzRead', [rffi.INTP, BZFILE, rffi.CCHARP, rffi.INT],
                      rffi.INT)
BZ2_bzWrite = external('BZ2_bzWrite', [rffi.INTP, BZFILE, rffi.CCHARP,
                                       rffi.INT], lltype.Void)
BZ2_bzCompressInit = external('BZ2_bzCompressInit', [bz_stream, rffi.INT,
                              rffi.INT, rffi.INT], rffi.INT)
BZ2_bzCompressEnd = external('BZ2_bzCompressEnd', [bz_stream], rffi.INT,
                             releasegil=False)
BZ2_bzCompress = external('BZ2_bzCompress', [bz_stream, rffi.INT], rffi.INT)
BZ2_bzDecompressInit = external('BZ2_bzDecompressInit', [bz_stream, rffi.INT,
                                                         rffi.INT], rffi.INT)
BZ2_bzDecompressEnd = external('BZ2_bzDecompressEnd', [bz_stream], rffi.INT,
                               releasegil=False)
BZ2_bzDecompress = external('BZ2_bzDecompress', [bz_stream], rffi.INT)

def _catch_bz2_error(space, bzerror):
    if BZ_CONFIG_ERROR and bzerror == BZ_CONFIG_ERROR:
        raise oefmt(space.w_SystemError,
                    "the bz2 library was not compiled correctly")
    if bzerror == BZ_PARAM_ERROR:
        raise oefmt(space.w_SystemError,
                    "the bz2 library has received wrong parameters")
    elif bzerror == BZ_MEM_ERROR:
        raise OperationError(space.w_MemoryError, space.w_None)
    elif bzerror in (BZ_DATA_ERROR, BZ_DATA_ERROR_MAGIC):
        raise oefmt(space.w_IOError, "invalid data stream")
    elif bzerror == BZ_IO_ERROR:
        raise oefmt(space.w_IOError, "unknown IO error")
    elif bzerror == BZ_UNEXPECTED_EOF:
        raise oefmt(space.w_EOFError,
                    "compressed file ended before the logical end-of-stream "
                    "was detected")
    elif bzerror == BZ_SEQUENCE_ERROR:
        raise oefmt(space.w_RuntimeError,
                    "wrong sequence of bz2 library commands used")

def _new_buffer_size(current_size):
    # keep doubling until we reach BIGCHUNK; then the buffer size is no
    # longer increased
    if current_size < BIGCHUNK:
        return current_size + current_size
    return current_size

# ____________________________________________________________

class OutBuffer(object):
    """Handler for the output buffer.  A bit custom code trying to
    encapsulate the logic of setting up the fields of 'bzs' and
    allocating raw memory as needed.
    """
    def __init__(self, bzs, initial_size=INITIAL_BUFFER_SIZE, max_length=-1):
        # when the constructor is called, allocate a piece of memory
        # of length 'piece_size' and make bzs ready to dump there.
        self.temp = []
        self.bzs = bzs
        self.max_length = max_length
        if max_length < 0 or max_length >= initial_size:
            size = initial_size
        else:
            size = max_length
        self._allocate_chunk(size)
        self.left = 0

    def get_data_size(self):
        curr_out = self.current_size - rffi.getintfield(self.bzs, 'c_avail_out')
        total_size = curr_out
        for s in self.temp:
            total_size += len(s)
        return total_size

    def _allocate_chunk(self, size):
        self.raw_buf, self.gc_buf, self.case_num = rffi.alloc_buffer(size)
        self.current_size = size
        self.bzs.c_next_out = self.raw_buf
        rffi.setintfield(self.bzs, 'c_avail_out', size)

    def _get_chunk(self, chunksize):
        assert 0 <= chunksize <= self.current_size
        raw_buf = self.raw_buf
        gc_buf = self.gc_buf
        case_num = self.case_num
        s = rffi.str_from_buffer(raw_buf, gc_buf, case_num,
                                 self.current_size, chunksize)
        rffi.keep_buffer_alive_until_here(raw_buf, gc_buf, case_num)
        self.current_size = 0
        return s

    def prepare_next_chunk(self):
        size = self.current_size
        self.temp.append(self._get_chunk(size))
        newsize = size
        if self.max_length == -1:
            newsize = _new_buffer_size(size)
        else:
            newsize = min(newsize, self.max_length - self.get_data_size())
        self._allocate_chunk(newsize)

    def make_result_string(self):
        count_unoccupied = rffi.getintfield(self.bzs, 'c_avail_out')
        s = self._get_chunk(self.current_size - count_unoccupied)
        if self.temp:
            self.temp.append(s)
            return ''.join(self.temp)
        else:
            return s

    def free(self):
        if self.current_size > 0:
            rffi.keep_buffer_alive_until_here(self.raw_buf, self.gc_buf,
                                              self.case_num)

    def __enter__(self):
        return self
    def __exit__(self, *args):
        self.free()


@unwrap_spec(compresslevel=int)
def descr_compressor__new__(space, w_subtype, compresslevel=9):
    x = space.allocate_instance(W_BZ2Compressor, w_subtype)
    W_BZ2Compressor.__init__(x, space, compresslevel)
    return x

class W_BZ2Compressor(W_Root):
    """BZ2Compressor([compresslevel=9]) -> compressor object

    Create a new compressor object. This object may be used to compress
    data sequentially. If you want to compress data in one shot, use the
    compress() function instead. The compresslevel parameter, if given,
    must be a number between 1 and 9."""
    def __init__(self, space, compresslevel):
        self.space = space
        self._lock = space.allocate_lock()
        self.bzs = lltype.malloc(bz_stream.TO, flavor='raw', zero=True)
        try:
            self.running = False
            self._init_bz2comp(compresslevel)
        except:
            lltype.free(self.bzs, flavor='raw')
            self.bzs = lltype.nullptr(bz_stream.TO)
            raise
        self.register_finalizer(space)

    def lock(self):
        if not self._lock.acquire(False):
            rgil.release()
            self._lock.acquire(True)
            rgil.acquire()

    def unlock(self):
        self._lock.release()

    def _init_bz2comp(self, compresslevel):
        if compresslevel < 1 or compresslevel > 9:
            raise oefmt(self.space.w_ValueError,
                        "compresslevel must be between 1 and 9")

        bzerror = intmask(BZ2_bzCompressInit(self.bzs, compresslevel, 0, 0))
        if bzerror != BZ_OK:
            _catch_bz2_error(self.space, bzerror)

        self.running = True

    def _finalize_(self):
        bzs = self.bzs
        if bzs:
            self.bzs = lltype.nullptr(bz_stream.TO)
            BZ2_bzCompressEnd(bzs)
            lltype.free(bzs, flavor='raw')

    def descr_getstate(self):
        raise oefmt(self.space.w_TypeError, "cannot serialize '%T' object", self)

    @unwrap_spec(data='bufferstr')
    def compress(self, data):
        """compress(data) -> string

        Provide more data to the compressor object. It will return chunks of
        compressed data whenever possible. When you've finished providing data
        to compress, call the flush() method to finish the compression process,
        and return what is left in the internal buffers."""

        assert data is not None
        try:
            self.lock()
            datasize = len(data)

            if datasize == 0:
                return self.space.newbytes("")

            if not self.running:
                raise oefmt(self.space.w_ValueError,
                            "this object was already flushed")

            in_bufsize = datasize

            with OutBuffer(self.bzs) as out:
                with rffi.scoped_nonmovingbuffer(data) as in_buf:

                    self.bzs.c_next_in = in_buf
                    rffi.setintfield(self.bzs, 'c_avail_in', in_bufsize)

                    while True:
                        bzerror = BZ2_bzCompress(self.bzs, BZ_RUN)
                        if bzerror != BZ_RUN_OK:
                            _catch_bz2_error(self.space, bzerror)

                        if rffi.getintfield(self.bzs, 'c_avail_in') == 0:
                            break
                        elif rffi.getintfield(self.bzs, 'c_avail_out') == 0:
                            out.prepare_next_chunk()

                    res = out.make_result_string()
                    return self.space.newbytes(res)
        finally:
            self.unlock()

    def flush(self):
        if not self.running:
            raise oefmt(self.space.w_ValueError,
                        "this object was already flushed")
        self.running = False

        try:
            self.lock()
            with OutBuffer(self.bzs) as out:
                while True:
                    bzerror = BZ2_bzCompress(self.bzs, BZ_FINISH)
                    if bzerror == BZ_STREAM_END:
                        break
                    elif bzerror != BZ_FINISH_OK:
                        _catch_bz2_error(self.space, bzerror)

                    if rffi.getintfield(self.bzs, 'c_avail_out') == 0:
                        out.prepare_next_chunk()

                res = out.make_result_string()
                return self.space.newbytes(res)
        finally:
            self.unlock()

W_BZ2Compressor.typedef = TypeDef("_bz2.BZ2Compressor",
    __doc__ = W_BZ2Compressor.__doc__,
    __new__ = interp2app(descr_compressor__new__),
    __getstate__ = interp2app(W_BZ2Compressor.descr_getstate),
    compress = interp2app(W_BZ2Compressor.compress),
    flush = interp2app(W_BZ2Compressor.flush),
)
W_BZ2Compressor.typedef.acceptable_as_base_class = False

def descr_decompressor__new__(space, w_subtype):
    x = space.allocate_instance(W_BZ2Decompressor, w_subtype)
    W_BZ2Decompressor.__init__(x, space)
    return x

class W_BZ2Decompressor(W_Root):
    """BZ2Decompressor() -> decompressor object

    Create a new decompressor object. This object may be used to decompress
    data sequentially. If you want to decompress data in one shot, use the
    decompress() function instead."""

    def __init__(self, space):
        self.space = space
        self._lock = space.allocate_lock()

        self.bzs = lltype.malloc(bz_stream.TO, flavor='raw', zero=True)
        try:
            self.running = False
            self.unused_data = ""
            self.needs_input = True
            self.input_buffer = ""
            self.left_to_process = 0

            self._init_bz2decomp()
        except:
            lltype.free(self.bzs, flavor='raw')
            self.bzs = lltype.nullptr(bz_stream.TO)
            raise
        self.register_finalizer(space)

    def lock(self):
        if not self._lock.acquire(False):
            rgil.release()
            self._lock.acquire(True)
            rgil.acquire()

    def unlock(self):
        self._lock.release()

    def _init_bz2decomp(self):
        bzerror = BZ2_bzDecompressInit(self.bzs, 0, 0)
        if bzerror != BZ_OK:
            _catch_bz2_error(self.space, bzerror)

        self.running = True

    def _finalize_(self):
        bzs = self.bzs
        if bzs:
            self.bzs = lltype.nullptr(bz_stream.TO)
            BZ2_bzDecompressEnd(bzs)
            lltype.free(bzs, flavor='raw')

    def descr_getstate(self):
        raise oefmt(self.space.w_TypeError, "cannot serialize '%T' object", self)

    def needs_input_w(self, space):
        """ True if more input is needed before more decompressed
            data can be produced. """
        return space.newbool(self.needs_input)

    def eof_w(self, space):
        if self.running:
            return space.w_False
        else:
            return space.w_True

    def _decompress_buf(self, data, max_length):
        total_in = len(data)
        in_bufsize = min(total_in, MAX_BUFSIZE)
        total_in -= in_bufsize
        with rffi.scoped_nonmovingbuffer(data) as in_buf:
            # setup the input and the size it can consume
            self.bzs.c_next_in = in_buf
            rffi.setintfield(self.bzs, 'c_avail_in', in_bufsize)
            self.left_to_process = in_bufsize

            with OutBuffer(self.bzs, max_length=max_length) as out:
                while True:
                    bzreturn = BZ2_bzDecompress(self.bzs)
                    # add up the size that has not been processed
                    avail_in = rffi.getintfield(self.bzs, 'c_avail_in')
                    self.left_to_process = avail_in
                    if bzreturn == BZ_STREAM_END:
                        self.running = False
                        break
                    if bzreturn != BZ_OK:
                        _catch_bz2_error(self.space, bzreturn)

                    if self.left_to_process == 0:
                        break
                    elif rffi.getintfield(self.bzs, 'c_avail_out') == 0:
                        if out.get_data_size() == max_length:
                            break
                        out.prepare_next_chunk()

                self.left_to_process += total_in
                res = out.make_result_string()
                return self.space.newbytes(res)

    @unwrap_spec(data='bufferstr', max_length=int)
    def decompress(self, data, max_length=-1):
        """decompress(data, max_length=-1) -> bytes

        Provide more data to the decompressor object. It will return chunks
        of decompressed data whenever possible. If you try to decompress data
        after the end of stream is found, EOFError will be raised. If any data
        was found after the end of stream, it'll be ignored and saved in
        unused_data attribute."""

        try:
            self.lock()
            if not self.running:
                raise oefmt(self.space.w_EOFError,
                            "end of stream was already found")
            datalen = len(data)
            if len(self.input_buffer) > 0:
                data = self.input_buffer + data
                datalen = len(data)
                self.input_buffer = ""

            result = self._decompress_buf(data, max_length)

            if not self.running: # eq. with eof == Ture
                self.needs_input = False
                if self.left_to_process != 0:
                    start = datalen - self.left_to_process
                    assert start > 0
                    self.unused_data = data[start:]
                    self.left_to_process = 0
            elif self.left_to_process == 0:
                self.input_buffer = ""
                self.needs_input = True
            else:
                self.needs_input = False
                if self.left_to_process > 0:
                    start = datalen-self.left_to_process
                    assert start >= 0
                    self.input_buffer = data[start:]
            return result
        finally:
            self.unlock()





W_BZ2Decompressor.typedef = TypeDef("_bz2.BZ2Decompressor",
    __doc__ = W_BZ2Decompressor.__doc__,
    __new__ = interp2app(descr_decompressor__new__),
    __getstate__ = interp2app(W_BZ2Decompressor.descr_getstate),
    unused_data = interp_attrproperty("unused_data", W_BZ2Decompressor,
        wrapfn="newbytes"),
    eof = GetSetProperty(W_BZ2Decompressor.eof_w),
    decompress = interp2app(W_BZ2Decompressor.decompress),
    needs_input = GetSetProperty(W_BZ2Decompressor.needs_input_w),
)
W_BZ2Decompressor.typedef.acceptable_as_base_class = False