-- export-loop-webp.lua
-- Press W to export the current A-B loop as both an animated WebP + GIF.
-- Frame-accurate (re-encodes). Exports raw + adjusted (gamma/brightness/etc).
--
-- WebP → ~/Videos/mpv-loops/
-- GIF  → ~/Videos/mpv-loops/gif/
--
-- Suggested input.conf binding:
--   W    script-binding export_loop_webp/export

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

-- ╔══════════════════════════════════════════╗
-- ║              CONFIGURATION               ║
-- ╚══════════════════════════════════════════╝

local OUTPUT_DIR = os.getenv("HOME") .. "/Videos/mpv-loops"
local GIF_DIR     = OUTPUT_DIR .. "/gif"
local ENABLE_OSD_MSG = true

-- Detect source frame rate (handles mf:// from webp-anim-bridge).
-- Returns at least 30fps.
local function source_fps()
    local mf_fps = mp.get_property_number("mf-fps", 0)
    if mf_fps and mf_fps > 0 then return math.max(30, mf_fps) end
    local vfps = mp.get_property_number("video-params/fps", 0)
    if vfps and vfps > 0 then return math.max(30, vfps) end
    local efps = mp.get_property_number("estimated-vf-fps", 0)
    if efps and efps > 0 then return math.max(30, efps) end
    return 30
end

-- Config — FPS is now dynamic via source_fps(). Fallback minimum for mf:// only.
local WEBP_QUALITY = 80
local WEBP_LOSSLESS = false
local WEBP_MAX_WIDTH = 800
local WEBP_LOOP = 0

-- GIF settings
local GIF_MAX_WIDTH = 800
local GIF_DITHER = "bayer"
local GIF_LOOP = 0

-- ╔══════════════════════════════════════════╗
-- ║           FORMAT HELPERS                ║
-- ╚══════════════════════════════════════════╝

local function format_ts(seconds)
    local h = math.floor(seconds / 3600)
    local m = math.floor((seconds % 3600) / 60)
    local s = seconds % 60
    return string.format("%02d:%02d:%06.3f", h, m, s)
end

local function format_label(seconds)
    local h = math.floor(seconds / 3600)
    local m = math.floor((seconds % 3600) / 60)
    local s = math.floor(seconds % 60)
    return string.format("%02d%02d%02d", h, m, s)
end

local function osd_msg(text, duration)
    if ENABLE_OSD_MSG then
        mp.osd_message(text, duration or 3)
    end
end

local function build_eq_filter()
    local brightness = mp.get_property_number("brightness", 0)
    local contrast   = mp.get_property_number("contrast", 0)
    local saturation = mp.get_property_number("saturation", 0)
    local gamma      = mp.get_property_number("gamma", 0)

    local parts = {}
    if brightness ~= 0 then table.insert(parts, string.format("brightness=%.2f", brightness / 100.0)) end
    if contrast ~= 0 then   table.insert(parts, string.format("contrast=%.2f", 1.0 + contrast / 100.0)) end
    if saturation ~= 0 then table.insert(parts, string.format("saturation=%.2f", 1.0 + saturation / 100.0)) end
    if gamma ~= 0 then      table.insert(parts, string.format("gamma=%.2f", 1.0 + gamma / 100.0)) end

    if #parts == 0 then return nil end
    return "eq=" .. table.concat(parts, ":")
end

-- Build crop filter string for current zoom/pan.
-- Returns nil if no zoom/pan is active (no filter needed).
local function zoom_pan_crop()
    local zoom = mp.get_property_number("video-zoom", -999)
    local pan_x = mp.get_property_number("video-pan-x", -999)
    local pan_y = mp.get_property_number("video-pan-y", -999)
    mp.msg.warn(string.format("[zp] raw: zoom=%+5.2f pan_x=%+5.2f pan_y=%+5.2f", zoom, pan_x, pan_y))
    -- Treat sentinel -999 or actual 0 as "no value"
    if zoom <= 0 then
        mp.msg.warn("[zp] zoom <= 0, skipping crop")
        return nil
    end
    if pan_x == -999 then pan_x = 0 end
    if pan_y == -999 then pan_y = 0 end
    local s = 2 ^ zoom
    local filter = string.format("crop=iw/%g:ih/%g:iw*(0.5-0.5/%g-%g):ih*(0.5-0.5/%g-%g)",
                                 s, s, s, pan_x, s, pan_y)
    mp.msg.warn(string.format("[zp] filter=%s", filter))
    mp.osd_message(string.format("ZOOM=%.2f PAN_X=%.2f PAN_Y=%.2f", zoom, pan_x, pan_y), 3)
    return filter
end

-- Build the WebP vf string (fps + optional eq + zoom/pan crop)
local function build_webp_vf(include_eq)
    local fps = source_fps()
    local parts = {}
    local crop = zoom_pan_crop()
    if crop then table.insert(parts, crop) end
    table.insert(parts, string.format("fps=%d", fps))
    if include_eq then
        local eq = build_eq_filter()
        if eq then table.insert(parts, eq) end
    end
    return table.concat(parts, ",")
end

-- Build the GIF vf string (fps + optional eq + palettegen/paletteuse + zoom/pan crop)
local function build_gif_vf(include_eq)
    local fps = source_fps()
    local parts = {}
    local crop = zoom_pan_crop()
    if crop then table.insert(parts, crop) end
    table.insert(parts, string.format("fps=%d", fps))
    if include_eq then
        local eq = build_eq_filter()
        if eq then table.insert(parts, eq) end
    end
    local pre = table.concat(parts, ",")
    -- palettegen+paletteuse pipeline (two-pass inside one vf)
    return string.format("%s,split[s0][s1];[s0]palettegen=max_colors=256[p];[s1][p]paletteuse=dither=%s",
        pre, GIF_DITHER)
end

local function dedup_path(base_path, suffix, ext)
    local candidate = string.format("%s_%s.%s", base_path, suffix, ext)
    local counter = 1
    while true do
        local f = io.open(candidate, "r")
        if not f then return candidate end
        f:close()
        candidate = string.format("%s_%s_%d.%s", base_path, suffix, counter, ext)
        counter = counter + 1
    end
end

-- Build the common ffmpeg prefix (input + seeking)
-- Handles mf:// sources from webp-anim-bridge (animated WebP frame sequences).
local function build_ffmpeg_seek_args(a, b)
    local args = {"ffmpeg", "-y"}
    local file_path = mp.get_property("stream-path") or mp.get_property("path", "")
    local mf_prefix = "mf:///tmp/mpv-webp-anim/"

    if file_path:sub(1, #mf_prefix) == mf_prefix or file_path:find("/tmp/mpv-webp-anim/") then
        -- Source is from webp-anim-bridge: use the extracted PNG frames
        local pattern = file_path
        if pattern:sub(1, 5) == "mf://" then
            pattern = pattern:sub(6)
        end
        local mf_fps = source_fps()
        table.insert(args, "-framerate"); table.insert(args, tostring(mf_fps))
        table.insert(args, "-start_number"); table.insert(args, "0")
        table.insert(args, "-i");           table.insert(args, pattern)
        -- Seek by frame number instead of timestamp
        local start_frame = math.floor(a * mf_fps)
        local end_frame   = math.ceil(b * mf_fps) - 1
        local num_frames  = math.max(1, end_frame - start_frame + 1)
        table.insert(args, "-start_number"); table.insert(args, tostring(start_frame))
        table.insert(args, "-vframes");      table.insert(args, tostring(num_frames))
    else
        -- Normal source: seek by timestamp
        table.insert(args, "-i");    table.insert(args, file_path)
        table.insert(args, "-ss");   table.insert(args, format_ts(a))
        table.insert(args, "-to");   table.insert(args, format_ts(b))
    end
    return args
end

local function format_size(path)
    local info = utils.file_info(path)
    if not info or not info.size then return "" end
    local s = info.size
    if s > 1024 * 1024 then
        return string.format(" (%.1f MB)", s / 1024 / 1024)
    elseif s > 1024 then
        return string.format(" (%.0f KB)", s / 1024)
    else
        return string.format(" (%d B)", s)
    end
end

-- ╔══════════════════════════════════════════╗
-- ║           MAIN EXPORT FUNCTION          ║
-- ╚══════════════════════════════════════════╝

local function export_both()
    -- 1. Read loop points
    local a = mp.get_property_number("ab-loop-a")
    local b = mp.get_property_number("ab-loop-b")

    if not a or not b then
        osd_msg("No A-B loop set. Press l twice to create one.", 3)
        return
    end
    if a > b then a, b = b, a end
    if a == b then
        osd_msg("Loop has zero duration. Set different A and B points.", 3)
        return
    end

    -- 2. Check a file is loaded
    local has_file = mp.get_property("path")
    if not has_file then
        osd_msg("No file loaded.", 3)
        return
    end

    -- 3. Filename stem
    local filename_stripped = mp.get_property("filename/no-ext") or "video"
    local basename = filename_stripped:gsub("[^%w._-]", "_")
    local label_a = format_label(a)
    local label_b = format_label(b)
    local stem = string.format("%s_loop_%s-%s", basename, label_a, label_b)

    local webp_base = utils.join_path(OUTPUT_DIR, stem)
    local gif_base  = utils.join_path(GIF_DIR, stem)

    -- 4. Ensure output directories exist
    mp.command_native({"run", "mkdir", "-p", OUTPUT_DIR, GIF_DIR})

    -- 5. Build job list
    local has_adjustments = build_eq_filter() ~= nil
    local jobs = {}

    local function add_job(base, suffix, ext, vf_string, encoder_fn)
        local out_path = dedup_path(base, suffix, ext)
        local args = build_ffmpeg_seek_args(a, b)
        table.insert(args, "-vf"); table.insert(args, vf_string)
        encoder_fn(args)
        table.insert(args, out_path)
        table.insert(jobs, {args = args, path = out_path, label = ("%s_%s.%s"):format(basename, suffix, ext)})
    end

    -- WebP encoder
    local function webp_enc(args)
        table.insert(args, "-c:v");     table.insert(args, "libwebp_anim")
        table.insert(args, "-lossless"); table.insert(args, WEBP_LOSSLESS and "1" or "0")
        table.insert(args, "-quality");  table.insert(args, string.format("%d", WEBP_QUALITY))
        table.insert(args, "-loop");     table.insert(args, string.format("%d", WEBP_LOOP))
        table.insert(args, "-an")
    end

    -- GIF encoder
    local function gif_enc(args)
        table.insert(args, "-loop"); table.insert(args, string.format("%d", GIF_LOOP))
    end

    -- WebP: normal
    add_job(webp_base, "normal", "webp", build_webp_vf(false), webp_enc)
    -- GIF: normal
    add_job(gif_base, "normal", "gif", build_gif_vf(false), gif_enc)

    if has_adjustments then
        add_job(webp_base, "modified", "webp", build_webp_vf(true), webp_enc)
        add_job(gif_base, "modified", "gif", build_gif_vf(true), gif_enc)
    end

    -- 6. Launch all jobs in parallel
    local total = #jobs
    local completed = {}
    local results = {}

    osd_msg(string.format("Exporting %d loops (WebP + GIF) ...", total), 2)
    for _, j in ipairs(jobs) do
        print(string.format("[export-loop-webp] %s", j.path))
    end

    local function on_all_done()
        local ok_count = 0
        local fail_list = {}
        local size_parts = {}
        for i = 1, total do
            if results[i].success then
                ok_count = ok_count + 1
                table.insert(size_parts, results[i].label .. format_size(results[i].path))
            else
                table.insert(fail_list, results[i].label)
            end
        end

        if ok_count == total then
            osd_msg(string.format("All %d exported: %s", total, table.concat(size_parts, ", ")), 5)
        elseif ok_count > 0 then
            osd_msg(string.format("%d/%d exported. Failed: %s — check console.",
                ok_count, total, table.concat(fail_list, ", ")), 5)
        else
            osd_msg("All exports failed. Check console.", 5)
        end

        for i = 1, total do
            if not results[i].success then
                print(string.format("[export-loop-webp] ERROR [%s]: %s",
                    results[i].label, results[i].error or "unknown"))
            end
        end
    end

    for i, j in ipairs(jobs) do
        mp.command_native({
            name = "subprocess", args = j.args,
            playback_only = false, capture_stdout = false, capture_stderr = false,
        }, function(success, result, err)
            completed[i] = true
            results[i] = {success = success, path = j.path, label = j.label, error = err}
            -- Check if all are done
            for k = 1, total do
                if not completed[k] then return end
            end
            on_all_done()
        end)
    end
end

-- ╔══════════════════════════════════════════╗
-- ║           KEYBINDING REGISTRATION       ║
-- ╚══════════════════════════════════════════╝

mp.add_key_binding(nil, "export", export_both)