using System.Diagnostics;
using System.Text;
using EncodingSampleTest;
using FFMpegCore;
using FFMpegCore.Enums;

const int MAX_JOBS = 8;
string ffmpegPath = Path.Combine(Environment.CurrentDirectory, "vendor");
string tempPath = Path.Combine(Environment.CurrentDirectory, "tmp");
string outputDir = Path.Combine(Environment.CurrentDirectory, "output");
const string inputFile = "testVid.mp4";
Directory.CreateDirectory(outputDir);
Directory.CreateDirectory(tempPath);
const string tableHeader = "Width\tHeight\tPreset\tCRF\tSize Ratio\tEncode Time\tSpeed\tVMAF\tSSIM\tPSNR_avg\tPSNR_MSE\n";
StringBuilder AV1GPUStats = new(tableHeader);
StringBuilder HEVCGPUStats = new(tableHeader);
StringBuilder H264GPUStats = new(tableHeader);
StringBuilder AV1CPUStats = new(tableHeader);
StringBuilder HEVCCPUStats = new(tableHeader);
StringBuilder H264CPUStats = new(tableHeader);

GlobalFFOptions.Configure(options => {
    options.BinaryFolder = ffmpegPath;
    options.TemporaryFilesFolder = tempPath;
    options.UseCache = true;
});

if(!File.Exists(inputFile)) return;

var originalFile = new FileInfo(inputFile);
var mediaInfo = FFProbe.Analyse(inputFile);
if (mediaInfo.PrimaryVideoStream == null) {
    Console.WriteLine("No video stream found.");
    return;
}
var W = mediaInfo.PrimaryVideoStream.Width;
var H = mediaInfo.PrimaryVideoStream.Height;

List<Task<EncodeStats>> GpuTasks = new();
List<Task<EncodeStats>> CpuTasks = new();
CancellationTokenSource cancellationTokenSource = new CancellationTokenSource();

CancellationToken cancellationToken = cancellationTokenSource.Token;
Console.WriteLine("Generating Tasks:");
for (int cq = 24; cq <= 34; cq++) {
    for (int preset = 12; preset <= 18; preset++) {
        //GpuTasks.Add(Encode(Encoder.AV1_NVENC, 1280, 720, cq, preset, cancellationToken));
        //GpuTasks.Add(Encode(Encoder.AV1_NVENC, 1920, 1080, cq, preset, cancellationToken));
        GpuTasks.Add(Encode(Encoder.AV1_NVENC, W, H, cq, preset, cancellationToken));

        //GpuTasks.Add(Encode(Encoder.HEVC_NVENC, 1280, 720, cq, preset, cancellationToken));
        //GpuTasks.Add(Encode(Encoder.HEVC_NVENC, 1920, 1080, cq, preset, cancellationToken));
        GpuTasks.Add(Encode(Encoder.HEVC_NVENC, W, H, cq, preset, cancellationToken));

        //GpuTasks.Add(Encode(Encoder.H264_NVENC, 1280, 720, cq, preset, cancellationToken));
        //GpuTasks.Add(Encode(Encoder.H264_NVENC, 1920, 1080, cq, preset, cancellationToken));
        //GpuTasks.Add(Encode(Encoder.H264_NVENC, W, H, cq, preset, cancellationToken));
    }
    for (int preset = 3; preset <= 10; preset++) {
        //CpuTasks.Add(Encode(Encoder.AV1_CPU,  1280, 720, cq, preset, cancellationToken));
        //CpuTasks.Add(Encode(Encoder.AV1_CPU,  1920, 1080, cq, preset, cancellationToken));
        CpuTasks.Add(Encode(Encoder.AV1_CPU,  W, H, cq, preset, cancellationToken));
    }
    for (int preset = 2; preset <= 5; preset++) {
        //CpuTasks.Add(Encode(Encoder.HEVC_CPU, 1280, 720, cq, preset, cancellationToken));
        //CpuTasks.Add(Encode(Encoder.HEVC_CPU, 1920, 1080, cq, preset, cancellationToken));
        CpuTasks.Add(Encode(Encoder.HEVC_CPU, W, H, cq, preset, cancellationToken));
        
        //CpuTasks.Add(Encode(Encoder.H264_CPU, 1280, 720, cq, preset, cancellationToken));
        //CpuTasks.Add(Encode(Encoder.H264_CPU, 1920, 1080, cq, preset, cancellationToken));
        //CpuTasks.Add(Encode(Encoder.H264_CPU, W, H, cq, preset, cancellationToken));
    }
}

