-- export-loop.lua
-- Export the current A-B loop to an MP4 via ffmpeg.
-- Frame-accurate (re-encodes). Exports both the raw loop and one with
-- mpv's gamma/brightness/contrast/saturation baked in.
--
-- Suggested input.conf binding:
--   e    script-binding export_loop/export
--
-- If both ab-loop-a and ab-loop-b are set, spawns ffmpeg with libx264
-- and saves to ~/Videos/mpv-loops/.

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

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

local OUTPUT_DIR = os.getenv("HOME") .. "/Videos/mpv-loops"
local ENABLE_OSD_MSG = true
local VIDEO_CRF = 18                  -- lower = better quality, 18 is visually lossless
local VIDEO_PRESET = "ultrafast"      -- encoding speed/compression tradeoff
local VIDEO_MAX_WIDTH = 1920          -- scale to fit this width (auto height)
local VIDEO_CONTAINER = "mp4"         -- output container format

-- ╔══════════════════════════════════════════╗
-- ║           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

-- Map mpv's video adjustment properties to ffmpeg eq filter string.
-- Returns nil if no adjustments are active.
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.
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))
    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 vf string. When include_eq is false, eq filters are omitted.
-- Zoom/pan crop is prepended before scale when active.
local function build_vf_string(include_eq)
    local parts = {}
    local crop = zoom_pan_crop()
    if crop then table.insert(parts, crop) end
    table.insert(parts, string.format("scale='min(%d,iw)':-1:flags=lanczos", VIDEO_MAX_WIDTH))
    if include_eq then
        local eq_str = build_eq_filter()
        if eq_str then
            table.insert(parts, eq_str)
        end
    end
    return table.concat(parts, ",")
end

-- Resolve a non-overwriting path: append _N before extension if needed.
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 ffmpeg arg prefix: input + seeking + vf.
-- Handles mf:// sources from webp-anim-bridge (animated WebP frame sequences).
local function build_ffmpeg_base_args(a, b, vf_string)
    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
        local pattern = file_path
        if pattern:sub(1, 5) == "mf://" then
            pattern = pattern:sub(6)
        end
        local mf_fps = mp.get_property_number("mf-fps", 10)
        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)
        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
        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
    table.insert(args, "-vf"); table.insert(args, vf_string)
    return args
end

-- Format a file size for display.
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_loop()
    -- 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. Build output filenames
    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 base_path = utils.join_path(OUTPUT_DIR, stem)

    local normal_path   = dedup_path(base_path, "normal",   VIDEO_CONTAINER)
    local modified_path = dedup_path(base_path, "modified", VIDEO_CONTAINER)

    -- 4. Ensure output directory exists
    mp.command_native({"run", "mkdir", "-p", OUTPUT_DIR})

    -- 5. Detect whether any adjustments are active
    local has_adjustments = build_eq_filter() ~= nil
    local ext = VIDEO_CONTAINER

    -- 6. Build ffmpeg arg lists
    local normal_vf = build_vf_string(false)
    local normal_args = build_ffmpeg_base_args(a, b, normal_vf)
    -- encoder
    table.insert(normal_args, "-c:v");     table.insert(normal_args, "libx264")
    table.insert(normal_args, "-preset");  table.insert(normal_args, VIDEO_PRESET)
    table.insert(normal_args, "-crf");     table.insert(normal_args, string.format("%d", VIDEO_CRF))
    table.insert(normal_args, "-pix_fmt"); table.insert(normal_args, "yuv420p")
    table.insert(normal_args, "-c:a");     table.insert(normal_args, "copy")
    table.insert(normal_args, normal_path)

    local modified_args
    if has_adjustments then
        local modified_vf = build_vf_string(true)
        modified_args = build_ffmpeg_base_args(a, b, modified_vf)
        table.insert(modified_args, "-c:v");     table.insert(modified_args, "libx264")
        table.insert(modified_args, "-preset");  table.insert(modified_args, VIDEO_PRESET)
        table.insert(modified_args, "-crf");     table.insert(modified_args, string.format("%d", VIDEO_CRF))
        table.insert(modified_args, "-pix_fmt"); table.insert(modified_args, "yuv420p")
        table.insert(modified_args, "-c:a");     table.insert(modified_args, "copy")
        table.insert(modified_args, modified_path)
    end

    -- 7. Launch exports in parallel
    local pending = has_adjustments and 2 or 1
    local successes = {}
    local errors = {}
    local names = {normal_path}

    if pending == 1 then
        -- Only normal (no adjustments — skip making an identical duplicate)
        osd_msg(string.format("Exporting loop (normal) ..."), 2)

        mp.command_native({
            name = "subprocess", args = normal_args,
            playback_only = false, capture_stdout = false, capture_stderr = false,
        }, function(success, result, err)
            if success then
                osd_msg(string.format("Loop exported: %s%s", normal_path, format_size(normal_path)), 4)
                print("[export-loop] Saved: " .. normal_path)
            else
                osd_msg("Export failed. Check console.", 4)
                print("[export-loop] ERROR: " .. (err or "unknown"))
            end
        end)
        return
    end

    names = {normal_path, modified_path}
    osd_msg("Exporting normal and modified ...", 2)
    print(string.format("[export-loop] Normal:   %s", normal_path))
    print(string.format("[export-loop] Modified: %s", modified_path))

    local function on_done()
        if successes[1] and successes[2] then
            osd_msg(string.format("Both exported: %s%s, %s%s",
                names[1], format_size(names[1]),
                names[2], format_size(names[2])), 5)
        elseif successes[1] then
            osd_msg(string.format("Normal exported%s. Modified failed. Check console.",
                format_size(names[1])), 5)
        elseif successes[2] then
            osd_msg(string.format("Modified exported%s. Normal failed. Check console.",
                format_size(names[2])), 5)
        else
            osd_msg("Both exports failed. Check console.", 5)
        end

        local failures_found = false
        for i = 1, 2 do
            if errors[i] then
                print(string.format("[export-loop] ERROR (job %d): %s", i, errors[i]))
                failures_found = true
            end
        end
        if not failures_found then
            print(string.format("[export-loop] Both saved: %s, %s", names[1], names[2]))
        end
    end

    local completed = {false, false}
    for i = 1, 2 do
        local idx = i
        local args = (idx == 1) and normal_args or modified_args
        mp.command_native({
            name = "subprocess", args = args,
            playback_only = false, capture_stdout = false, capture_stderr = false,
        }, function(success, result, err)
            completed[idx] = true
            successes[idx] = success
            if not success then errors[idx] = err or "unknown" end
            if completed[1] and completed[2] then
                on_done()
            end
        end)
    end
end

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

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