#!/usr/bin/env python3

import mmap

from PIL import Image, ImageFile

from .filebuf import FileBuffer
from ._libflif import lib as libflif, ffi


# ---------------------------------------------------------
# Image file
def _accept(prefix):
    return prefix[:4] == b"FLIF"


class FLImageFile(ImageFile.ImageFile):
    format = "FLIF"
    format_description = "Free Lossless Image Format"

    def _open(self):
        header = self.fp.read(30)
        info = libflif.flif_read_info_from_memory(header, 30)
        try:
            # Mode/Channels
            channels = libflif.flif_info_get_nb_channels(info)
            if channels == 1:
                if libflif.flif_info_get_depth(info) == 1:
                    self.mode = "1"
                else:
                    self.mode = "L"
            elif channels == 3:
                self.mode = "RGB"
            elif channels == 4:
                self.mode = "RGBA"
            else:
                raise ValueError("Cannot open FLIF image with %i channels" % channels)

            # Size
            self.size = (
                libflif.flif_info_get_width(info),
                libflif.flif_info_get_height(info)
            )
        finally:
            libflif.flif_destroy_info(info)

        self.tile = [("libflif", (0, 0) + self.size, 0, ())]


Image.register_open(FLImageFile.format, FLImageFile, _accept)

Image.register_extension(FLImageFile.format, ".flif")

Image.register_mime(FLImageFile.format, "image/flif")
Image.register_mime(FLImageFile.format, "application/x-flif")


# ---------------------------------------------------------
# Decoder
class LibFLIFDecoder(ImageFile.PyDecoder):
    _pulls_fd = True

    def init(self, args):
        self.context = libflif.flif_create_decoder()

    def _import_image(self, image, bufsize, mode, readinto_fn):
        # FIXME: Don't copy shit around TWICE; Probably want to do it in C
        buf = bytearray(bufsize)
        readinto_fn(image, ffi.from_buffer(buf), bufsize)
        self.set_as_raw(bytes(buf), mode)

    def decode(self, _):
        # Figure out best way to get pointer to data in memory

        with FileBuffer(self.fd) as fb:
            # Decode
            libflif.flif_decoder_decode_memory(self.context, ffi.from_buffer(fb.buffer), fb.size)

        # Get image
        image = libflif.flif_decoder_get_image(self.context, 0)

        # Convert to PIL representation
        # flif_image_read_into_* defined in libflif_build.py
        if self.mode in ("L", "1"):
            self._import_image(image, self.state.xsize * self.state.ysize, "L", libflif.flif_image_read_into_GRAY8)
        elif self.mode in ("RGB",):
            self._import_image(image, self.state.xsize * self.state.ysize * 3 + self.state.xsize, "RGB", libflif.flif_image_read_into_RGB8)
        elif self.mode in ("RGBA",):
            self._import_image(image, self.state.xsize + self.state.ysize * 4, "RGBA", libflif.flif_image_read_into_RGBA8)

        return 0, 0

    def cleanup(self):
        libflif.flif_destroy_decoder(self.context)


Image.register_decoder("libflif", LibFLIFDecoder)


# ---------------------------------------------------------
# Writing
def _save(image, fp, filename):
    lossy = image.encoderinfo.get("lossy", 0)

    # Create FLIF_IMAGE
    # FIXME: Probably should do it in C
    mode = image.mode
    if mode == "L":
        fimage = libflif.flif_import_image_GRAY(image.size[0], image.size[1], image.tobytes(), image.size[0])
    elif mode == "RGB":
        fimage = libflif.flif_import_image_RGB(image.size[0], image.size[1], image.tobytes(), image.size[0] * 3)
    elif mode == "RGBA":
        fimage = libflif.flif_import_image_RGBA(image.size[0], image.size[1], image.tobytes(), image.size[0] * 4)
    else:
        raise IOError("cannot write mode %s as FLIF (yet)" % mode)

    if fimage == ffi.NULL:
        raise RuntimeError("Could not create FLIF_IMAGE")

    # Set up encoder
    context = libflif.flif_create_encoder()

    libflif.flif_encoder_set_lossy(context, lossy)

    libflif.flif_encoder_add_image_move(context, fimage)

    # Encode
    out_buf = ffi.new("void **")
    out_size = ffi.new("size_t *")

    res = libflif.flif_encoder_encode_memory(context, out_buf, out_size)

    fp.write(ffi.buffer(out_buf[0], out_size[0]))

    libflif.flif_destroy_encoder(context)
    libflif.flif_free_memory(out_buf[0])


Image.register_save(FLImageFile.format, _save)