-- Usage:
--    Shift + B - create bilingual subtitles (and automatically select as default subtitles with visibility set to true)
-- Note: 
--    Requires the original subtitles and translated subtitles (*.srt) alongside with the video file. 
-- Status:
--    Experimental & not interested.

------- Script Options -------
srt_original_file_extensions = {".srt", ".en.srt", ".eng.srt"}
srt_translated_file_extensions = {".ru.srt", ".rus.srt"}
------------------------------

function srt_time_to_seconds(time)
    local major, minor = time:match("(%d%d:%d%d:%d%d),(%d%d%d)")
    local hours, mins, secs = major:match("(%d%d):(%d%d):(%d%d)")
    return hours * 3600 + mins * 60 + secs + minor / 1000
end

function seconds_to_ass_time(time)
    local hours = math.floor(time / 3600)
    local mins = math.floor(time / 60) % 60
    local secs = math.floor(time % 60)
    local milliseconds = (time * 1000) % 1000

    return string.format("%d:%02d:%02d.%02d", hours, mins, secs, milliseconds / 10)
end

function open_subtitles_file(srt_file_exts)
    local srt_base_filename = mp.get_property("working-directory") .. "/" .. mp.get_property("filename/no-ext")
    
    for i, ext in ipairs(srt_file_exts) do
        local f, err = io.open(srt_base_filename .. ext, "r")
        
        if f then
            return f
        end
    end

    return false
end

function read_subtitles(srt_file_exts)
    local f = open_subtitles_file(srt_file_exts)

    local subs = {}
    local subs_start = {}
    local subs_end = {}

    if f then
        local data = f:read("*all")
        data = string.gsub(data, "\r\n", "\n")
        f:close()
        
        for start_time, end_time, text in string.gmatch(data, "(%d%d:%d%d:%d%d,%d%d%d) %-%-> (%d%d:%d%d:%d%d,%d%d%d)\n(.-)\n\n") do
          table.insert(subs, text)
          table.insert(subs_start, srt_time_to_seconds(start_time))
          table.insert(subs_end, srt_time_to_seconds(end_time))
        end
    end

    return subs, subs_start, subs_end
end