//Run all GPU tasks sequentially
while (GpuTasks.Count(x => x.IsCompleted) < GpuTasks.Count) {
    var task = GpuTasks.First(x => !x.IsCompleted);
    task.Start();
    await task;
    Console.WriteLine("Computing Scores...");
    await ComputeScores(task.Result);
    Console.WriteLine(GpuTasks.Count(x => !x.IsCompleted) + " GPU tasks remaining...");
}

//Run all CPU tasks sequentially
while (CpuTasks.Count(x => x.IsCompleted) < CpuTasks.Count) {
    var task = CpuTasks.First(x => !x.IsCompleted);
    task.Start();
    await task;
    Console.WriteLine("Computing Scores...");
    await ComputeScores(task.Result);
    Console.WriteLine(CpuTasks.Count(x => !x.IsCompleted) + " CPU tasks remaining...");
}

cancellationTokenSource.Cancel();
return;

async Task ComputeScores(EncodeStats es) {
    Console.WriteLine("Computing VMAF...");
    var vmaf = ComputeVMAFScore(inputFile, es.Filename);//Task.FromResult(-1f);
    Console.WriteLine($"VMAF= {vmaf.Result}");
    var psnr = ComputePSNRScore(inputFile, es.Filename);
    Console.WriteLine($"PSNR= {psnr.Result.Item1} - {psnr.Result.Item2}");
    var ssim = ComputeSSIMScore(inputFile, es.Filename);
    Console.WriteLine($"SSIM= {ssim.Result}");
    await Task.WhenAll([vmaf, psnr, ssim]);
    string line = $"{es.Width}\t{es.Height}\t{es.preset}\t{es.CRF}\t{es.SizeRatio:F3}\t{es.EncodeTime:g}\t"+
                  $"{(mediaInfo.Duration.TotalSeconds / es.EncodeTime.TotalSeconds):F3}x\t"+
                  $"{vmaf.Result:F2}\t" + //VMAF
                  $"{ssim.Result:F4}\t" + //SSIM
                  $"{psnr.Result.Item1:F2}\t" + //PSNR_avg
                  $"{psnr.Result.Item2:F2}\n";    //PSNR_MSE
    
    Console.WriteLine(line);
    switch (es.EncoderType) {
        case Encoder.AV1_NVENC: AV1GPUStats.Append(line); break;
        case Encoder.HEVC_NVENC: HEVCGPUStats.Append(line); break;
        case Encoder.H264_NVENC: H264GPUStats.Append(line); break;
        case Encoder.AV1_CPU: AV1CPUStats.Append(line); break;
        case Encoder.HEVC_CPU: HEVCCPUStats.Append(line); break;
        case Encoder.H264_CPU: H264CPUStats.Append(line); break;
    }
    WriteAllStats();
}

void WriteAllStats() {
    File.WriteAllText(Path.Combine(outputDir, "AV1_NVENC_Stats.tsv"), AV1GPUStats.ToString());
    File.WriteAllText(Path.Combine(outputDir, "HEVC_NVENC_Stats.tsv"), HEVCGPUStats.ToString());
    //File.WriteAllText(Path.Combine(outputDir, "H264_NVENC_Stats.tsv"), H264GPUStats.ToString());
    File.WriteAllText(Path.Combine(outputDir, "AV1_CPU_Stats.tsv"), AV1CPUStats.ToString());
    File.WriteAllText(Path.Combine(outputDir, "HEVC_CPU_Stats.tsv"), HEVCCPUStats.ToString());
    //File.WriteAllText(Path.Combine(outputDir, "H264_CPU_Stats.tsv"), H264CPUStats.ToString());
}

