#!/usr/bin/env python3

import subprocess
import sys
import os
import time
import parse
import argparse
import json
from ansi.colour import fg, fx
from datetime import datetime


class Probe:
    def __init__(self):
        parser = argparse.ArgumentParser(
            description="Readable version of ffprobe. Give filename as argument."
        )
        parser.add_argument(
            "-c", dest="colorize", action="store_true", default=False, help="Colorize"
        )
        parser.add_argument("filename", help="File to probe")
        parsed = parser.parse_args()

        self.filename = parsed.filename
        self.colorize = parsed.colorize
        self.run_probe()
        self.print()

    def run_probe(self):
        process = subprocess.Popen(
            [
                "",
                "-v",
                "error",
                "-hide_banner",
                "-of",
                "default=noprint_wrappers=0",
                "-print_format",
                "json",
                "-show_format",
                "-show_streams",
                self.filename,
            ],
            executable="ffprobe",
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        self.raw_data, errors = process.communicate()
        if len(errors) > 0:
            print(errors)
        self.json_data = json.loads(self.raw_data)

    def print(self):
        # ~ print(self.json_data['format'])
        self.json_data["format"]["time"] = self._timestr(
            self.json_data["format"].get("duration", "0")
        )
        self.json_data["format"]["hrate"] = self._speedfmt(
            self.json_data["format"].get("bit_rate", "0")
        )
        self.json_data["format"]["hsize"] = self._sizefmt(
            float(self.json_data["format"].get("size", 0))
        )
        self.json_data["format"].update(self._get_colors())

        msg = """{Y}==== Format ===={z}
   File: {filename}
 Format: {format_long_name} ({format_name})
 Length: {time}
   Size: {hsize}
Bitrate: {hrate}
Streams: {nb_streams}""".format(
            **self.json_data["format"]
        )
        print(msg)
        for stream in self.json_data["streams"]:
            original = stream.copy()
            stream.update(self._get_colors())

            if stream["codec_type"] == "video":
                msg = self.format_video(stream)
            elif stream["codec_type"] == "audio":
                msg = self.format_audio(stream)
            elif stream["codec_type"] == "subtitle":
                msg = self.format_subtitle(stream)
            else:
                msg = "Unrecognized stream\n{}".format(json.dumps(original, indent=2))

            print(msg)

    def format_subtitle(self, stream):

        stream["tags-lang"] = stream.get("tags", {}).get("language", "NA")
        msg = """{Y}==== Stream:{index} ===={z}
    Type: {G}{codec_type}{z}
   Codec: {codec_long_name} ({codec_name})
Language: {tags-lang}""".format(
            **stream
        )
        return msg

    def format_video(self, stream):
        stream["hrate"] = self._speedfmt(stream.get("bit_rate", 0))
        stream["fps"] = "{:.2f}".format(
            float(stream["r_frame_rate"].split("/")[0])
            / float(stream["r_frame_rate"].split("/")[1])
        )
        stream["display_aspect_ratio"] = stream.get("display_aspect_ratio", "NA")
        msg = """{Y}==== Stream:{index} ===={z}
       Type: {G}{codec_type}{z}
      Codec: {codec_long_name} ({codec_name})
    Bitrate: {hrate}
 Resolution: {width}x{height}
     Aspect: {display_aspect_ratio}
    Profile: {profile}
        FPS: {r_frame_rate} ({fps})""".format(
            **stream
        )

        return msg

    def format_audio(self, stream):
        stream["hrate"] = self._speedfmt(stream.get("bit_rate", 0))
        msg = """{Y}==== Stream:{index} ===={z}
       Type: {G}{codec_type}{z}
      Codec: {codec_long_name} ({codec_name})
    Bitrate: {hrate}
Sample Rate: {sample_rate} Hz
   Channels: {channels} ({channel_layout})""".format(
            **stream
        )

        return msg

    def _timestr(self, sec):
        msec = int((float(sec) - int(float(sec))) * 1000)
        sec = int(float(sec))
        try:
            hours = int(sec) // 3600 % 24
            minutes = int(sec) // 60 % 60
            seconds = int(sec) % 60

            return "{:02d}:{:02d}:{:02d}.{:03d}".format(hours, minutes, seconds, msec)
        except Exception:
            return sec

    def _mbstr(self, b):

        try:
            return "{:.1f}".format((float(b) / (1024 ** 2)))
        except Exception:
            return b

    def _speedfmt(self, b):

        return self._sizefmt(float(b), suffix="bit/s")
        try:
            return "{:.1f}".format((float(b) / (1024 ** 2)))
        except Exception:
            return b

    def _sizefmt(self, num, suffix="B"):

        num = float(num)
        for unit in ["", "K", "M", "G", "T", "P", "E", "Z"]:
            if abs(num) < 1024.0:
                return "%3.1f %s%s" % (num, unit, suffix)
            num /= 1024.0
        return "%.1f %s%s" % (num, "Y", suffix)

    def _get_colors(self):

        if self.colorize:
            return {
                "Y": fg.boldyellow,
                "G": fg.boldgreen,
                "z": fx.reset,
            }
        else:
            return {
                "Y": "",
                "G": "",
                "z": "",
            }


def main():
    Probe()


if __name__ == "__main__":
    main()