function convert_into_sentences(subs, subs_start, subs_end)
    local sentences = {}
    local sentences_start = {}
    local sentences_end = {}

    for i, sub_text in ipairs(subs) do
        local sub_start = subs_start[i]
        local sub_end = subs_end[i]
        local sub_text = string.gsub(sub_text, "<[^>]+>", "")

        if sub_text:find("^- ") ~= nil and sub_text:sub(3,3) ~= sub_text:sub(3,3):upper() then
           sub_text = string.gsub(sub_text, "^- ", "")
        end

        if #sentences > 0 then
            local prev_sub_start = sentences_start[#sentences]
            local prev_sub_end = sentences_end[#sentences]
            local prev_sub_text = sentences[#sentences]

            if (sub_start - prev_sub_end) <= 2 and sub_text:sub(1,1) ~= '-' and 
                    sub_text:sub(1,1) ~= '"' and sub_text:sub(1,1) ~= "'" and sub_text:sub(1,1) ~= '(' and
                    (prev_sub_text:sub(prev_sub_text:len()) ~= "." or prev_sub_text:sub(prev_sub_text:len()-2) == "...") and
                    prev_sub_text:sub(prev_sub_text:len()) ~= "?" and prev_sub_text:sub(prev_sub_text:len()) ~= "!" and
                    (sub_text:sub(1,1) == sub_text:sub(1,1):lower() or prev_sub_text:sub(prev_sub_text:len()) == ",") then
                
                local text = sentences[#sentences] .. " " .. sub_text
                text = string.gsub(text, "\n", "#")
                text = string.gsub(text, "%.%.%. %.%.%.", " ")
                text = string.gsub(text, "#%-", "\n-")
                text = string.gsub(text, "#", " ")
                if text:match("\n%-") ~= nil and text:match("^%-") == nil then
                    text = "- " .. text
                end
                
                sentences[#sentences] = text
                sentences_end[#sentences] = sub_end
            else
                table.insert(sentences, sub_text)
                table.insert(sentences_start, sub_start)
                table.insert(sentences_end, sub_end)    
            end
        else
            table.insert(sentences, sub_text)
            table.insert(sentences_start, sub_start)
            table.insert(sentences_end, sub_end)
        end
    end

    return sentences, sentences_start, sentences_end
end

function string.starts(s, substring)
   return string.sub(s, 1, string.len(substring)) == substring
end

function write_bilingual_subtitles(subs_original, subs_original_start, subs_original_end, subs_translated, subs_translated_start, subs_translated_end, subtitles_filename)
    if #subs_original == 0 or #subs_translated == 0 then
        return false
    end

    local screenx, screeny, aspect = mp.get_osd_size()
    mp.set_osd_ass(screenx, screeny, "{\\an9}● ")

    local f = assert(io.open(subtitles_filename, "w"))

    f:write("[Script Info]", "\n\n")
    f:write("[V4+ Styles]", "\n")
    f:write("Format: Name, Fontname, Fontsize, PrimaryColour, SecondaryColour, OutlineColour, BackColour, Bold, Italic, Underline, StrikeOut, ScaleX, ScaleY, Spacing, Angle, BorderStyle, Outline, Shadow, Alignment, MarginL, MarginR, MarginV, Encoding", "\n")
    f:write("Style: English,Arial,20,&H00FFFFFF,&H000000FF,&H00000000,&H50000000,0,0,0,0,100,100,0,0,4,1,6.25,2,10,10,10,1", "\n")
    f:write("[Events]", "\n")
    f:write("Format: Layer, Start, End, Style, Name, MarginL, MarginR, MarginV, Effect, Text", "\n")
    
    local translated_sub_lines = {}
    local translates_sub_style = "\\N{\\fs2} \\N{\\fs14}{\\c&H009DDFF1&}"
    
    for i, sub_text in ipairs(subs_original) do
        local sub_start = subs_original_start[i]
        local sub_end = subs_original_end[i]
        sub_text = string.gsub(sub_text, "\n", " ")

        local flag = false
        local sub_line = ""
        for j, second_sub_text in ipairs(subs_translated) do
            local second_sub_start = subs_translated_start[j]
            local second_sub_end = subs_translated_end[j]
            second_sub_text = string.gsub(second_sub_text, "\n", " ")
            second_sub_text = string.gsub(second_sub_text, "<[^>]+>", "")

            if second_sub_end > sub_start and second_sub_start < sub_end then
                local s_sub_start = second_sub_start > sub_start and second_sub_start or sub_start
                local s_sub_end = second_sub_end > sub_end and sub_end or second_sub_end
                local s_sub_length = s_sub_end - s_sub_start

                if ((second_sub_start <= sub_start and second_sub_end >= sub_end) or (second_sub_start >= sub_start and second_sub_end <= sub_end)) or 
                   (s_sub_length > 0.5) then

                    if flag then
                        sub_line = sub_line .. " "
                    end
                    sub_line = sub_line .. second_sub_text
                    flag = true
                end
            end

            if second_sub_end > sub_end then
                break
            end
        end
        
        if flag then 
            table.insert(translated_sub_lines, sub_line)
        else
            table.insert(translated_sub_lines, "")
        end
    end

    local sub_line = ""
    local sub_line_start, sub_line_end
    for i, sub_text in ipairs(subs_original) do
        local sub_start = subs_original_start[i]
        local sub_end = subs_original_end[i]
        local sub_text = string.gsub(sub_text, "\n", " ")
        
        if i == 1 then
            sub_line = sub_text
            sub_line_start = sub_start
            sub_line_end = sub_end
        else
            if (translated_sub_lines[i] == translated_sub_lines[i-1] or string.starts(translated_sub_lines[i], translated_sub_lines[i-1])) and translated_sub_lines[i] ~= "" then
                sub_line = sub_line .. " " .. sub_text
                sub_line_end = sub_end
            else
                f:write(string.format("Dialogue: 0,%s,%s,English,,0,0,0,,%s\n", seconds_to_ass_time(sub_line_start), seconds_to_ass_time(sub_line_end), sub_line .. translates_sub_style .. translated_sub_lines[i-1]))
    
                sub_line = sub_text
                sub_line_start = sub_start
                sub_line_end = sub_end
            end
        end
    end
    
    f:write(string.format("Dialogue: 0,%s,%s,English,,0,0,0,,%s\n", seconds_to_ass_time(sub_line_start), seconds_to_ass_time(sub_line_end), sub_line .. translates_sub_style .. translated_sub_lines[#translated_sub_lines]))

    f:close()

    mp.set_osd_ass(screenx, screeny, "")

    return true
end

function create_bilingual_subtitles()
    local subs_translated, subs_translated_start, subs_translated_end = read_subtitles(srt_translated_file_extensions)

    local subs_original, subs_original_start, subs_original_end = read_subtitles(srt_original_file_extensions)
    local sentences_original, sentences_original_start, sentences_original_end = convert_into_sentences(subs_original, subs_original_start, subs_original_end)

    local subtitles_filename = mp.get_property("working-directory") .. "/" .. mp.get_property("filename/no-ext") .. ".ass"
    
    local ret = write_bilingual_subtitles(sentences_original, sentences_original_start, sentences_original_end, subs_translated, subs_translated_start, subs_translated_end, subtitles_filename)
    if ret then
        mp.commandv("sub-add", subtitles_filename)
        mp.set_property("sub-visibility", "yes")
        mp.osd_message("Finished creating bilingual subtitles")
    else
        mp.osd_message("Failed to create bilingual subtitles")
    end
end

mp.add_key_binding("B", "create-bilingual-subtitles", create_bilingual_subtitles)