Task<EncodeStats> Encode(Encoder encoder, int W = -1, int H = -1, int CQ = 23, int preset = 0, CancellationToken cancellationToken = default) {
    var outputFile = Path.Combine(outputDir, $"output_{encoder}-{W}x{H}_p{preset}_cq{CQ}.mp4");
    
    FFMpegArgumentProcessor ffmpegArgs = encoder switch {
        Encoder.AV1_NVENC   => AV1Encode(     outputFile, W, H, CQ, preset),
        Encoder.HEVC_NVENC  => HEVCEncode(    outputFile, W, H, CQ, preset),
        Encoder.H264_NVENC  => H264Encode(    outputFile, W, H, CQ, preset),
        Encoder.AV1_CPU     => AV1EncodeCPU(  outputFile, W, H, CQ, preset),
        Encoder.HEVC_CPU    => HEVCEncodeCPU( outputFile, W, H, CQ, preset),
        Encoder.H264_CPU    => H264EncodeCPU( outputFile, W, H, CQ, preset),
        _ => throw new ArgumentOutOfRangeException(nameof(encoder), encoder, null)
    };
    
    ffmpegArgs.WithLogLevel(FFMpegLogLevel.Info)
        .NotifyOnProgress(progress => Console.WriteLine($"Encoding {outputFile}: {progress:g}/{mediaInfo.Duration:g} {progress / mediaInfo.Duration:P}"))
        .CancellableThrough(cancellationToken);
    
    Console.WriteLine(ffmpegArgs.Arguments);
    return new Task<EncodeStats>(_ => { 
        Console.WriteLine("Starting encoding: " + outputFile);
        var sw = Stopwatch.StartNew();
        ffmpegArgs.ProcessSynchronously();
        sw.Stop();
        var fileSize = new FileInfo(outputFile).Length;
        return new EncodeStats() {
            Filename = outputFile,
            EncoderType = encoder,
            Width = W,
            Height = H,
            CRF = CQ,
            preset = preset,
            SizeRatio = originalFile.Length / (float)fileSize,
            EncodeTime = sw.Elapsed
        };
    }, cancellationToken);
}

FFMpegArgumentProcessor AV1Encode(string filePath, int W = -1, int H = -1, int CQ = 23, int preset = 0) { // AV1 is visually lossless at CQ 23
    return FFMpegArguments
        .FromFileInput(inputFile, false, options => options.WithHardwareAcceleration())
        .OutputToFile(filePath, true, options => options
            .CopyChannel(Channel.Audio)
            .CopyChannel(Channel.Subtitle)
            .WithVideoCodec("av1_nvenc")
            .WithArgument(new NvencSpeedPreset((NvencSpeed)preset))
            .WithArgument(new NvencTuneArgument(NvencTune.hq))
            .WithArgument(new NvencHighBitDepthArgument(true))
            .WithCustomArgument("-rc vbr")
            .WithCustomArgument($"-cq {CQ}")
            .WithVideoFilters(filterOptions => { if (W > 0 && H > 0) filterOptions.Scale(W, H); })
            .WithFastStart()
        );
}

FFMpegArgumentProcessor HEVCEncode(string filePath, int W = -1, int H = -1, int CQ = 23, int preset = 0) {
    return FFMpegArguments
        .FromFileInput(inputFile, true, options => options.WithHardwareAcceleration())
        .OutputToFile(filePath, true, options => options
            .CopyChannel(Channel.Audio)
            .CopyChannel(Channel.Subtitle)
            .WithVideoCodec("hevc_nvenc")
            .WithArgument(new NvencSpeedPreset((NvencSpeed)preset))
            .WithArgument(new NvencTuneArgument(NvencTune.hq))
            .WithCustomArgument("-rc vbr_hq")
            .WithCustomArgument($"-cq {(byte)CQ}")
            .WithVideoFilters(filterOptions => { if (W > 0 && H > 0) filterOptions.Scale(W, H); })
            .WithFastStart()
        );
}

