#!/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]))
        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()