-- README:
-- original version by Zehkul https://gist.github.com/Zehkul/25ea7ae77b30af959be0
-- needs: yad, libnotify (and --ytdl if you want to encode from streams)
-- press any of alt + w, g or x to set the start frame
-- press again to set the end frame and use
-- alt + w to make a webm
-- alt + g to make a gif
-- alt + x to make a x264 encoded mkv

-- Note: encoding a webm to a specific filesize will only work if audio is disabled
--       and including subs will make them lose their styling

local msg = require 'mp.msg'
local opt = require 'mp.options'
local mputils = require 'mp.utils'

-- default options, convert_script.conf is read
local options = {
    bitrate_multiplier = 0.975,         -- to make sure the file won’t go over the target file size, set it to 1 if you don’t care
}

read_options(options, "convert_script_webm")
read_options(options, "convert_script_gif")
read_options(options, "convert_script_x264")

-----------------
-- Main script --
-----------------

function set_timepos(func)

    if timepos1 then
        
        timepos2 = mp.get_property("time-pos")
        timepos2_humanreadable = mp.get_property_osd("time-pos")
    
        if tonumber(timepos1) > tonumber(timepos2) then
            
            length = timepos1-timepos2
            start = timepos2
            start_humanreadable = timepos2_humanreadable
            end_humanreadable = timepos1_humanreadable
            msg.info("End frame set")
        
        elseif tonumber(timepos2) > tonumber(timepos1) then
        
            length = timepos2-timepos1
            start = timepos1
            start_humanreadable = timepos1_humanreadable
            end_humanreadable = timepos2_humanreadable
            msg.info("End frame set")
    
        else
    
            msg.error("Both frames are the same, ignoring the second one")
            mp.osd_message("Both frames are the same, ignoring the second one")
            timepos2 = nil
            return
           
        end
    
        timepos1 = nil
        func.call_gui()
    
    else
        
        timepos1 = mp.get_property("time-pos")
        timepos1_humanreadable = mp.get_property_osd("time-pos")
        msg.info("Start frame set")
        mp.osd_message("Start frame set")
    
    end

end

function convert_script_hotkey_call_webm ()
    
    func = { 
      call_gui = call_gui_webm
    }

    set_timepos(func)

end

function convert_script_hotkey_call_gif ()
    
    func = { 
      call_gui = call_gui_gif
    }

    set_timepos(func)

end

function convert_script_hotkey_call_x264 ()
    
    func = { 
      call_gui = call_gui_x264
    }

    set_timepos(func)

end

------------
-- Encode --
------------

function generate_filenames(ext)   
    
    local filename_ext = mp.get_property_osd("filename")
    filename_ext = string.gsub(filename_ext, "'", "'\\''")
    local filename = string.gsub(filename_ext, "%....$","")
    
    if string.len(filename) > 230 then
    
        filename = mp.get_property("options/title")
        if filename == 'mpv - ${media-title}' or string.len(filename) > 230 then
            filename = 'output'
        end
        
    end
    
    local path = mp.get_property("path", "")
    local dir, fil = mputils.split_path(path)

    file_in = dir .. fil
    file_out = dir .. filename .. "-out." .. ext

    file_in = string.gsub(file_in, "'", "'\\''")
    file_out = string.gsub(file_out, "'", "'\\''")

    if string.sub(file_in,1,string.len("http")) == "http" then

        file_in = mp.get_property("stream-open-filename")
        file_out = os.getenv("HOME") .. "/" .. mp.get_property("media-title") .. "-out." .. ext
        file_out = string.gsub(file_out, "'", "'\\''")

    end

end

function encode_webm ()
    
    generate_filenames("webm")
    
    local full_command = '(ffmpeg -ss ' .. start .. ' -i \'' .. file_in .. '\' -t ' .. length .. ' -b:v ' .. bitrate .. bitsize .. ' -vf scale=' .. scale .. ':-1 -quality good -cpu-used 0 -pass 1 ' .. audio .. ' ' .. subs .. ' -f webm /dev/null -y' 

    full_command = full_command .. ' && ffmpeg -ss ' .. start .. ' -i \'' .. file_in .. '\' -t ' .. length .. ' -b:v ' .. bitrate .. bitsize .. ' -vf scale=' .. scale .. ':-1 -quality good -cpu-used 0 -pass 2 ' .. audio .. ' ' .. subs .. ' -f webm \'' .. file_out .. '\' -y'

    full_command = full_command .. ') && notify-send "Encode done!"'

    msg.info(full_command)
    os.execute(full_command)
    
end

