MPV Lua scripts to interface with Stash

ordered_table.lua

local ordered_table = {}
--[[
This implementation of ordered table does not hold performance above functionality.
It invokes a metamethod `__newindex` for every access, and
while this is not ideal performance wise, the resulting table behaves very closely to
a standard Lua table, with the quirk that the keys are ordered by when they are first seen
(unless deleted and then reinserted.)
--]]
-- private unique keys
local _values = {}
local _keys = {}

function ordered_table.insert(t, k, v)
  if v == nil then
    ordered_table.remove(t, k)
  else                           -- update/store value
    if t[_values][k] == nil then -- new key
      t[_keys][#t[_keys] + 1] = k
    end
    t[_values][k] = v
  end
end

local function find(t, v)
  for i, val in ipairs(t) do
    if v == val then
      return i
    end
  end
end

function ordered_table.remove(t, k)
  local tv = t[_values]
  local v = tv[k]
  if v ~= nil then
    local tk = t[_keys]
    table.remove(tk, assert(find(tk, k)))
    tv[k] = nil
  end
  return v
end

function ordered_table.pairs(t)
  local i = 0
  return function()
    i = i + 1
    local key = t[_keys][i]
    if key ~= nil then
      return key, t[_values][key]
    end
  end
end

-- set metamethods for ordered_table "class"
ordered_table.__newindex = ordered_table.insert -- called for updates too since we store the value in `_values` instead.
ordered_table.__len = function(t) return #t[_keys] end
ordered_table.__pairs = ordered_table.pairs
ordered_table.__index = function(t, k) return t[_values][k] end -- function so we can share between tables

function ordered_table:new(init)
  init = init or {}
  local key_table = {}
  local value_table = {}
  local t = { [_keys] = key_table,[_values] = value_table }
  local n = #init
  if n % 2 ~= 0 then
    error("key: " .. tostring(init[#init]) .. " is missing value", 2)
  end
  for i = 1, n / 2 do
    local k = init[i * 2 - 1]
    local v = init[i * 2]
    if value_table[k] ~= nil then
      error("duplicated key:" .. tostring(k), 2)
    end
    key_table[#key_table + 1] = k
    value_table[k]            = v
  end
  return setmetatable(t, self)
end

return setmetatable(ordered_table, { __call = ordered_table.new })

--[[ Example Usage:

ordered_table = require"ordered_table"

local t = ordered_table{
  "hello", 1,  -- key, value pairs
  2, 2,
  50, 3,
  "bye", 4,
  200, 5,
  "_values", 20,
  "_keys", 4,
}

print(#t, "items")
print("hello is", t.hello)
print()
for k, v in pairs(t) do print(k, v) end
print()

t.bye = nil -- delete it
t.hello = 0 -- updates
t.bye = 4 -- bring it back!
for k, v in pairs(t) do print(k, v) end
print(#t, "items")

--]]

stash-interface.lua

local msg = require "mp.msg"
local utils = require 'mp.utils'

local ordered_table = require 'ordered_table'

local mod = {
  filters = nil,
  tags = nil,
}

local scene_query =
"query FindScenes($filter: FindFilterType, $scene_filter: SceneFilterType, $scene_ids: [Int!]) { findScenes(filter: $filter, scene_filter: $scene_filter, scene_ids: $scene_ids) {count scenes { id title files { path } rating100 tags { name id aliases } performers { name } } } }"
local filter_query = "query FindSavedFilters($mode: FilterMode) { findSavedFilters(mode: $mode) { name filter } }"
local tag_query = "query allTags { allTags { name id aliases } }"
local scene_update_tags_mutation =
"mutation SceneUpdate($id: ID! $tag_ids: [ID!]) { sceneUpdate(input: {id: $id, tag_ids: $tag_ids}) { id title tags { name id } } }"
local delete_mutation =
"mutation ScenesDestroy($ids: [ID!]!) { scenesDestroy(input: {ids: $ids, delete_file: true, delete_generated: true} ) }"
local organized_mutation =
"mutation SceneUpdate($input: SceneUpdateInput!) { sceneUpdate(input: $input) { id organized } }"
local rating_mutation = "mutation SceneUpdate($input: SceneUpdateInput!) { sceneUpdate(input: $input) { id rating100 } }"
local increment_o_mutation = "mutation IncrementO($id: ID!) { sceneIncrementO(id: $id) }"

local function call_graphql(request)
  local json = utils.format_json(request)

  local command = { 'curl', '127.0.0.1:9999/graphql', '-X', 'POST', '-H', 'content-type: application/json', '-d', json }

  local response = mp.command_native({
    name = 'subprocess',
    args = command,
    capture_stdout = true,
    capture_stderr = true,
    playback_only = false,
  })

  if response == nil then return end

  if response.status ~= 0 then
    msg.error(response.stderr);
    return nil
  end

  local result = utils.parse_json(response.stdout)
  if result.errors ~= nil then
    for _, error in pairs(result.errors) do
      msg.error(error["message"])
    end

    return nil
  end
  return result.data
end

local function trim(s)
  return (s:gsub("^%s*(.-)%s*$", "%1"))
end

local function multicriterion_converter(filter)
  local values = {}
  local res = {
    value = values,
    modifier = filter.modifier
  }
  for _, value in pairs(filter.value.items) do
    table.insert(values, value.id)
  end
  if filter.value.depth ~= nil then
    res["depth"] = filter.value.depth
  end
  return res
end

local function boolean_converter(filter)
  return string.lower(filter.value) == "true"
end

local function criterion_converter(filter)
  local res = {
    modifier = filter.modifier,
    value = ""
  }
  if filter.value ~= nil then
    res["value"] = filter.value.value
    if filter.value.value2 ~= nil then
      res["value2"] = filter.value.value2
    end
  end
  return res
end

local function string_converter(filter)
  return filter.value
end

local function field_converter(field_name)
  -- Missing, in order of importance:
  -- DateCriterion
  -- TimestampCriterion

  if field_name == "organized"
      or field_name == "performer_favorite"
      or field_name == "interactive" then
    return boolean_converter
  end

  if field_name == "tags"
      or field_name == "performers"
      or field_name == "performer_tags"
      or field_name == "studios"
  then
    return multicriterion_converter
  end

  if field_name == "duration"
      or field_name == "file_count"
      or field_name == "interactive_speed"
      or field_name == "o_counter"
      or field_name == "performer_age"
      or field_name == "performer_count"
      or field_name == "play_count"
      or field_name == "play_duration"
      or field_name == "rating"
      or field_name == "rating100"
      or field_name == "resume_time"
      or field_name == "tag_count"
  then
    return criterion_converter
  end

  if field_name == "sceneIsMissing" or field_name == "hasMarkers" then
    return string_converter
  end

  if field_name == "title"
      or field_name == "code"
      or field_name == "details"
      or field_name == "director"
      or field_name == "oshash"
      or field_name == "checksum"
      or field_name == "phash"
      or field_name == "path"
      or field_name == "stash_id"
      or field_name == "url"
      or field_name == "captions"
  then
    return criterion_converter
  end
end

-- The filters are stored in the format used by the TypeScript
-- frontend, but we need them in GraphQL format: this translates them
local function filter_to_query(filter_string)
  local parsed_filter = utils.parse_json(filter_string)

  local scene_filter = {}
  setmetatable(scene_filter, getmetatable(parsed_filter))

  for _, v in ipairs(parsed_filter.c) do
    local filter = utils.parse_json(v)
    local field_name = filter.type
    local converter = field_converter(field_name)
    scene_filter[field_name] = converter(filter)
  end

  return {
    query = scene_query,
    variables = {
      filter = {
        page = 1,
        per_page = parsed_filter.perPage,
        sort = parsed_filter.sortby,
        direction = parsed_filter.sortdir:upper(),
      },
      scene_filter = scene_filter
    }
  }
end

local function search_filter(str)
  return {
    query = scene_query,
    variables = {
      filter = {
        q = str,
        page = 1,
        per_page = 40,
        sort = "date",
        direction = "DESC",
      },
    }
  }
end

local function generate_random_seed()
  math.randomseed(os.time())
  math.random()
  local new_seed = string.format("random_%d", math.random(10000000, 99999999))
  return new_seed
end

function mod.get_filters()
  -- Filters are cached per session
  if mod.filters ~= nil then
    return mod.filters
  end

  local raw_filters = call_graphql({
    query = filter_query,
    variables = {
      mode = "SCENES"
    }
  })

  mod.filters = ordered_table {}

  if raw_filters == nil then
    msg.error("Failed to fetch filters")
    return mod.filters
  end

  for _, filter in pairs(raw_filters["findSavedFilters"]) do
    local name = filter["name"]
    -- Replace static random seed with a new, very random seed
    local vars = filter["filter"]:gsub("random_%d+", generate_random_seed())
    msg.debug(string.format("Found filter '%s': %s", name, vars))

    ordered_table.insert(mod.filters, name, filter_to_query(vars))
  end

  return mod.filters
end

local function map(tbl, func)
  local t = {}
  for k, v in pairs(tbl) do
    t[k] = func(v)
  end
  return t
end

local function flatten(tbl)
  return map(tbl, function(v) return v.name end)
end

local function first(tbl)
  for _, v in pairs(tbl) do
    return v
  end
end

local function to_tag(tag)
  -- The next three lines are only necessary because I format my tags in a special way
  -- but it should still work with plain tag names
  local category, localDisplayName, lastPos = tag.name:match("(%a+)%d*%.?%s*([%a%s]+)()")
  local stashDbName = tag.name:sub(lastPos + 1, -2)
  local aliases = tag.aliases or {}
  return {
    name = trim(localDisplayName),
    aliases = aliases,
    id = tag.id,
    searchName = tag.name .. " " .. table.concat(aliases, " "),
  }
end

-- Returns a map of paths to scenes, scenes contain title and id
local function scenes_to_info(raw_scenes)
  local scene_map = ordered_table {}
  if raw_scenes == nil then return scene_map end

  for _, v in pairs(raw_scenes.findScenes.scenes) do
    local tags = table.concat(flatten(map(v.tags, to_tag)), ", ")
    local performers = table.concat(flatten(v.performers), ", ")
    local path = first(v.files).path
    local rating = (v.rating100 or 0)
    ordered_table.insert(scene_map, path, {
      id = v.id,
      title = v.title,
      rating = rating,
      tags = tags,
      fullTags = v.tags,
      performers = performers,
    })
  end
  return scene_map
end

function mod.get_scenes_by_ids(ids)
  -- split the ids string into an array of ints
  local ids_array = {}
  for id in string.gmatch(ids, "%d+") do
    table.insert(ids_array, tonumber(id))
  end

  local request = {
    query = scene_query,
    variables = {
      scene_ids = ids_array
    }
  }

  local result = call_graphql(request)
  return scenes_to_info(result)
end

function mod.from_filter(filter)
  return scenes_to_info(call_graphql(filter))
end

function mod.search_scenes(query)
  return mod.from_filter(search_filter(query))
end

function mod.mark_organized(scene)
  local result = call_graphql({
    query = organized_mutation,
    variables = {
      input = {
        id = scene.id,
        organized = true
      }
    }
  })
  return result
end

-- Adds a tag to the given scene
function mod.add_tag(scene, tag)
  if mod.tags == nil then
    local fresh_tags = call_graphql({
      query = tag_query
    }).allTags
    mod.tags = map(fresh_tags, to_tag)
  end

  local actualTag = nil
  local new_tag_ids = {}
  for _, v in pairs(scene.fullTags) do
    table.insert(new_tag_ids, v.id)
  end
  for _, v in pairs(mod.tags) do
    if v.name:lower() == tag:lower() then
      actualTag = v.name
      table.insert(new_tag_ids, v.id)
      break
    end
  end

  local result = call_graphql({
    query = scene_update_tags_mutation,
    variables = {
      id = scene.id,
      tag_ids = new_tag_ids
    }
  })
  if result == nil then return nil end
  return actualTag
end

-- Sets the given scene rating on a scale from 0 to 5
function mod.set_rating(scene, rating)
  local result = call_graphql({
    query = rating_mutation,
    variables = {
      input = {
        id = scene.id,
        rating100 = rating * 20
      }
    }
  })
  if result == nil then return nil end
  return result.sceneUpdate.rating100
end

-- Delete a given scene through Stash
function mod.mark_deleted(scene)
  local result = call_graphql({
    query = delete_mutation,
    variables = {
      ids = { scene.id },
    }
  })
  return result
end

-- Increment O-counter for given scene
function mod.increment_O(scene)
  return call_graphql({
    query = increment_o_mutation,
    variables = { id = scene.id }
  })
end

return mod

stash.lua

package.path = mp.command_native({ "expand-path", "~~/script-modules/?.lua;" }) .. package.path
local opts = require 'mp.options'
local util = require 'mp.utils'
local ordered_table = require 'ordered_table'
local uin = require 'user-input-module'
local stash = require 'stash-interface'

local options = {
    ids = "",
}
opts.read_options(options)

local move_list = {}

local function move_to_stash()
    local finalpath = "F:\\Porn\\_unsorted\\"
    local args = { "cmd", "/c", "move", "/Y" }
    for _, path in pairs(move_list) do
        table.insert(args, path)
    end
    table.insert(args, finalpath)

    util.subprocess_detached({ args = args })
end

local function count(t)
    local total = 0
    for _ in pairs(t) do total = total + 1 end
    return total
end

local function remove_currently_playing()
    local path = mp.get_property("path")
    if path == nil then return end

    Scenes[path] = nil
    mp.commandv("playlist-remove", "current")
end

local function organize_in_stash()
    local path = mp.get_property("path")
    if path == nil then return end

    if Scenes == nil then
        mp.osd_message("Moving to stash folder")
        local dirname, _ = util.split_path(path)
        if dirname ~= "D:\\Downloads\\sorted\\" then return end
        table.insert(move_list, path)
        mp.commandv("playlist-remove", "current")
        return
    end

    local currently_playing = Scenes[path]
    if currently_playing == nil then
        mp.osd_message("Not organizing: scene not fetched from stash")
        return
    end

    mp.osd_message("Marking file as organized")
    stash.mark_organized(currently_playing)
    remove_currently_playing()
end

local function delete_from_stash()
    local path = mp.get_property("path")
    if path == nil then return end

    if Scenes == nil then
        mp.commandv("script-message", "delete_file")
        mp.commandv("playlist-remove", "current")
        return
    end

    local currently_playing = Scenes[path]
    if currently_playing == nil then
        mp.osd_message("Not deleting: scene not fetched from stash")
        return
    end

    mp.osd_message("Marking file for deletion")
    remove_currently_playing()
    stash.mark_deleted(currently_playing)
end

local function get_stars(rating)
    if rating == 0 then return "" end
    return string.format("[%.1f]", rating / 10)
end

local function set_window_title()
    local path = mp.get_property("path")
    if path == nil then return end

    if Scenes == nil or Scenes[path] == nil then
        mp.set_property("force-media-title", "")
        return
    end

    local currently_playing = Scenes[path]
    local current = mp.get_property("force-media-title")
    local sub = current:gmatch("%[([^%]]*)%]")()

    if sub ~= nil then
        mp.set_property("force-media-title", string.format("[%s] %s", sub, currently_playing.title))
    else
        mp.set_property("force-media-title", currently_playing.title)
    end
end

local function show_details()
    local path = mp.get_property("path")
    if path == nil then return end

    if Scenes == nil or Scenes[path] == nil then
        local filename = mp.get_property("filename/no-ext")
        mp.osd_message(filename, 5)
        return
    end

    local currently_playing = Scenes[path]
    local rating_stars = get_stars(currently_playing.rating)
    local title = currently_playing.title
    if title == "" then title = mp.get_property("filename/no-ext") end
    local details = string.format("%s %s\n%s\n%s", title, rating_stars, currently_playing.performers,
        currently_playing.tags)
    mp.osd_message(details, 5)
end

local function open_in_stash()
    local path = mp.get_property("path")
    if path == nil then return end

    if Scenes == nil then return end

    local currently_playing = Scenes[path]
    if currently_playing == nil then return end

    local cmd = string.format("start http://localhost:9999/scenes/%s", currently_playing.id)
    mp.commandv("run", "cmd", "/c", cmd)
end

local function replace_currently_playing_scenes(new_title)
    if new_title ~= nil then
        mp.set_property("force-media-title", string.format("[%s]", new_title))
    end

    mp.commandv("playlist-clear")
    mp.commandv("playlist-remove", "current")

    for path, _ in ordered_table.pairs(Scenes) do
        mp.commandv("loadfile", path, "append")
    end

    mp.commandv("playlist-play-index", 0)
end

local function search_scenes(query)
    if query == nil or query == "" then
        mp.osd_message("No query provided", 5)
        return
    end
    local results = stash.search_scenes(query)
    local total = count(results)
    if total > 0 then
        mp.osd_message(string.format("Search for '%s' returned %d results", query, total), 5)
        Scenes = results
        replace_currently_playing_scenes(query)
    else
        mp.osd_message(string.format("Search for '%s' returned no results", query), 5)
    end
end

local function search_stash()
    mp.commandv("script-message-to", "autoload", "disable-autoload")
    uin.get_user_input(search_scenes, {
        request_text = "Enter Stash search query:",
        replace = true,
    })
end

local function scenes_from_filter(filter)
    mp.commandv("script-message-to", "autoload", "disable-autoload")
    if filter == nil then
        mp.osd_message("No filter provided", 5)
        return
    end
    local filters = stash.get_filters()
    if filters == nil then
        mp.osd_message("No filters found", 5)
        return
    end
    local wantedFilter = filters[filter]
    if wantedFilter == nil then
        mp.osd_message(string.format("Filter '%s' not found", filter), 5)
        return
    end

    local results = stash.from_filter(wantedFilter)
    local total = count(results)
    if total > 0 then
        mp.osd_message(string.format("Filter '%s' returned %d results", filter, total), 5)
        Scenes = results
        replace_currently_playing_scenes(filter)
    else
        mp.osd_message(string.format("Filter '%s' returned no results", filter), 5)
    end
end

local function pick_filter()
    uin.get_user_input(scenes_from_filter, {
        request_text = "Enter filter name:",
        default_input = "Unorganized",
        cursor_pos = 12,
        replace = true,
    })
end

local function increment_O()
    local path = mp.get_property("path")
    if path == nil then return end

    if Scenes == nil then return end

    local currently_playing = Scenes[path]
    if currently_playing == nil then return end

    local result = stash.increment_O(currently_playing)
    if result ~= nil then
        mp.osd_message(string.format("New count: %d", result.sceneIncrementO), 5)
    else
        mp.osd_message("Failed to increment O", 5)
    end
end

local function set_rating(rating)
    local path = mp.get_property("path")
    if path == nil then return end

    if Scenes == nil then return end

    local currently_playing = Scenes[path]
    if currently_playing == nil then return end

    local result = stash.set_rating(currently_playing, rating)
    if result ~= nil then
        currently_playing.rating = result
        mp.osd_message(string.format("New rating: %d", result), 5)
    else
        mp.osd_message("Failed to set rating", 5)
    end
end

local function add_tag(tag)
    local path = mp.get_property("path")
    if path == nil then return end

    if Scenes == nil then return end

    local currently_playing = Scenes[path]
    if currently_playing == nil then return end

    local result = stash.add_tag(currently_playing, tag)
    if result ~= nil then
        mp.osd_message(string.format("Added tag '%s'", result), 5)
    else
        mp.osd_message("Failed to add tag", 5)
    end
end

mp.add_key_binding("x", "show_scene_details", show_details)
mp.add_key_binding("shift+x", "open_in_stash", open_in_stash)
mp.add_key_binding("shift+z", "delete_from_stash", delete_from_stash)
mp.add_key_binding("shift+c", "organize", organize_in_stash)
mp.add_key_binding("shift+r", "increment_o", increment_O)

mp.add_key_binding("1", "set_rating_1", function() set_rating(1) end)
mp.add_key_binding("2", "set_rating_2", function() set_rating(2) end)
mp.add_key_binding("3", "set_rating_3", function() set_rating(3) end)
mp.add_key_binding("4", "set_rating_4", function() set_rating(4) end)
mp.add_key_binding("5", "set_rating_5", function() set_rating(5) end)
mp.add_key_binding("6", "set_rating_0", function() set_rating(0) end)
mp.add_key_binding("ctrl+a", "add_anal_tag", function() add_tag("anal") end)


mp.add_key_binding("ctrl+f", "pick_filter", pick_filter)
mp.add_key_binding("ctrl+s", "search_stash", search_stash)

mp.register_event("file-loaded", set_window_title)
mp.register_event("shutdown", move_to_stash)

-- This means the script was started with a list of stash ids
-- so we just load them right away
if options.ids ~= "" then
    mp.commandv("script-message-to", "autoload", "disable-autoload")
    Scenes = stash.get_scenes_by_ids(options.ids)
    replace_currently_playing_scenes()
    mp.commandv("playlist-play-index", 0)
end

user-input-module.lua

--[[
    This is a module designed to interface with mpv-user-input
    https://github.com/CogentRedTester/mpv-user-input

    Loading this script as a module will return a table with two functions to format
    requests to get and cancel user-input requests. See the README for details.

    Alternatively, developers can just paste these functions directly into their script,
    however this is not recommended as there is no guarantee that the formatting of
    these requests will remain the same for future versions of user-input.
]]
local API_VERSION = "0.1.0"

local mp = require 'mp'
local msg = require "mp.msg"
local utils = require 'mp.utils'
local mod = {}

local name = mp.get_script_name()
local counter = 1

local function pack(...)
  local t = { ... }
  t.n = select("#", ...)
  return t
end

local request_mt = {}

-- ensures the option tables are correctly formatted based on the input
local function format_options(options, response_string)
  return {
    response = response_string,
    version = API_VERSION,
    id = name .. '/' .. (options.id or ""),
    source = name,
    request_text = ("[%s] %s"):format(options.source or name,
      options.request_text or options.text or "requesting user input:"),
    default_input = options.default_input,
    cursor_pos = tonumber(options.cursor_pos),
    queueable = options.queueable and true,
    replace = options.replace and true
  }
end

-- cancels the request
function request_mt:cancel()
  assert(self.uid, "request object missing UID")
  mp.commandv("script-message-to", "user_input", "cancel-user-input/uid", self.uid)
end

-- updates the options for the request
function request_mt:update(options)
  assert(self.uid, "request object missing UID")
  options = utils.format_json(format_options(options))
  mp.commandv("script-message-to", "user_input", "update-user-input/uid", self.uid, options)
end

-- sends a request to ask the user for input using formatted options provided
-- creates a script message to recieve the response and call fn
function mod.get_user_input(fn, options, ...)
  options = options or {}
  local response_string = name .. "/__user_input_request/" .. counter
  counter = counter + 1

  local request = {
    uid = response_string,
    passthrough_args = pack(...),
    callback = fn,
    pending = true
  }

  -- create a callback for user-input to respond to
  mp.register_script_message(response_string, function(response)
    mp.unregister_script_message(response_string)
    request.pending = false

    response = utils.parse_json(response)
    request.callback(response.line, response.err, unpack(request.passthrough_args, 1, request.passthrough_args.n))
  end)

  -- send the input command
  options = utils.format_json(format_options(options, response_string))
  mp.commandv("script-message-to", "user_input", "request-user-input", options)

  return setmetatable(request, { __index = request_mt })
end

-- runs the request synchronously using coroutines
-- takes the option table and an optional coroutine resume function
function mod.get_user_input_co(options, co_resume)
  local co, main = coroutine.running()
  assert(not main and co, "get_user_input_co must be run from within a coroutine")

  local uid = {}
  local request = mod.get_user_input(function(line, err)
    if co_resume then
      co_resume(uid, line, err)
    else
      local success, er = coroutine.resume(co, uid, line, err)
      if not success then
        msg.warn(debug.traceback(co))
        msg.error(er)
      end
    end
  end, options)

  -- if the uid was not sent then the coroutine was resumed by the user.
  -- we will treat this as a cancellation request
  local success, line, err = coroutine.yield(request)
  if success ~= uid then
    request:cancel()
    request.callback = function()
    end
    return nil, "cancelled"
  end

  return line, err
end

-- sends a request to cancel all input requests with the given id
function mod.cancel_user_input(id)
  id = name .. '/' .. (id or "")
  mp.commandv("script-message-to", "user_input", "cancel-user-input/id", id)
end

return mod