-- transcribe-subtitles.lua
-- Press T to generate subtitles for the current video via faster-whisper.
-- Spawns transcription in background, polls for SRT file, loads it when done.
--
-- input.conf:  T    script-binding transcribe_subtitles/transcribe

local mp = require("mp")
local utils = require("mp.utils")

local function load_srt(srt_path)
    mp.commandv("sub-add", srt_path)
    local count = mp.get_property_number("track-list/count", 0)
    for i = 0, count - 1 do
        if mp.get_property("track-list/" .. i .. "/type") == "sub"
        and mp.get_property("track-list/" .. i .. "/external") == "yes"
        and mp.get_property("track-list/" .. i .. "/external-filename") == srt_path then
            mp.set_property_number("sub-track-id", i + 1)
            break
        end
    end
    mp.osd_message("Subtitles loaded", 2)
end

local function poll_srt(srt_path, elapsed, max_wait)
    if elapsed >= max_wait then
        mp.osd_message("Transcription timed out", 3)
        return
    end

    local f = io.open(srt_path, "r")
    if f then
        f:close()
        load_srt(srt_path)
        return
    end

    mp.add_timeout(2, function()
        poll_srt(srt_path, elapsed + 2, max_wait)
    end)
end

local function transcribe()
    local file_path = mp.get_property("path")
    if not file_path then mp.osd_message("No file loaded.", 2); return end

    local abs_path = mp.get_property("stream-path")
        or utils.join_path(mp.get_property("working-directory", "."), file_path)

    local base = abs_path:gsub("%.[^./]+$", "")
    local srt_path = base .. ".srt"

    -- If SRT already exists, just load it
    local existing = io.open(srt_path, "r")
    if existing then
        existing:close()
        load_srt(srt_path)
        return
    end

    local transcriber = os.getenv("HOME") .. "/.local/bin/transcribe-video"
    mp.osd_message("Transcribing subtitles ...", 1)

    local shell_cmd = "export LD_LIBRARY_PATH=/opt/cuda/targets/x86_64-linux/lib; "
        .. transcriber .. " --srt \"" .. abs_path .. "\" \"" .. srt_path .. "\" &>/dev/null &"

    mp.command_native({
        name = "subprocess",
        args = { "bash", "-c", shell_cmd },
        playback_only = false,
    })

    poll_srt(srt_path, 0, 900)
end

mp.add_key_binding(nil, "transcribe", transcribe)