FFMpegArgumentProcessor H264Encode(string filePath, int W = -1, int H = -1, int CQ = 23, int preset = 0) {
    return FFMpegArguments
        .FromFileInput(inputFile, true, options => options.WithHardwareAcceleration())
        .OutputToFile(filePath, true, options => options
            .CopyChannel(Channel.Audio)
            .CopyChannel(Channel.Subtitle)
            .WithVideoCodec("h264_nvenc")
            .WithArgument(new NvencSpeedPreset((NvencSpeed)preset))
            .WithArgument(new NvencTuneArgument(NvencTune.hq))
            .WithCustomArgument("-rc vbr_hq")
            .WithCustomArgument($"-cq {(byte)CQ}") // approximate CRF to CQ
            .WithVideoFilters(filterOptions => { if (W > 0 && H > 0) filterOptions.Scale(W, H); })
            .WithFastStart()
        );
}

FFMpegArgumentProcessor AV1EncodeCPU(string filePath, int W = -1, int H = -1, int CQ = 23, int preset = 0) {
    return FFMpegArguments
        .FromFileInput(inputFile)
        .OutputToFile(filePath, true, options => options
            .CopyChannel(Channel.Audio)
            .CopyChannel(Channel.Subtitle)
            .WithVideoCodec("libsvtav1")
            .WithCustomArgument($"-preset {preset}")
            .WithConstantRateFactor(CQ)
            .WithCustomArgument("-svtav1-params tune=0")
            .WithVideoFilters(filterOptions => {if (W > 0 && H > 0) filterOptions.Scale(W, H);})
            .WithFastStart()
        );
}

FFMpegArgumentProcessor HEVCEncodeCPU(string filePath, int W = -1, int H = -1, int CQ = 23, int preset = 0) {
    return FFMpegArguments
        .FromFileInput(inputFile)
        .OutputToFile(filePath, true, options => options
            .CopyChannel(Channel.Audio)
            .CopyChannel(Channel.Subtitle)
            .WithVideoCodec(VideoCodec.LibX265)
            .WithSpeedPreset(ToX264Speed(preset))
            .WithConstantRateFactor(CQ)
            .WithCustomArgument("-x265-params profile=main10:high-tier=1:aq-mode=3:psy-rd=1:rc-lookahead=60:bframes=6")
            .WithVideoFilters(filterOptions => {if (W > 0 && H > 0) filterOptions.Scale(W, H);})
            .WithFastStart()
        );
}

FFMpegArgumentProcessor H264EncodeCPU(string filePath, int W = -1, int H = -1, int CQ = 23, int preset = 0) {
    return FFMpegArguments
        .FromFileInput(inputFile)
        .OutputToFile(filePath, true, options => options
            .CopyChannel(Channel.Audio)
            .CopyChannel(Channel.Subtitle)
            .WithVideoCodec(VideoCodec.LibX264)
            .WithSpeedPreset(ToX264Speed(preset))
            .WithConstantRateFactor(CQ)
            .WithCustomArgument("-x265-params profile=main10:high-tier=1:aq-mode=3:psy-rd=1:rc-lookahead=60:bframes=6")
            .WithVideoFilters(filterOptions => {if (W > 0 && H > 0) filterOptions.Scale(W, H);})
            .WithFastStart()
        );
}

async Task<float> ComputeVMAFScore(string referenceFile, string distortedFile) {
    float VMAF = -1;
    ProcessStartInfo startInfo = new ProcessStartInfo {
        FileName = Path.Combine(ffmpegPath, "ffmpeg"),
        Arguments = $"-i \"{distortedFile}\" -i \"{referenceFile}\"  -lavfi libvmaf=n_threads=32 -f null -",
        CreateNoWindow = true,
        UseShellExecute = false,
        RedirectStandardOutput = false,
        RedirectStandardError = true
    };
    using Process ffmpeg = Process.Start(startInfo)!;
    //string output = await ffmpeg.StandardOutput.ReadToEndAsync();
    string output = await ffmpeg.StandardError.ReadToEndAsync();
    await ffmpeg.WaitForExitAsync();
    // Parse VMAF score from output
    string marker = "VMAF score: ";
    int index = output.LastIndexOf(marker, StringComparison.Ordinal);
    if (index != -1) {
        int start = index + marker.Length;
        int end = output.IndexOf("\n", start, StringComparison.Ordinal);
        string scoreStr = output.Substring(start, end - start).Trim();
        if (float.TryParse(scoreStr, out float score)) {
            VMAF = score;
        }
    }
    return VMAF;
}

