# -*- coding: utf-8 -*-
# Copyright (c) 2017-2023 Richard Hull and contributors
# See LICENSE.rst for details.

import os
import sys
import struct

try:
    import fcntl
    import termios
    import curses
    ASCII_AVAILABLE = True
except ImportError:
    # If running on windows, these package are not available!
    ASCII_AVAILABLE = False
import atexit
import logging
import string
import collections
from io import StringIO
from PIL import Image, ImageFont, ImageDraw

from luma.core.device import device
from luma.core.interface.serial import noop
from luma.emulator.render import transformer
from luma.emulator.clut import rgb2short
from luma.emulator.segment_mapper import regular


logger = logging.getLogger(__name__)

__all__ = ["capture", "gifanim", "pygame"]


class emulator(device):
    """
    Base class for emulated display driver classes
    """
    def __init__(self, width, height, rotate, mode, transform, scale):
        super(emulator, self).__init__(serial_interface=noop())
        try:
            import pygame
        except ImportError:
            raise RuntimeError("Emulator requires pygame to be installed")
        self._pygame = pygame
        self.capabilities(width, height, rotate, mode)
        self.scale = 1 if transform == "none" else scale
        self._transform = getattr(transformer(pygame, width, height, scale),
                                  "none" if scale == 1 else transform)
        self._contrast = 1.0
        self._last_image = None
        self.segment_mapper = regular

    def show(self):
        self.contrast(0xFF)

    def hide(self):
        self.contrast(0x00)

    def contrast(self, value):
        assert 0 <= value <= 255
        self._contrast = value / 255.0
        if self._last_image is not None:
            self.display(self._last_image)

    def cleanup(self):
        pass

    def to_surface(self, image, alpha=1.0):
        """
        Converts a :py:mod:`PIL.Image` into a :class:`pygame.Surface`,
        transforming it according to the ``transform`` and ``scale``
        constructor arguments.
        """
        assert 0.0 <= alpha <= 1.0
        if alpha < 1.0:
            im = image.convert("RGBA")
            black = Image.new(im.mode, im.size, "black")
            im = Image.blend(black, im, alpha)
        else:
            im = image.convert("RGB")

        mode = im.mode
        size = im.size
        data = im.tobytes()
        del im

        surface = self._pygame.image.fromstring(data, size, mode)
        return self._transform(surface)


class capture(emulator):
    """
    Pseudo-device that acts like a physical display, except that it writes the
    image to a numbered PNG file when the :func:`display` method is called.
    Supports 24-bit color depth.
    """
    def __init__(self, width=128, height=64, rotate=0, mode="RGB",
                 transform="scale2x", scale=2, file_template="luma_{0:06}.png",
                 **kwargs):
        super(capture, self).__init__(width, height, rotate, mode, transform, scale)
        self._count = 0
        self._file_template = file_template

    def display(self, image):
        """
        Takes a :py:mod:`PIL.Image` and dumps it to a numbered PNG file.
        """
        assert image.size == self.size
        self._last_image = image

        self._count += 1
        filename = self._file_template.format(self._count)
        image = self.preprocess(image)
        surface = self.to_surface(image, alpha=self._contrast)
        logger.debug(f"Writing: {filename}")
        self._pygame.image.save(surface, filename)


class gifanim(emulator):
    """
    Pseudo-device that acts like a physical display, except that it collects
    the images when the :func:`display` method is called, and on exit,
    assembles them into an animated GIF image. Supports 24-bit color depth,
    albeit with an indexed color palette.
    """
    def __init__(self, width=128, height=64, rotate=0, mode="RGB",
                 transform="scale2x", scale=2, filename="luma_anim.gif",
                 duration=0.01, loop=0, max_frames=None, **kwargs):
        super(gifanim, self).__init__(width, height, rotate, mode, transform, scale)
        self._images = []
        self._count = 0
        self._max_frames = max_frames
        self._filename = filename
        self._loop = loop
        self._duration = int(duration * 1000)
        atexit.register(self.write_animation)

    def display(self, image):
        """
        Takes an image, scales it according to the nominated transform, and
        stores it for later building into an animated GIF.
        """
        assert image.size == self.size
        self._last_image = image

        image = self.preprocess(image)
        surface = self.to_surface(image, alpha=self._contrast)
        rawbytes = self._pygame.image.tostring(surface, "RGB", False)
        im = Image.frombytes("RGB", surface.get_size(), rawbytes)
        self._images.append(im)

        self._count += 1
        logger.debug(f"Recording frame: {self._count}")

        if self._max_frames and self._count >= self._max_frames:
            sys.exit(0)

    def write_animation(self):
        if len(self._images) > 0:
            logger.debug("Please wait... building animated GIF")
            with open(self._filename, "w+b") as fp:
                self._images[0].save(fp, save_all=True, loop=self._loop,
                                     duration=self._duration,
                                     append_images=self._images[1:],
                                     optimize=True, format="GIF")

            file_size = os.stat(self._filename).st_size
            logger.debug(f"Wrote {self._count} frames to file: {self._filename} ({file_size} bytes)")


