// ==UserScript==
// @name         Plex downloader
// @description  Adds a download button the Plex desktop interface. Works on episodes, movies, whole seasons, and entire shows.
// @author       Mow
// @version      1.0
// @license      MIT
// @grant        none
// @match        https://app.plex.tv/desktop/*
// @run-at       document-start
// @namespace https://greasyfork.org/users/1260133
// @downloadURL none
// ==/UserScript==


// This code is a heavy modification of the existing PlxDwnld project
// https://sharedriches.com/plex-scripts/piplongrun/




const metadataIdRegex  = new RegExp("key=%2Flibrary%2Fmetadata%2F(\\d+)");
const clientIdRegex    = new RegExp("server\/([a-f0-9]{40})\/");


const logPrefix = "[USERJS Plex Downloader]";
const domPrefix = "USERJSINJECTED_";

const xmlParser = new DOMParser();

const playBtnSelector = "button[data-testid=preplay-play]";


// Server idenfitiers and their respective data loaded over API request
const serverData = {
	servers : {
		// Example data
		/*
		"fd174cfae71eba992435d781704afe857609471b" : {
			"baseUri"     : "https://1-1-1-1.e38c3319c1a4a0f67c5cc173d314d74cb19e862b.plex.direct:13100",
			"accessToken" : "fH5dn-HgT7Ihb3S-p9-k",
			"mediaData"   : {}
		}
		*/
	},
	
	// Promise for loading server data, ensure it is loaded before we try to pull data
	promise : false
};

// Merge new data object into serverData
function updateServerData(newData, serverDataScope) {
	if (!serverDataScope) {
		serverDataScope = serverData;
	}
	
	for (let key in newData) {
		if (!serverDataScope.hasOwnProperty(key)) {
			serverDataScope[key] = newData[key];
		} else {
			if (typeof newData[key] === "object") {
				updateServerData(newData[key], serverDataScope[key]);
			} else {
				serverDataScope[key] = newData[key];
			}
		}
	}
}


// Should not be visible in normal operation
function errorHandle(msg) {
	console.log(logPrefix + " " + msg.toString());
}


// Async fetch XML and return parsed body
async function fetchXml(url) {
	const response     = await fetch(url);
	const responseText = await response.text();
	const responseXml  = xmlParser.parseFromString(responseText, "text/xml");
	return responseXml;
}

// Async fetch JSON and return parsed body
async function fetchJSON(url) {
	const response     = await fetch(url, { headers : { accept : "application/json" } });
	const responseJSON = await response.json();
	return responseJSON;
}


// Async load server information for this user account from plex.tv api
async function loadServerData() {
	// Ensure access token
	if (!localStorage.hasOwnProperty("myPlexAccessToken")) {
		errorHandle("Cannot find a valid access token (localStorage Plex token missing).");
		return;
	}
	
	const apiResourceUrl = "https://plex.tv/api/resources?includeHttps=1&X-Plex-Token={token}";
	const resourceXml = await fetchXml(apiResourceUrl.replace("{token}", localStorage["myPlexAccessToken"]));
	
	const serverInfoXPath  = "//Device[@provides='server']";
	const servers = resourceXml.evaluate(serverInfoXPath, resourceXml, null, XPathResult.ORDERED_NODE_ITERATOR_TYPE, null);
	// Stupid ugly iterator pattern. Yes this is how you're supposed to do this
	// https://developer.mozilla.org/en-US/docs/Web/API/XPathResult/iterateNext
	let server;
	while (server = servers.iterateNext()) {
		const clientId    = server.getAttribute("clientIdentifier");
		const accessToken = server.getAttribute("accessToken");
		if (!clientId || !accessToken) {
			errorHandle("Cannot find valid server information (missing ID or token in API response).");
			continue;
		}
		
		const connectionXPath  = "//Connection[@local='0']";
		const conn = resourceXml.evaluate(connectionXPath, server, null, XPathResult.FIRST_ORDERED_NODE_TYPE, null);
		if (!conn.singleNodeValue || !conn.singleNodeValue.getAttribute("uri")) {
			errorHandle("Cannot find valid server information (no connection data for server " + clientId + ").");
			continue;
		}
		
		const baseUri = conn.singleNodeValue.getAttribute("uri");
		
		updateServerData({
			servers : {
				[clientId] : {
					baseUri     : baseUri,
					accessToken : accessToken,
					mediaData   : {},
				}
			}
		});
	}
}