function encode_gif ()
    
    generate_filenames("gif")
    
    local full_command = '(ffmpeg -ss ' .. start .. ' -i \'' .. file_in .. '\' -t ' .. length .. ' -vf scale=' .. scale .. ':-1 -r ' .. fps .. ' \'' .. file_out .. '\' -y'

    full_command = full_command .. ') && notify-send "Encode done!"'
   
    msg.info(full_command)
    os.execute(full_command)
    
end

function encode_x264 ()
    
    generate_filenames("mkv")
    
    local full_command = '(ffmpeg -ss ' .. start .. ' -i \'' .. file_in .. '\' -t ' .. length .. ' -c:v libx264 -vf scale=' .. sheight .. ':' .. swidth .. ' -preset slow -qp ' .. quality .. ' -c:a copy -c:s copy -map 0 \'' .. file_out .. '\' -y'

    full_command = full_command .. ') && notify-send "Encode done! (' .. quality .. ')"'
   
    msg.info(full_command)
    os.execute(full_command)
    
end

---------
-- GUI --
---------

function call_gui_webm ()

    mp.resume_all()
    local handle = io.popen('yad --title="Convert Script" --center --form --separator="\n" --field="Resize to:NUM" "720" --field="Don’t resize at all:CHK" "false" --field="Include audio:CHK" "false" --field="Include subs:CHK" "false" --field="Bitrate/Filesize (M):NUM" "3" --field="Filesize > Bitrate:CHK" "true" --button="gtk-cancel:2" --button="gtk-ok:0" && echo "$?"')
    local yad = handle:read("*a")
    handle:close()
   
    if yad == "" then
        return
    end
    
    local yad_table = {}
    
    local i = 0
    for k in string.gmatch(yad, "[%a%d]+") do
       yad_table[i] = k
       i = i + 1
    end
    
    if (yad_table[2] == "FALSE") then
        scale = yad_table[0]
    else
        scale = "-1"
    end
    
    if yad_table[3] == "FALSE" then
        audio = '-an'
    else
        audio = ""
    end

    if yad_table[4] == "FALSE" then
        subs = '-sn'
    else
        subs = ""
    end

    if yad_table[7] == "TRUE" then
        bitrate = math.floor(yad_table[5]*1024*8/length*options.bitrate_multiplier)
        bitsize = "K"
    else
        bitrate = yad_table[5]
        bitsize = "M"
    end
    
    if yad_table[8] == "0" then
        encode_webm()
    end

end

function call_gui_gif ()

    mp.resume_all()
    local handle = io.popen('yad --title="Convert Script" --center --form --separator="\n" --field="Resize to:NUM" "320" --field="Don’t resize at all:CHK" "false" --field="FPS:NUM" "24" --button="gtk-cancel:2" --button="gtk-ok:0" && echo "$?"')
    local yad = handle:read("*a")
    handle:close()
   
    if yad == "" then
        return
    end
    
    local yad_table = {}
    
    local i = 0
    for k in string.gmatch(yad, "[%a%d]+") do
       yad_table[i] = k
       i = i + 1
    end
    
    if (yad_table[2] == "FALSE") then
        scale = yad_table[0]
    else
        scale = "-1"
    end

    fps = yad_table[3]
    
    if yad_table[5] == "0" then
        encode_gif()
    end

end

function call_gui_x264 ()

    mp.resume_all()
    local handle = io.popen('yad --title="Convert Script" --center --form --separator="\n" --field="Resize height:NUM" "720" --field="Resize width:NUM" "480" --field="Don’t resize at all:CHK" "true" --field="Quality (0-51):NUM" "18" --button="gtk-cancel:2" --button="gtk-ok:0" && echo "$?"')
    local yad = handle:read("*a")
    handle:close()
   
    if yad == "" then
        return
    end
    
    local yad_table = {}
    
    local i = 0
    for k in string.gmatch(yad, "[%a%d]+") do
       yad_table[i] = k
       i = i + 1
    end

    if (yad_table[4] == "FALSE") then
        sheight = yad_table[0]
    else
        sheight = "-1"
    end

    if (yad_table[4] == "FALSE") then
        swidth = yad_table[2]
    else
        swidth = "-1"
    end
    
    if ((sheight ~= "-1" and swidth ~= "-1") and (tonumber(sheight)%2 ~= 0 or tonumber(swidth)%2 ~= 0)) then
        os.execute('notify-send "' .. sheight ..':' .. swidth .. '"')
        return
    end

    quality = yad_table[5];

    if yad_table[7] == "0" then
        encode_x264()
    end

end

mp.add_key_binding("alt+w", "convert_script_webm", convert_script_hotkey_call_webm)
mp.add_key_binding("alt+g", "convert_script_gif", convert_script_hotkey_call_gif)
mp.add_key_binding("alt+x", "convert_script_x264", convert_script_hotkey_call_x264)