from .conversions import pg_video_convert
from .scripts.crop_video import video_crop
from contextlib import contextmanager
from procgraph import Block
from procgraph.block_utils import expand, make_sure_dir_exists, input_check_rgb_or_grayscale
from procgraph.utils import friendly_path, indent
import numpy
import os
import signal
import subprocess
import tempfile


__all__ = ["MEncoder"]


class MEncoder(Block):
    """
        Encodes a video stream using ``mencoder``.

        Note that allowed codec and bitrate depend on
        your version of mencoder.

        MP4 output: currently it works by creating a .AVI with mencoder
        and then converting to MP4 using ffmpeg.

        RGBA videos: not working fully. For now it outputs .AVI with
        mencoder and encoded using png.


    """

    Block.alias("mencoder")

    Block.input(
        "image",
        "Either a HxWx3 uint8 numpy array representing " "an RGB image, or a HxW representing grayscale. ",
    )

    Block.config("file", "Output file (AVI/MP4 format)")
    Block.config(
        "fps",
        "Framerate of resulting movie. If not specified, " "it will be guessed from data.",
        default=None,
    )
    Block.config(
        "fps_safe",
        "If the frame autodetect gives strange results, " "we use this safe value instead.",
        default=10,
    )

    Block.config("quiet", "If True, suppress mencoder's messages", default=True)
    Block.config(
        "timestamps",
        "If True, also writes <file>.timestamps that" " includes a line with the timestamp for" " each frame",
        default=True,
    )

    Block.config("crop", "If true, the video will be " "post-processed and cropped", default=False)

    Block.config("md", "Metadata for the file.", default={})

    Block.config("container", "Which container to use; " "if None, it will be guessed.", default=None)
    Block.config("vcodec", "Codec to use. If None, it will be guessed", default=None)
    Block.config("vcodec_params", "Codec-dependent params.", default={})

    Block.config("firstpass_bitrate", default=3 * 1000 * 1000)

    def init(self):
        self.process = None
        self.buffer = []
        self.image_shape = None  # Shape of image being encoded

        self.first_frame_timestamp = None

        from .programs_existence import check_programs_existence

        programs = ["mencoder", "ffmpeg"]
        self.programs = check_programs_existence(programs=programs)

        for p in self.programs:
            self.info("Using %13s = %s" % (p, self.programs[p]))

    def update(self):
        input_check_rgb_or_grayscale(self, 0)

        # Put image in a buffer -- we don't use it right away
        image = self.get_input(0)

        timestamp = self.get_input_timestamp(0)
        self.buffer.append((timestamp, image))

        if self.process is None:
            self.try_initialization()

        if self.process is not None:
            # initialization was succesful
            while self.buffer:
                timestamp, image = self.buffer.pop(0)

                if self.first_frame_timestamp is None:
                    self.first_frame_timestamp = timestamp

                self.write_value(timestamp, image)

    def try_initialization(self):
        # If we don't have at least two frames, continue
        if len(self.buffer) < 2:
            return

        # Get height and width from first image
        first_image = self.buffer[0][1]

        self.shape = first_image.shape
        self.height = self.shape[0]
        self.width = self.shape[1]

        if self.height > 8192 or self.width > 8192:
            msg = "Mencoder cannot support movies this big (%sx%s)"
            msg = msg % (self.height, self.width)
            raise Exception(msg)
        self.ndim = len(self.shape)

        # Format for mencoder's rawvideo "format" option
        if self.ndim == 2:
            format = "y8"  # @ReservedAssignment
        else:
            if self.shape[2] == 3:
                format = "rgb24"  # @ReservedAssignment
            elif self.shape[2] == 4:
                # Note: did not try this yet
                format = "rgba"  # @ReservedAssignment
                msg = "I detected that you are trying to write a transparent"
                msg += "video. This does not work well yet (and besides,"
                msg += "it is not supported in many applications, like "
                msg += "Keynote). Anyway, the plan is to use mencoder "
                msg += 'to write a .AVI with codec "png". This will fail '
                msg += "for now, but perhaps in the future it will be "
                msg += ".better"
                self.error(msg)

        # guess the fps if we are not given the config
        if self.config.fps is None:
            delta = self.buffer[-1][0] - self.buffer[0][0]

            if delta == 0:
                timestamps = [x[0] for x in self.buffer]
                self.debug("Got 0 delta: timestamps: %s" % timestamps)
                fps = 0
            else:
                fps = (len(self.buffer) - 1) / delta

            # Check for very wrong results
            if not (3 < fps < 60):
                self.error(
                    "Detected fps is %.2f; this seems strange to me,"
                    " so I will use the safe choice fps = %.2f." % (fps, self.config.fps_safe)
                )
                fps = self.config.fps_safe
        else:
            fps = self.config.fps

        # adapt the bitrate to the size of the image
        vbitrate0 = self.config.firstpass_bitrate
        shape0 = (640, 480)

        n1 = self.width * self.height
        n0 = shape0[0] * shape0[1]
        vbitrate = vbitrate0 * n1 / n0

        self.info("Estimated bitrate %r" % vbitrate)

        max_bitrate = 90 * 1000 * 1000
        if vbitrate > max_bitrate:

            vbitrate = max_bitrate
            self.info("Estimated bitrate too high, capping at %s" % vbitrate)

        self.filename = expand(self.config.file)
        if os.path.exists(self.filename):
            self.info("Removing previous version of %s." % friendly_path(self.filename))
            os.unlink(self.filename)

        self.tmp_filename = f"{self.filename}-active.avi"

        make_sure_dir_exists(self.filename)

        self.info(
            "Writing %dx%d %s video stream at %.3f fps to %r."
            % (self.width, self.height, format, fps, friendly_path(self.filename))
        )

        if format == "rgba":
            ovc = ["-ovc", "lavc", "-lavcopts", "vcodec=png"]
        else:
            ovc = ["-ovc", "lavc", "-lavcopts", "vcodec=%s:vbitrate=%d" % ("mpeg4", vbitrate)]

        out = ["-o", self.tmp_filename]
        args = (
            [
                self.programs["mencoder"],
                "/dev/stdin",
                "-demuxer",
                "rawvideo",
                "-rawvideo",
                "w=%d:h=%d:fps=%f:format=%s" % (self.width, self.height, fps, format),
            ]
            + ovc
            + out
        )

        # '-v', "0", # verbosity level (1 prints stats \r)

        # self.debug('$ %s' % " ".join(args))
        # Note: mp4 encoding is currently broken in mencoder :-(
        #       so we have to use ffmpeg as a second step.
        # These would be the options to add:
        # '-of', 'lavf', '-lavfopts', 'format=mp4'

        self.tmp_stdout = tempfile.TemporaryFile()
        self.tmp_stderr = tempfile.TemporaryFile()

        quiet = self.config.quiet
        if quiet:
            # XXX /dev/null not portable
            #            self.debug('stderr: %s' % self.tmp_stderr)
            #            self.debug('stdout: %s' % self.tmp_stdout)
            # TODO: write error
            self.process = subprocess.Popen(
                args,
                preexec_fn=ignore_sigint,
                stdin=subprocess.PIPE,
                stdout=self.tmp_stdout.fileno(),
                stderr=self.tmp_stderr.fileno(),
            )
        else:
            self.process = subprocess.Popen(args=args, stdin=subprocess.PIPE)

        if self.config.timestamps:
            self.timestamps_filename = self.filename + ".timestamps"
            self.timestamps_file = open(self.timestamps_filename, "w")

    def _get_metadata(self):
        """ Returns the user-given metadata as well as some extra created by us. """
        metadata = self.config.md
        # metadata['author'] = 'procgraph'
        # metadata['show'] = 'procgraph'
        # metadata['title'] = 'A video by me.'
        # metadata['comment'] = 'Another video by me.'
        return metadata

    def finish(self):
        if self.process is None:
            msg = "Finish() before starting to encode."
            self.error(msg)
            raise Exception(msg)
            return

        timestamp = self.first_frame_timestamp
        metadata = self._get_metadata()
        container = self.config.container
        vcodec = self.config.vcodec
        vcodec_params = self.config.vcodec_params

        self.info("Transcoding %s" % friendly_path(self.filename))
        pg_video_convert(
            self.tmp_filename,
            self.filename,
            container=container,
            vcodec=vcodec,
            vcodec_params=vcodec_params,
            timestamp=timestamp,
            metadata=metadata,
        )

        #         if False:  # XXX
        if os.path.exists(self.tmp_filename):
            os.unlink(self.tmp_filename)

        if True:
            T = self.first_frame_timestamp
            os.utime(self.filename, (T, T))

        self.info("Finished %s" % friendly_path(self.filename))

        # TODO: skip mp4
        if self.config.crop:
            base, ext = os.path.splitext(self.filename)
            cropped = "%s-crop%s" % (base, ext)
            video_crop(self.filename, cropped)
            os.rename(cropped, self.filename)

    def cleanup(self):
        # TODO: remove timestamps
        #         if 'timestamps_filename' in self.__dict__:
        #             if os.path.exists(self.timestamps_filename):
        #                 os.unlink(self.timestamps_filename)

        if "tmp_filename" in self.__dict__:
            #             if False:  # XXX
            if os.path.exists(self.tmp_filename):
                os.unlink(self.tmp_filename)

        self.cleanup_mencoder()

    def cleanup_mencoder(self):
        # Try to cleanup as well as possible
        if (not "process" in self.__dict__) or self.process is None:
            return

        self.process.stdin.close()

        try:
            if False:
                with timeout(5):
                    self.process.terminate()
                    self.process.wait()
            else:
                self.process.terminate()
                self.process.wait()
        except Timeout as e:
            self.error(e)
        except (OSError, AttributeError):
            # Exception AttributeError: AttributeError("'NoneType' object
            # has no attribute 'SIGTERM'",)
            # in <bound method RangefinderUniform.__del__
            # of RangefinderUniform> ignored
            # http://stackoverflow.com/questions/2572172/
            pass

    #            self.error('Error while cleanup: %s' % e)

    def write_value(self, timestamp, image):
        if self.image_shape is None:
            self.image_shape = image.shape

        if self.image_shape != image.shape:
            msg = "The image has changed shape, from %s to %s." % (self.image_shape, image.shape)
            raise Exception(msg)  # TODO: badinput

        # very important! make sure we are using a reasonable array
        if not image.flags["C_CONTIGUOUS"]:
            image = numpy.ascontiguousarray(image)  # @UndefinedVariable

        try:
            if False:
                with timeout(5):
                    self.process.stdin.write(image.data)
                    self.process.stdin.flush()
            else:
                self.process.stdin.write(image.data)
                self.process.stdin.flush()
        except (Exception, KeyboardInterrupt) as e:
            # IOError = broken pipe
            msg = "Could not write data to mencoder: %s." % e
            #            self.error(msg)
            #            msg += '\n' + indent(self.process.stdout.read(), 'stdout> ')
            #            msg += '\n' + indent(self.process.stderr.read(), 'stderr> ')

            def read_all(f):
                os.lseek(f.fileno(), 0, 0)
                return f.read()

            stderr = read_all(self.tmp_stderr)
            stdout = read_all(self.tmp_stdout)

            msg += indent("stdout>", stderr)
            msg += indent("stderr>", stdout)

            raise Exception(msg)

        if self.config.timestamps:
            self.timestamps_file.write("%.4f\n" % timestamp)
            self.timestamps_file.flush()


class Timeout(Exception):
    pass


@contextmanager
def timeout(delta=5):
    def handler(signum, frame):
        raise Timeout("operation took more than %s seconds" % delta)

    # Set the signal handler and a 5-second alarm
    signal.signal(signal.SIGALRM, handler)
    signal.alarm(delta)
    try:
        yield
    finally:
        signal.alarm(0)


def ignore_sigint():
    signal.signal(signal.SIGINT, signal.SIG_IGN)
