#!/usr/bin/python
# Filename: ffprobe.py
"""
Based on Python wrapper for ffprobe command line tool. ffprobe must exist in the path.
Author: <NAME>

"""

version = "0.5"

import os
import platform
import re
import subprocess


class FFProbe:
    """
    FFProbe wraps the ffprobe command and pulls the data into an object form::
        metadata=FFProbe('multimedia-file.mov')
    """

    def __init__(self, video_file):
        self.video_file = video_file
        try:
            with open(os.devnull, "w") as tempf:
                subprocess.check_call(["ffprobe", "-h"], stdout=tempf, stderr=tempf)
        except:
            raise IOError("ffprobe not found.")
        if os.path.isfile(video_file):
            video_file = self.video_file.replace("\\", "\\\\")

            if str(platform.system()) == "Windows":
                cmd = ["ffprobe", "-show_streams", '"' + video_file + '"']
            else:
                cmd = ["ffprobe -show_streams " + '"' + video_file + '"']

            p = subprocess.Popen(
                cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True
            )
            self.format = None
            self.created = None
            self.duration = None
            self.start = None
            self.bitrate = None
            self.creation_time = None
            self.streams = []
            self.video = []
            self.audio = []
            datalines = []

            for b in iter(p.stdout.readline, b""):
                a = b.decode("utf-8")
                if re.match("\[STREAM\]", a):
                    datalines = []
                elif re.match("\[\/STREAM\]", a):
                    self.streams.append(FFStream(datalines))
                    datalines = []
                else:
                    datalines.append(a)
            for b in iter(p.stderr.readline, b""):
                a = b.decode("utf-8")
                if re.match("\[STREAM\]", a):
                    datalines = []
                elif re.match("\[\/STREAM\]", a):
                    self.streams.append(FFStream(datalines))
                    datalines = []
                else:
                    datalines.append(a)
            for a in self.streams:
                if a.isAudio():
                    self.audio.append(a)
                if a.isVideo():
                    self.video.append(a)
        else:
            raise IOError("No such media file " + video_file)


class FFStream:
    """
    An object representation of an individual stream in a multimedia file.
    """

    def __init__(self, datalines):
        for a in datalines:
            if re.match(r"^.+=.+$", a) is None:
                print(f"Warning: detected incorrect stream metadata line format: {a}")
            else:
                (key, val) = a.strip().split("=")
                key = key.lstrip("TAG:")
                self.__dict__[key] = val

    def isAudio(self):
        """
        Is this stream labelled as an audio stream?
        """
        val = False
        if self.__dict__["codec_type"]:
            if str(self.__dict__["codec_type"]) == "audio":
                val = True
        return val

    def isVideo(self):
        """
        Is the stream labelled as a video stream.
        """
        val = False
        if self.__dict__["codec_type"]:
            if self.codec_type == "video":
                val = True
        return val

    def isSubtitle(self):
        """
        Is the stream labelled as a subtitle stream.
        """
        val = False
        if self.__dict__["codec_type"]:
            if str(self.codec_type) == "subtitle":
                val = True
        return val

    def frameSize(self):
        """
        Returns the pixel frame size as an integer tuple (width,height) if the stream is a video stream.
        Returns None if it is not a video stream.
        """
        size = None
        if self.isVideo():
            if self.__dict__["width"] and self.__dict__["height"]:
                try:
                    size = (int(self.__dict__["width"]), int(self.__dict__["height"]))
                except Exception as e:
                    print(
                        f"None integer size {str(self.__dict__['width'])}:{str(+self.__dict__['height'])}"
                    )
                    size = (0, 0)
        return size

    def pixelFormat(self):
        """
        Returns a string representing the pixel format of the video stream. e.g. yuv420p.
        Returns none is it is not a video stream.
        """
        f = None
        if self.isVideo():
            if self.__dict__["pix_fmt"]:
                f = self.__dict__["pix_fmt"]
        return f

    def frames(self):
        """
        Returns the length of a video stream in frames. Returns 0 if not a video stream.
        """
        f = 0
        if self.isVideo() or self.isAudio():
            if self.__dict__["nb_frames"]:
                try:
                    f = int(self.__dict__["nb_frames"])
                except Exception as e:
                    print("None integer frame count")
        return f

    def durationSeconds(self):
        """
        Returns the runtime duration of the video stream as a floating point number of seconds.
        Returns 0.0 if not a video stream.
        """
        f = 0.0
        if self.isVideo() or self.isAudio():
            if self.__dict__["duration"]:
                try:
                    f = float(self.__dict__["duration"])
                except Exception as e:
                    print("None numeric duration")
        return f

    def language(self):
        """
        Returns language tag of stream. e.g. eng
        """
        lang = None
        if self.__dict__["TAG:language"]:
            lang = self.__dict__["TAG:language"]
        return lang

    def codec(self):
        """
        Returns a string representation of the stream codec.
        """
        codec_name = None
        if self.__dict__["codec_name"]:
            codec_name = self.__dict__["codec_name"]
        return codec_name

    def codecDescription(self):
        """
        Returns a long representation of the stream codec.
        """
        codec_d = None
        if self.__dict__["codec_long_name"]:
            codec_d = self.__dict__["codec_long_name"]
        return codec_d

    def codecTag(self):
        """
        Returns a short representative tag of the stream codec.
        """
        codec_t = None
        if self.__dict__["codec_tag_string"]:
            codec_t = self.__dict__["codec_tag_string"]
        return codec_t

    def bitrate(self):
        """
        Returns bitrate as an integer in bps
        """
        b = 0
        if self.__dict__["bit_rate"]:
            try:
                b = int(self.__dict__["bit_rate"])
            except Exception as e:
                print("None integer bitrate")
        return b


if __name__ == "__main__":
    print("Module ffprobe")
