import subprocess
import os
import time
import re
import csv

temp_cmd = ["sensors", "coretemp-isa-0000"]
power_cmd = ["turbostat", "--quiet", "--hide", "CorWatt,Die,Core,CPU,Busy%,Bzy_MHz,TSC_MHz,IPC,IRQ,SMI,POLL,C1,C2,C3,POLL%,C1%,C2%,C3%,C1ACPI,C2ACPI,C3ACPI,C1ACPI%,C2ACPI%,C3ACPI%,CPU%c1,CPU%c6,CPU%c7,CoreTmp,CoreThr,PkgTmp,Totl%C0,Any%C0,GFX%C0,CPUGFX%,Pkg%pc2,Pkg%pc3,Pkg%pc6,Pkg%pc8,Pk%pc10,CPU%LPI,SYS%LPI,GFXWatt,RAMWatt,PKG_%,RAM_%,UncMHz",
              "-i", "0.001", "--num_iterations", "1", "--Summary"]
csv_fieldnames = ["AVX512", "Resolution", "Encoder", "Preset", "Run Time", "Temperature", "MHz", "Watts"]
cur_dir = os.getcwd()

def ffmpeg_cmd(resolution: str, encoder: str, preset: str, avx: str = "") -> None:
    cmd = f"ssh -t knowledge@nazrin \"ffmpeg -i '{cur_dir}/{resolution}_input.mkv' -c:v {encoder} -preset {preset} {avx} /tmp/output.mkv -y\""
    print(cmd)
    num_tests = 5
    run_time_accumulation = 0
    temp_accumulation = 0
    mhz_accumulation = 0
    watt_accumulation = 0
    
    stat_track_iterations = 0
    for x in range(num_tests):
        start_time = time.time()
        with subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, shell=True) as process:
                while not process.stdout.readline().startswith("frame="):
                    continue
                while process.poll() is None:
                    stdout_line = process.stdout.readline().split()
                    temperature = subprocess.run(temp_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
                    temp_reading = re.search(r"[0-9][0-9][.][0-9]",temperature.stdout)
                    temp_accumulation += float(temp_reading[0])
                    power = subprocess.run(power_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
                    power_reading = power.stdout.split()
                    mhz_accumulation += int(power_reading[2])
                    watt_accumulation += float(power_reading[3])
                    stat_track_iterations += 1
                    if len(stdout_line) > 6:
                        print(f"Temp: {temp_reading[0]} MHz: {power_reading[2]} Pwr: {power_reading[3]} FPS: {stdout_line[2]} Time: {stdout_line[6]}", end='\r')

        run_time_accumulation += time.time() - start_time

        os.remove("/tmp/output.mkv")

    run_time_avg = run_time_accumulation / num_tests
    temp_avg = temp_accumulation / stat_track_iterations
    mhz_avg = mhz_accumulation / stat_track_iterations
    watt_avg = watt_accumulation / stat_track_iterations
    print(f"run time: {run_time_avg}")
    print(f"temperature: {temp_avg}")
    print(f"mhz: {mhz_avg}")
    print(f"power: {watt_avg}")

    avx512 = "Off"

    if (encoder == "libx265" and avx != "") or (encoder == "libsvtav1" and avx == ""):
        avx512 = "On"

    results_dict = {
        "AVX512": avx512,
        "Resolution": resolution,
        "Encoder": encoder,
        "Preset": preset,
        "Run Time": round(run_time_avg, 2),
        "Temperature": round(temp_avg, 2),
        "MHz": round(mhz_avg, 2),
        "Watts": round(watt_avg, 2)
    }

    with open('results.csv', 'a', newline='') as results_file:
        writer = csv.DictWriter(results_file, csv_fieldnames)
        writer.writerow(results_dict)

def main() -> None:
    with open('results.csv', 'w', newline='') as results_file:
        writer = csv.DictWriter(results_file, csv_fieldnames)
        writer.writeheader()

    x265_presets = ["ultrafast", "superfast", "veryfast", "faster", "fast",
                    "medium", "slow", "slower", "veryslow", "placebo"]
    av1_presets = ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9", "10", "11", "12", "13"]
    resolutions = ["720p", "1080p", "4k"]

    for preset in av1_presets:
        for resolution in resolutions:
            ffmpeg_cmd(resolution, "libsvtav1", preset, "-svtav1-params \"asm=9\"")
            #ffmpeg_cmd(resolution, "libsvtav1", preset)

    for preset in x265_presets:
        for resolution in resolutions:
            ffmpeg_cmd(resolution, "libx265", preset)
            #ffmpeg_cmd(resolution, "libx265", preset, "-x265-params \"asm=avx512\"")

if __name__ == "__main__":
    main()