class pygame(emulator):
    """
    Pseudo-device that acts like a physical display, except that it renders
    to a displayed window. The frame rate is limited to 60FPS (much faster
    than a Raspberry Pi can achieve, but this can be overridden as necessary).
    Supports 24-bit color depth.

    :mod:`pygame` is used to render the emulated display window, and it's
    event loop is checked to see if the ESC key was pressed or the window
    was dismissed: if so :func:`sys.exit()` is called.
    """
    def __init__(self, width=128, height=64, rotate=0, mode="RGB", transform="scale2x",
                 scale=2, frame_rate=60, **kwargs):
        super(pygame, self).__init__(width, height, rotate, mode, transform, scale)
        self._pygame.display.init()
        self._pygame.font.init()
        self._pygame.display.set_caption("Luma.core Display Emulator")
        self._clock = self._pygame.time.Clock()
        self._fps = frame_rate
        self._screen = None

    def _abort(self):
        keystate = self._pygame.key.get_pressed()
        return keystate[self._pygame.K_ESCAPE] or self._pygame.event.peek(self._pygame.QUIT)

    def display(self, image):
        """
        Takes a :py:mod:`PIL.Image` and renders it to a pygame display surface.
        """
        assert image.size == self.size
        self._last_image = image

        image = self.preprocess(image)
        self._clock.tick(self._fps)
        self._pygame.event.pump()

        if self._abort():
            self._pygame.quit()
            sys.exit()

        surface = self.to_surface(image, alpha=self._contrast)
        if self._screen is None:
            self._screen = self._pygame.display.set_mode(surface.get_size())
        self._screen.blit(surface, (0, 0))
        self._pygame.display.flip()