/// <summary>
/// Computes the PSNR score between a reference and distorted video file.
/// </summary>
/// <param name="referenceFile">The path to the reference video file.</param>
/// <param name="distortedFile">The path to the distorted video file.</param>
/// <returns>A tuple containing the average PSNR and average MSE values.</returns>
async Task<Tuple<float,float>> ComputePSNRScore(string referenceFile, string distortedFile) {
    string dir = Environment.CurrentDirectory;
    Environment.CurrentDirectory = Path.GetDirectoryName(distortedFile)!;
    string psnrLogFile = Path.Combine(Path.GetDirectoryName(distortedFile)!, $"{Path.GetFileNameWithoutExtension(distortedFile)}.psnr");
    
    float PSNR = -1;
    float MSE = -1;
    ProcessStartInfo startInfo = new ProcessStartInfo {
        FileName = Path.Combine(ffmpegPath, "ffmpeg"),
        //TODO: This is a damn hack.
        Arguments = $"-i \"{distortedFile}\" -i \"..\\{referenceFile}\" -lavfi psnr=stats_file={Path.GetFileName(psnrLogFile)}:stats_version=2 -f null -",
        UseShellExecute = false,
        CreateNoWindow = true
    };
    using Process ffmpeg = Process.Start(startInfo)!;
    await ffmpeg.WaitForExitAsync();
    // Parse PSNR score from stats file
    var PsnrLog = File.ReadAllLines(psnrLogFile);
    
    List<float> mseVals = new();
    List<float> psnrVals = new();
    foreach (var line in PsnrLog) {
        if (!line.Contains("mse_avg:")) continue; // Skip header and irrelevant lines
        string[] parts = line.Split(' ');
        foreach (var part in parts) { 
            if (part.StartsWith("psnr_avg:")) psnrVals.Add(float.Parse(part.Split(':')[1]));
            if (part.StartsWith("mse_avg:")) mseVals.Add(float.Parse(part.Split(':')[1]));
        }
    }
    PSNR = psnrVals.Count > 0 ? psnrVals.Average() : -1;
    MSE = mseVals.Count > 0 ? mseVals.Average() : -1;
    Environment.CurrentDirectory = dir;
    return new (PSNR, MSE);
}

async Task<float> ComputeSSIMScore(string referenceFile, string distortedFile) {
    float SSIM = -1;
    ProcessStartInfo startInfo = new ProcessStartInfo {
        FileName = Path.Combine(ffmpegPath, "ffmpeg"),
        Arguments = $"-i \"{distortedFile}\" -i \"{referenceFile}\" -lavfi ssim -f null -",
        RedirectStandardOutput = false,
        RedirectStandardError = true,
        UseShellExecute = false,
        CreateNoWindow = true
    };
    using Process ffmpeg = Process.Start(startInfo)!;
    string output = await ffmpeg.StandardError.ReadToEndAsync();
    await ffmpeg.WaitForExitAsync();
    
    // Parse SSIM score from output
    string marker = "All:";
    int index = output.LastIndexOf(marker, StringComparison.Ordinal);
    if (index == -1) return SSIM;
    int start = index + marker.Length;
    int end = output.IndexOf(" ", start, StringComparison.Ordinal);
    string scoreStr = output.Substring(start, end - start).Trim();
    if (float.TryParse(scoreStr, out float score)) SSIM = score;
    return SSIM;
}

Speed ToX264Speed(int p) => p switch {
    1 => Speed.VerySlow,
    2 => Speed.Slower,
    3 => Speed.Slow,
    4 => Speed.Medium,
    5 => Speed.Fast,
    6 => Speed.Faster,
    7 => Speed.VeryFast,
    _ => Speed.Medium
};

public enum Encoder {
    AV1_NVENC,
    HEVC_NVENC,
    H264_NVENC,
    AV1_CPU,
    HEVC_CPU,
    H264_CPU
}

record struct EncodeStats {
    public string Filename;
    public Encoder EncoderType;
    public int Width;
    public int Height;
    public int CRF;
    public int preset;
    public float SizeRatio;
    public TimeSpan EncodeTime;
}