// Merge video node data from API response into the serverData media cache
async function updateServerDataMedia(clientId, videoNode) {
	await serverData.promise;
	
	const key         = videoNode.Media[0].Part[0].key;
	const baseUri     = serverData.servers[clientId].baseUri;
	const accessToken = serverData.servers[clientId].accessToken
	
	if (!serverData.servers.hasOwnProperty(clientId)) {
		errorHandle("Cannot find valid server information (no data for clientId).");
		return;
	}
	
	// Build download URL using file path
	const downloadUrl = "{baseuri}{key}?X-Plex-Token={token}&download=1";
	const dlurl = downloadUrl.replace("{baseuri}", baseUri)
							 .replace("{key}",     key)
							 .replace("{token}",   accessToken);
	
	updateServerData({
		servers : {
			[clientId] : {
				mediaData : {
					[videoNode.ratingKey] : {
						dlurl : dlurl,
					}
				}
			}
		}
	});
}


// Pull API response for this media item and handle parents/grandparents
async function fetchMediaData(clientId, metadataId) {
	
	// Make sure server data has loaded in
	await serverData.promise;
	
	// If no data for this server, something has gone wrong
	if (!serverData.servers.hasOwnProperty(clientId)) {
		errorHandle("Cannot find valid server information (no data for clientId).");
		return;
	}
	
	// Should already have an entry for this metadataId
	if (!serverData.servers[clientId].mediaData.hasOwnProperty(metadataId)) {
		serverData.servers[clientId].mediaData[metadataId] = {
			promise : Promise.resolve()
		};
	} else if (serverData.servers[clientId].mediaData[metadataId].promise.state === "resolved") {
		// Already have this media item in cache
		return;
	}
	
	// Get access token and base URI for this server
	const baseUri     = serverData.servers[clientId].baseUri;
	const accessToken = serverData.servers[clientId].accessToken;
	
	// Request library data from this server using metadata ID
	const apiLibraryUrl = "{baseuri}/library/metadata/{id}?X-Plex-Token={token}";
	const libraryUrl = apiLibraryUrl.replace("{baseuri}", baseUri)
	                                .replace("{id}",      metadataId)
	                                .replace("{token}",   accessToken);
	
	const libraryJSON = await fetchJSON(libraryUrl);
	
	// Determine if this is media or just a parent to media
	let leafCount = false;
	if (libraryJSON.MediaContainer.Metadata[0].hasOwnProperty("leafCount")) {
		leafCount = libraryJSON.MediaContainer.Metadata[0].leafCount;
	}
	
	let childCount = false;
	if (libraryJSON.MediaContainer.Metadata[0].hasOwnProperty("childCount")) {
		childCount = libraryJSON.MediaContainer.Metadata[0].childCount;
	}
	
	if (leafCount || childCount) {
		// This is a group media item (show, season)
		
		// Get all of its children, either by leaves or children directly
		let childrenUrl;
		if (childCount && leafCount && (childCount !== leafCount)) {
			const apiLeavesUrl = "{baseuri}/library/metadata/{id}/allLeaves?X-Plex-Token={token}";
			childrenUrl = apiLeavesUrl.replace("{baseuri}", baseUri)
		                              .replace("{id}",      metadataId)
		                              .replace("{token}",   accessToken);
		} else {
			const apiChildrenUrl = "{baseuri}/library/metadata/{id}/children?X-Plex-Token={token}";
			childrenUrl = apiChildrenUrl.replace("{baseuri}", baseUri)
		                                .replace("{id}",      metadataId)
		                                .replace("{token}",   accessToken);
		}
		
		const childrenJSON = await fetchJSON(childrenUrl);
		const childVideoNodes = childrenJSON.MediaContainer.Metadata;
		
		// Iterate over the children of this media item and gather their data
		updateServerData({
			servers : {
				[clientId] : {
					mediaData : {
						[metadataId] : {
							children : [],
						}
					}
				}
			}
		});
		
		for (let i = 0; i < childVideoNodes.length; i++) {
			childMetadataId = childVideoNodes[i].ratingKey;
			await updateServerDataMedia(clientId, childVideoNodes[i]);
			
			serverData.servers[clientId].mediaData[metadataId].children.push(childMetadataId);
			
			// Copy promise to child
			serverData.servers[clientId].mediaData[childMetadataId].promise = serverData.servers[clientId].mediaData[metadataId].promise;
		}
		
	} else {
		// This is a regular media item (episode, movie)
		const videoNode = libraryJSON.MediaContainer.Metadata[0];
		await updateServerDataMedia(clientId, videoNode);
	}
}


