#!/usr/bin/python3
import subprocess
import pandas as pd
import logging


def run_psnr_check(original, encoded, video_info):
    out = ""
    # bad practice, but idgaf
    # -f rawvideo {video_info}
    options = f"-f rawvideo {video_info} -i {original} -i {encoded} -filter_complex psnr -f null /dev/null"
    with open("ffmpeg-log.txt", "w") as f:
        proc = subprocess.run(["ffmpeg", *options.split()],
                              stdout=f, stderr=subprocess.STDOUT, text=True)
        logging.info(f"Return code: {proc.returncode}")
    with open("ffmpeg-log.txt", "r") as f:
        out = f.read()
    return out


def run_ssim_check(original, encoded, video_info):
    # bad practice, but idgaf
    # -f rawvideo {video_info}
    # we don't  need additional information with h264 encoded files
    options = f"-f rawvideo {video_info} -i {original} -i {encoded} -filter_complex ssim -f null /dev/null"
    with open("ffmpeg-log.txt", "w") as f:
        proc = subprocess.run(["ffmpeg", *options.split()],
                              stdout=f, stderr=subprocess.STDOUT, text=True)
        logging.info(f"Return code: {proc.returncode}")
    with open("ffmpeg-log.txt", "r") as f:
        out = f.read()
    return out


def parse_psnr_output(output):
    for line in output.splitlines():
        if "[Parsed_psnr" in line and "PSNR" in line:
            parts = line.split()
            y = parts[4].split(":")[1]
            u = parts[5].split(":")[1]
            v = parts[6].split(":")[1]
            avg = parts[7].split(":")[1]
            minYUV = parts[8].split(":")[1]
            maxYUV = parts[9].split(":")[1]
            return {
                "Y": y,
                "U": u,
                "V": v,
                "Average": avg,
                "MinYUV": minYUV,
                "MaxYUV": maxYUV
            }
    return {}


def parse_ssim_output(output):
    for line in output.splitlines():
        if "[Parsed_ssim" in line and "SSIM" in line:
            parts = line.split()
            all_value = parts[10].split(":")[1]
            y = parts[4].split(":")[1]
            u = parts[6].split(":")[1]
            v = parts[8].split(":")[1]
            return {
                "Y": y,
                "U": u,
                "V": v,
                "Average": all_value
            }
    return {}


def run_quality_check(original, encoded, option):
    psnr_result = run_psnr_check(original, encoded, option)
    ssim_result = run_ssim_check(original, encoded, option)
    psnr_metrics = parse_psnr_output(psnr_result)
    ssim_metrics = parse_ssim_output(ssim_result)
    logging.info("PSNR Metrics:", psnr_metrics)
    logging.info("SSIM Metrics:", ssim_metrics)
    return psnr_metrics, ssim_metrics


def parse_quality_report(psnr_metrics, ssim_metrics):
    psnrSeries = pd.Series(psnr_metrics)
    ssimSeries = pd.Series(ssim_metrics)
    combined = pd.concat([psnrSeries, ssimSeries], axis=1)
    combined.columns = ["PSNR", "SSIM"]
    combined = combined.fillna(0)
    return combined


if __name__ == "__main__":
    psnr, ssim = run_quality_check(
        "base-x264enc-kpop-test-10.yuv",
        "encoded-x264enc-kpop-test-10.mp4",
        "-pixel_format yuv420p -color_range tv -video_size 1920x1080 -framerate 23.98 "
    )

    combined = parse_quality_report(
        psnr,
        ssim
    )

    encoder = "x264enc"
    profile = "main"
    params = "bitrate=5000"

    columns = pd.MultiIndex.from_tuples(
        [(encoder, profile, params, col) for col in combined.columns]
    )

    combined.columns = columns

    main_df = combined
    profile = "baseline"

    combined2 = parse_quality_report(
        psnr,
        ssim
    )
    columns = pd.MultiIndex.from_tuples(
        [(encoder, profile, params, col) for col in combined2.columns]
    )
    combined2.columns = columns
    main_df = pd.concat([main_df, combined2], axis=1)
    logging.info(main_df)

    main_df.to_csv("quality_report.csv")