#!/usr/bin/env python3
# -*- coding: utf-8 -*-

""" Author:                     Hendrik Schutter, localhorst@mosad.xyz
    Date of creation:           2022/01/22
    Date of last modification:  2022/01/25
"""

import os
import sys
import time
import subprocess
import shlex
import shutil
import glob
import PIL
from PIL import Image
from PIL import ImageFont
from PIL import ImageDraw 

temp_dir = os.path.join(os.getcwd(), "_tempCodecVisualizer/")

def create_temp_dir():
    try:
        os.mkdir(temp_dir)
    except OSError:
        print("Unable to create temp dir:  %s" % temp_dir)
        sys.exit()

def delete_temp_dir():
    try:
        shutil.rmtree(temp_dir)
    except OSError:
        pass

def get_length(filename):
    result = subprocess.run(["ffprobe", "-v", "error", "-show_entries",
                             "format=duration", "-of",
                             "default=noprint_wrappers=1:nokey=1", filename],
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT)
    return float(result.stdout)

def get_codec(filename):
    result = subprocess.run(["ffprobe", "-v", "error", "-select_streams", "v:0",
                             "-show_entries", "stream=codec_name",
                             "-of", "default=noprint_wrappers=1:nokey=1", filename],
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT)
    return str(result.stdout.decode("utf-8")).rstrip("\n")

def extract_frame(video_file, time_offset, output_file):
    cmd = "ffmpeg -y -ss " + time.strftime('%H:%M:%S', time.gmtime(time_offset)) + ".00 -i " + str(video_file) + " -frames:v 1 " + str(output_file)
    devnull = open(os.devnull, 'w')
    subprocess.call(shlex.split(cmd),stdout=devnull, stderr=devnull)

def get_font_path():
    path = "/usr/share/fonts"
    files = glob.glob(path + "/**/DejaVuSans-Bold.ttf", recursive = True)
    return files[0]

def zoom_at(img, crop_factor):
    image_width, image_height = img.size
    img = img.crop((int(image_width/2) - int(image_width/crop_factor) / (crop_factor * 2), int(image_height/2) - int(image_height/crop_factor) / (crop_factor * 2), 
                    int(image_width/2) + int(image_width/crop_factor) / (crop_factor * 2), int(image_height/2) + int(image_height/crop_factor) / (crop_factor * 2)))
    return img.resize((int(image_width/crop_factor*2), int(image_height/crop_factor*2)), Image.LANCZOS)

def create_collage(images_A, images_B, statistics,  output_file):
    image_width, image_height = images_A[0].size
    x_offset = int(image_width/15)
    y_offset = int(image_height/1.5)
    output_width = image_width *2 + x_offset
    output_height = image_height *4 + y_offset
    output_image = Image.new('RGB', (output_width, output_height))

    font_file = get_font_path()
    draw = ImageDraw.Draw(output_image)

    text = "codec visualizer"
    draw.text((x_offset, 0), text,(255,255,255),font=ImageFont.truetype(font_file, int(image_width/10)))

    text = "reduced data size: " + str(statistics["compression_rate"]) + "%"
    draw.text((x_offset + image_width , int(image_height/18)), text,(255,255,255),font=ImageFont.truetype(font_file, int(image_width/18)))

    text = str(statistics["codec"][0])
    draw.text((x_offset , int(image_height/18) * 5), text,(255,255,255),font=ImageFont.truetype(font_file, int(image_width/25)))
    text = str(statistics["codec"][1])
    draw.text((x_offset + image_width , int(image_height/18) * 5), text,(255,255,255),font=ImageFont.truetype(font_file, int(image_width/25)))

    text = str(statistics["filename"][0])
    draw.text((x_offset , int(image_height/18) * 7), text,(255,255,255),font=ImageFont.truetype(font_file, int(image_width/25)))
    text = str(statistics["filename"][1])
    draw.text((x_offset + image_width , int(image_height/18) * 7), text,(255,255,255),font=ImageFont.truetype(font_file, int(image_width/25)))

    text = str(statistics["size"][0])
    draw.text((x_offset , int(image_height/18) * 9), text,(255,255,255),font=ImageFont.truetype(font_file, int(image_width/25)))
    text = str(statistics["size"][1])
    draw.text((x_offset + image_width , int(image_height/18) * 9), text,(255,255,255),font=ImageFont.truetype(font_file, int(image_width/25)))

    i = 0
    y = y_offset
    for row in range(4):
            output_image.paste(images_A[i], (int(x_offset/2), y))
            output_image.paste(images_B[i], (int(x_offset/2) + image_width, y))
            i += 1
            y += image_height

    i = 0
    y = y_offset
    crop_factor = 6
    for row in range(4):
            cropped_image = zoom_at(images_A[i], crop_factor)
            cropped_image_with, _ = cropped_image.size
            output_image.paste(cropped_image, (int(x_offset/2) + image_width - cropped_image_with , y))
            cropped_image = zoom_at(images_B[i].transpose(PIL.Image.FLIP_LEFT_RIGHT), crop_factor)
            output_image.paste(cropped_image, (int(x_offset/2) + image_width, y))
            i += 1
            y += image_height

    output_image.save(output_file)

def calc_compression_rate(size_A, size_B):
    if size_A == size_B:
        return 0
    if size_A > size_B:
        return int(100 - 100/size_A*size_B)    
    if size_B > size_A:
        return int(100 - 100/size_B*size_A)

def human_readable_size(size, decimal_places=2):
    for unit in ['B', 'KiB', 'MiB', 'GiB', 'TiB', 'PiB']:
        if size < 1024.0 or unit == 'PiB':
            break
        size /= 1024.0
    return f"{size:.{decimal_places}f} {unit}"

def main() -> None:
    if(len(sys.argv) != 3):
        print("Bad usage!")
        print("Usage: python ./codec_analyzer.py file file")
        sys.exit()

    file_A = sys.argv[1]
    file_B = sys.argv[2]

    if (int(get_length(file_A)) != int(get_length(file_B))):
        print("video files have different lengths")
        sys.exit()

    video_lenght = int(get_length(file_A))

    delete_temp_dir()
    create_temp_dir()
 
    stills_timestamps = []
    stills_file_A = []
    stills_file_B = []

    stills_timestamps.append(int(video_lenght*0.15)) #don't get the intro
    stills_timestamps.append(int(video_lenght*0.33))
    stills_timestamps.append(int(video_lenght*0.50))
    stills_timestamps.append(int(video_lenght*0.75)) #don't get the outro

    for frame_timestamp in stills_timestamps:
        still_file_A = temp_dir +  "A_"+ str(frame_timestamp) + ".tif"
        extract_frame(file_A, frame_timestamp, still_file_A)
        stills_file_A.append(Image.open(still_file_A))

        still_file_B = temp_dir +  "B_"+ str(frame_timestamp) + ".tif"
        extract_frame(file_B, frame_timestamp, still_file_B)
        stills_file_B.append(Image.open(still_file_B))

    file_statistics = {
        "filename": [os.path.basename(file_A),os.path.basename(file_B)],
        "size": [human_readable_size(os.path.getsize(file_A)), human_readable_size(os.path.getsize(file_B))],
        "codec": [get_codec(file_A), get_codec(file_B)],
        "compression_rate": calc_compression_rate(os.path.getsize(file_A), os.path.getsize(file_B))
    }

    create_collage(stills_file_A, stills_file_B, file_statistics, "output.png")

    delete_temp_dir()

if __name__ == "__main__":
    main()