// Start fetching a media item from the URL parameters
function initFetchMediaData() {
	// Get client ID for current server
	const clientIdMatch = clientIdRegex.exec(window.location.href);
	if (!clientIdMatch || clientIdMatch.length !== 2) {
		return;
	}
	
	// Get metadata ID for current media item
	const metadataIdMatch = metadataIdRegex.exec(window.location.href);
	if (!metadataIdMatch || metadataIdMatch.length !== 2) {
		return;
	}
	
	// Get rid of extra regex matches
	const clientId   = clientIdMatch[1];
	const metadataId = metadataIdMatch[1];
	
	updateServerData({
		servers: {
			[clientId] : {
				mediaData : {
					[metadataId] : {
						promise : fetchMediaData(clientId, metadataId)
					}
				}
			}
		}
	});
}


// Initiate a download of a URI using iframes
function downloadUri(uri) {
	let frame = document.createElement("frame");
	frame.name = domPrefix + "downloadFrame";
	frame.style = "display: none; !important";
	document.body.appendChild(frame);
	frame.src = uri;
}

// Clean up old DOM elements from previous downloads, if needed
function cleanUpOldDownloads() {
	// There is no way to detect when the download dialog is closed, so just clean up here to prevent DOM clutter
	let oldFrames = document.getElementsByName(domPrefix + "downloadFrame");
	while (oldFrames.length != 0) {
		oldFrames[0].parentNode.removeChild(oldFrames[0]);
	}
}

// Download a media item, handling parents/grandparents
function downloadMedia(clientId, metadataId) {
	// Should not need to wait for these when this is called
	//await serverData.promise;
	//await serverData.servers[clientId].mediaData[metadataId].promise();
	
	if (serverData.servers[clientId].mediaData[metadataId].hasOwnProperty("children")) {
		for (let i = 0; i < serverData.servers[clientId].mediaData[metadataId].children.length; i++) {
			let childId = serverData.servers[clientId].mediaData[metadataId].children[i];
			downloadMedia(clientId, childId);
		}
	} else {
		let uri = serverData.servers[clientId].mediaData[metadataId].dlurl;
		downloadUri(uri);
	}
}


// Create and add the new DOM elements, return an object with references to them
function modifyDom(injectionPoint) {
	// Steal CSS from the inection point element by copying its class name
	const downloadButton = document.createElement(injectionPoint.tagName);
	downloadButton.id = domPrefix + "DownloadButton";
	downloadButton.textContent = "Download";
	downloadButton.className = domPrefix + "element" + " " + injectionPoint.className;
	downloadButton.style.fontWeight = "bold";
	
	// Starts disabled
	downloadButton.style.opacity = 0.5;
	downloadButton.disabled = true;
	
	injectionPoint.after(downloadButton);
	
	return downloadButton;
}


// Activate DOM elements and hook clicking with function
async function domCallback(domElement, clientId, metadataId) {
	
	// Make sure server data has loaded in
	await serverData.promise;
	
	// Make sure we have media data for this item
	await serverData.servers[clientId].mediaData[metadataId].promise;
	
	const downloadFunction = function() {
		cleanUpOldDownloads();
		downloadMedia(clientId, metadataId);
	}
	
	domElement.addEventListener("click", downloadFunction);
	domElement.disabled = false;
	domElement.style.opacity = 1;
}


function checkStateAndRun() {
	// We detect the prescence of the play button (and absence of our injected button) after each page mutation
	const playBtn = document.querySelector(playBtnSelector);  
	if (!playBtn) return;
	if (playBtn.nextSibling.id.startsWith(domPrefix)) return;
	
	// Get client ID for current server
	const clientIdMatch = clientIdRegex.exec(window.location.href);
	if (!clientIdMatch || clientIdMatch.length !== 2) {
		return;
	}
	
	// Get metadata ID for current media item
	const metadataIdMatch = metadataIdRegex.exec(window.location.href);
	if (!metadataIdMatch || metadataIdMatch.length !== 2) {
		return;
	}
	
	// Get rid of extra regex matches
	const clientId   = clientIdMatch[1];
	const metadataId = metadataIdMatch[1];
	
	let domElement = modifyDom(playBtn);
	domCallback(domElement, clientId, metadataId);
}


(async function() {
	// Begin loading server data immediately
	serverData.promise = loadServerData();
	
	// Try to eager load media info
	initFetchMediaData();
	window.addEventListener("hashchange", initFetchMediaData);
	
	// Use a mutation observer to detect pages loading in
	const mo = new MutationObserver(checkStateAndRun);
	mo.observe(document.documentElement, { childList: true, subtree: true });
})();