if ASCII_AVAILABLE:
    __all__.extend(["asciiart", "asciiblock"])

    class asciiart(emulator):
        """
        Pseudo-device that acts like a physical display, except that it converts the
        image to display into an ASCII-art representation and downscales colors to
        match the xterm-256 color scheme. Supports 24-bit color depth.

        This device takes hold of the terminal window (using curses), and any output
        for sysout and syserr is captured and stored, and is replayed when the
        cleanup method is called.

        Loosely based on https://github.com/ajalt/pyasciigen/blob/master/asciigen.py

        .. versionadded:: 0.2.0
        """
        def __init__(self, width=128, height=64, rotate=0, mode="RGB", transform="scale2x",
                     scale=2, **kwargs):

            super(asciiart, self).__init__(width, height, rotate, mode, transform, scale)
            self._stdscr = curses.initscr()
            curses.start_color()
            curses.use_default_colors()
            for i in range(0, curses.COLORS):
                curses.init_pair(i, i, -1)
            curses.noecho()
            curses.cbreak()

            # Capture all stdout, stderr
            self._old_stdX = (sys.stdout, sys.stderr)
            self._captured = (StringIO(), StringIO())
            sys.stdout, sys.stderr = self._captured

            # Sort printable characters according to the number of black pixels present.
            # Don't use string.printable, since we don't want any whitespace except spaces.
            charset = (string.ascii_letters + string.digits + string.punctuation + "  ")
            self._chars = list(reversed(sorted(charset, key=self._char_density)))
            self._char_width, self._char_height = ImageFont.load_default().getsize("X")
            self._contrast = 1.0

        def _char_density(self, c, font=ImageFont.load_default()):
            """
            Count the number of black pixels in a rendered character.
            """
            image = Image.new('1', font.getsize(c), color=255)
            draw = ImageDraw.Draw(image)
            draw.text((0, 0), c, fill="white", font=font)
            return collections.Counter(image.getdata())[0]  # 0 is black

        def _generate_art(self, image, width, height):
            """
            Return an iterator that produces the ascii art.
            """
            # Characters aren't square, so scale the output by the aspect ratio of a charater
            height = int(height * self._char_width / float(self._char_height))
            image = image.resize((width, height), Image.ANTIALIAS).convert("RGB")

            for (r, g, b) in image.getdata():
                greyscale = int(0.299 * r + 0.587 * g + 0.114 * b)
                ch = self._chars[int(greyscale / 255. * (len(self._chars) - 1) + 0.5)]
                yield (ch, rgb2short(r, g, b))

        def display(self, image):
            """
            Takes a :py:mod:`PIL.Image` and renders it to the current terminal as
            ASCII-art.
            """
            assert image.size == self.size
            self._last_image = image

            surface = self.to_surface(self.preprocess(image), alpha=self._contrast)
            rawbytes = self._pygame.image.tostring(surface, "RGB", False)
            image = Image.frombytes("RGB", surface.get_size(), rawbytes)

            scr_width = self._stdscr.getmaxyx()[1]
            scale = float(scr_width) / image.width

            self._stdscr.erase()
            self._stdscr.move(0, 0)
            try:
                for (ch, color) in self._generate_art(image, int(image.width * scale), int(image.height * scale)):
                    self._stdscr.addstr(ch, curses.color_pair(color))

            except curses.error:
                # End of screen reached
                pass

            self._stdscr.refresh()

        def cleanup(self):
            super(asciiart, self).cleanup()

            # Stty sane
            curses.nocbreak()
            curses.echo()
            curses.endwin()

            # Restore stdout & stderr, then print out captured
            sys.stdout, sys.stderr = self._old_stdX
            sys.stdout.write(self._captured[0].getvalue())
            sys.stdout.flush()
            sys.stderr.write(self._captured[1].getvalue())
            sys.stderr.flush()

    class asciiblock(emulator):
        """
        Pseudo-device that acts like a physical display, except that it converts
        the image pixels to display into colored ASCII half-blocks (ASCII code 220,
        '▄'), where the upper part background is used for one pixel, and the lower
        part foreground is used for the pixel on the next row. As most terminal
        display characters are in ratio 2:1, the half-block appears square.

        Inspired by `Command Line Curiosities - Making the Terminal Sing by Hamza Haiken <https://www.youtube.com/watch?v=j5zA5Xi_ph8>`__

        .. versionadded:: 1.1.0
        """
        def __init__(self, width=128, height=64, rotate=0, mode="RGB", transform="scale2x",
                     scale=2, **kwargs):

            super(asciiblock, self).__init__(width, height, rotate, mode, transform, scale)
            self._CSI("2J")

        def _terminal_size(self):
            s = struct.pack('HHHH', 0, 0, 0, 0)
            t = fcntl.ioctl(sys.stdout.fileno(), termios.TIOCGWINSZ, s)
            return struct.unpack('HHHH', t)

        def _generate_art(self, image, width, height):
            """
            Return an iterator that produces the ascii art.
            """
            image = image.resize((width, height), Image.LANCZOS).convert("RGB")
            pixels = list(image.getdata())

            for y in range(0, height - 1, 2):
                for x in range(width):
                    i = y * width + x
                    bg = rgb2short(*(pixels[i]))
                    fg = rgb2short(*(pixels[i + width]))
                    yield (fg, bg)

        def _CSI(self, cmd):
            """
            Control sequence introducer
            """
            sys.stdout.write('\x1b[')
            sys.stdout.write(cmd)

        def display(self, image):
            """
            Takes a :py:mod:`PIL.Image` and renders it to the current terminal as
            ASCII-blocks.
            """
            assert image.size == self.size
            self._last_image = image

            surface = self.to_surface(self.preprocess(image), alpha=self._contrast)
            rawbytes = self._pygame.image.tostring(surface, "RGB", False)
            image = Image.frombytes("RGB", surface.get_size(), rawbytes)

            scr_width = self._terminal_size()[1]
            scale = float(scr_width) / image.width

            self._CSI('1;1H')  # Move to top/left

            for (fg, bg) in self._generate_art(image, int(image.width * scale), int(image.height * scale)):
                self._CSI(f'38;5;{fg};48;5;{bg}m')
                sys.stdout.write('▄')

            self._CSI('0m')
            sys.stdout.flush()

        def cleanup(self):
            super(asciiblock, self).cleanup()
            self._CSI('0m')
            self._CSI('2J')
