tpbproxy/controller/search.js

'use strict';

const conf = require('>/conf');
const tmdb = require('>/controller/tmdb');
const { tmdbSearch, tmdbDetails } = tmdb;
const library = require('>/controller/library');
const ollama = require('>/controller/ollama');

// TPB numeric categories that count as Movies / TV( used for the q.php search and to
// keep the LLM from ever seeing non-video results).
const MOVIE_CATS = [201, 202, 207, 209, 211];
const TV_CATS = [205, 208, 212];

function humanSize(bytes){
	bytes = Number(bytes) || 0;
	let units = ['B', 'KB', 'MB', 'GB', 'TB'];
	let i = 0;
	while(bytes >= 1024 && i < units.length - 1){
		bytes /= 1024;
		i++;
	}
	return `${bytes.toFixed(bytes < 10 && i > 0 ? 1 : 0)} ${units[i]}`;
}

async function fetchJSON(url, options){
	let res = await fetch(url, options);
	if(!res.ok){
		let error = new Error(`UpstreamError`);
		error.message = `Request to ${url.split('?')[0]} failed( ${res.status})`;
		error.status = 502;
		throw error;
	}
	return await res.json();
}

// Turn a raw TPB category id into 'movie' | 'tv' | null.
function videoKind(category){
	category = parseInt(category, 10);
	if(TV_CATS.includes(category)) return 'tv';
	if(MOVIE_CATS.includes(category) || (category >= 200 && category < 300)) return 'movie';
	return null;
}

// --- TPB HTML search ------------------------------------------------------

const SIZE_UNITS = { B: 1, KIB: 1024, MIB: 1024 ** 2, GIB: 1024 ** 3, TIB: 1024 ** 4 };

function sizeToBytes(value, unit){
	return Math.round(parseFloat(value) * (SIZE_UNITS[String(unit).toUpperCase()] || 1));
}

function decodeEntities(text){
	return String(text)
		.replace(/&amp;/g, '&').replace(/&lt;/g, '<').replace(/&gt;/g, '>')
		.replace(/&quot;/g, '"').replace(/&#0?39;/g, "'").replace(/&apos;/g, "'");
}

// Parse the classic TPB HTML result table( category cell, details link name, magnet,
// then right-aligned size / seeders / leechers cells) into apibay-shaped objects.
function parseTPBHtml(html){
	let out = [];
	for(let row of html.split(/<tr/i)){
		let hash = row.match(/magnet:\?xt=urn:btih:([A-Fa-f0-9]+)/i);
		if(!hash) continue;

		let cat = row.match(/\/browse\/(\d+)/);
		let name = row.match(/title="Details for ([^"]+)"/i);
		let size = row.match(/<td align="right">\s*([\d.]+)(?:&nbsp;|\s)+(GiB|MiB|KiB|TiB)\s*<\/td>/i);
		let nums = [...row.matchAll(/<td align="right">\s*([\d,]+)\s*<\/td>/gi)].map(m => Number(m[1].replace(/,/g, '')));

		if(!name) continue;
		// nums are [size, seeders, leechers]; size cell also matched the regex above so
		// seeders/leechers are the last two numeric right-aligned cells.
		let seeders = nums.length >= 2 ? nums[nums.length - 2] : 0;
		let leechers = nums.length >= 1 ? nums[nums.length - 1] : 0;

		out.push({
			name: decodeEntities(name[1]),
			info_hash: hash[1].toLowerCase(),
			category: cat ? Number(cat[1]) : 0,
			size: size ? sizeToBytes(size[1], size[2]) : 0,
			seeders: seeders,
			leechers: leechers,
			imdb: null,
		});
	}
	return out;
}

async function tpbSearch(title){
	// Category 200 = all Video; we narrow to movie/tv subcats in prefilter().
	let url = `${conf.search.tpbBase}/search/${encodeURIComponent(title)}/1/99/200`;
	let res = await fetch(url, { headers: { 'User-Agent': 'Mozilla/5.0' } });
	if(!res.ok){
		let error = new Error('UpstreamError');
		error.message = `TPB search failed( ${res.status})`;
		error.status = 502;
		throw error;
	}
	return parseTPBHtml(await res.text());
}

// Low-quality source tags we never want to recommend( token match to avoid false hits
// like "GHOSTS" matching "TS").
const JUNK = new Set(['cam', 'hdcam', 'camrip', 'ts', 'hdts', 'telesync', 'tc', 'telecine', 'scr', 'screener', 'dvdscr', 'workprint']);
function isJunk(name){
	return String(name).toLowerCase().split(/[.\s_\-\[\]()]+/).some(tok => JUNK.has(tok));
}

// Deterministic pre-filter before the torrents ever reach the LLM.
function prefilter(torrents, imdbId){
	let out = torrents.filter(t => videoKind(t.category) && Number(t.seeders) > 0 && !isJunk(t.name));

	// If we can positively identify the title by IMDB id, trust it and drop the rest.
	if(imdbId){
		let matches = out.filter(t => t.imdb && t.imdb === imdbId);
		if(matches.length) out = matches;
	}

	return out
		.sort((a, b) => Number(b.seeders) - Number(a.seeders))
		.slice(0, 40);
}

function buildMagnet(infoHash, name){
	let magnet = `magnet:?xt=urn:btih:${infoHash}&dn=${encodeURIComponent(name)}`;
	for(let tracker of conf.search.trackers){
		magnet += `&tr=${encodeURIComponent(tracker)}`;
	}
	return magnet;
}

// --- LLM ranking ----------------------------------------------------------

function buildSystemPrompt(title, year, today){
	return [
		`You are a torrent-selection assistant. Today's date is ${today}.`,
		`Recent releases are legitimate: NEVER reject a torrent for being too new or assume`,
		`a current-year release is fake.`,
		``,
		`The user wants "${title}" (${year || 'unknown year'}). From the candidate list,`,
		`select the best releases.`,
		year ? `IMPORTANT: only pick releases whose name contains the year ${year}; reject other years or remakes of a same-named title.` : ``,
		`Preferences, in order:`,
		`- video codec HEVC/x265 (over x264/AVC)`,
		`- 1080p resolution`,
		`- total size around 1.5 GB (avoid needlessly huge files)`,
		`- English audio and English subtitles present`,
		`- prefer EXTENDED / UNRATED / UNCUT / DIRECTOR'S CUT editions`,
		`Reject: CAM, TS, TC, TELESYNC, HDCAM, SCREENER/SCR, and non-English-only copies.`,
		``,
		`Return STRICT JSON only, shape:`,
		`{"options":[{"role":"recommended|uhd|other","label":"...","info_hash":"...",`,
		`"name":"...","resolution":"...","codec":"...","hasSubs":true,"cut":"...",`,
		`"warning":"...","why":"..."}]}`,
		`Rules: exactly one "recommended" (the 1080p sweet spot). Include one "uhd" ONLY`,
		`if a 4K/2160p release exists, and set its "warning" to note it is a much larger,`,
		`slower download. Add "other" entries for genuinely distinct useful releases`,
		`(e.g. a different cut or a much smaller copy). "why" is one short sentence.`,
		`Copy info_hash verbatim from the chosen candidate.`,
	].join('\n');
}

// Pull the first JSON object out of a model response( tolerates ```json fences / prose).
function extractJSON(content){
	let text = String(content).replace(/```json/gi, '').replace(/```/g, '').trim();
	let start = text.indexOf('{');
	let end = text.lastIndexOf('}');
	if(start === -1 || end === -1) throw new Error('No JSON in model response');
	return JSON.parse(text.slice(start, end + 1));
}

async function rankReleases(candidates, meta){
	let compact = candidates.map(t => ({
		name: t.name,
		sizeHuman: humanSize(t.size),
		seeders: Number(t.seeders),
		info_hash: t.info_hash,
		imdb: t.imdb || null,
	}));

	let body = {
		model: conf.ollama.model,
		stream: false,
		format: 'json',
		messages: [
			{ role: 'system', content: buildSystemPrompt(meta.title, meta.year, meta.today) },
			{ role: 'user', content: JSON.stringify(compact) },
		],
	};

	let data = await fetchJSON(`${conf.ollama.url}/api/chat`, {
		method: 'POST',
		headers: {
			'Content-Type': 'application/json',
			'Authorization': `Bearer ${conf.ollama.apiKey}`,
			'User-Agent': 'tpbproxy/1.0 (search; +https://github.com/wmantly/tpbproxy)',
		},
		body: JSON.stringify(body),
	});

	let parsed = extractJSON(data.message && data.message.content);
	let options = Array.isArray(parsed.options) ? parsed.options : [];
	if(!options.length) throw new Error('Model returned no options');
	return options;
}

// Deterministic fallback if the LLM is unavailable or returns junk: pick the most-seeded
// candidate, preferring x265 + 1080p, so the feature degrades instead of failing.
function fallbackReleases(candidates){
	let score = t => (/x265|hevc/i.test(t.name) ? 2 : 0) + (/1080p/i.test(t.name) ? 1 : 0);
	let best = [...candidates].sort((a, b) =>
		(score(b) - score(a)) || (Number(b.seeders) - Number(a.seeders))
	)[0];
	if(!best) return [];
	return [{
		role: 'recommended',
		label: 'Most seeded',
		info_hash: best.info_hash,
		name: best.name,
		sizeHuman: humanSize(best.size),
		seeders: Number(best.seeders),
		why: 'Automatically chosen (AI ranking unavailable): most seeders.',
	}];
}

// Attach the fields the front end / add flow need to each option and drop any the model
// hallucinated( info_hash must exist in the candidate set).
function decorateOptions(options, candidates){
	let byHash = {};
	for(let t of candidates) byHash[String(t.info_hash).toLowerCase()] = t;

	let out = [];
	for(let opt of options){
		let src = byHash[String(opt.info_hash).toLowerCase()];
		if(!src) continue;
		out.push({
			...opt,
			info_hash: src.info_hash,
			name: opt.name || src.name,
			category: Number(src.category),
			sizeHuman: opt.sizeHuman || humanSize(src.size),
			seeders: Number(src.seeders),
			magnetLink: buildMagnet(src.info_hash, src.name),
		});
	}
	return out;
}

async function findReleases(tmdbId, mediaType){
	let meta = await tmdbDetails(tmdbId, mediaType);
	let torrents = await tpbSearch(meta.title);
	let candidates = prefilter(torrents, meta.imdbId);

	// For movies, require the confirmed release year in the name. TPB search is fuzzy
	// (a "Toy Story 5" query returns Toy Story 4/3/…), and without this the fallback can
	// cross to the wrong movie. If nothing matches, we genuinely have no good copy yet.
	if(mediaType !== 'tv' && meta.year){
		candidates = candidates.filter(t => t.name.includes(meta.year));
	}

	if(!candidates.length){
		return { title: meta.title, year: meta.year, tmdbId, mediaType, options: [] };
	}

	let today = new Date().toISOString().slice(0, 10);
	let options;
	try{
		options = await rankReleases(candidates, { title: meta.title, year: meta.year, today });
	}catch(error){
		console.error('rankReleases failed, using fallback:', error.message);
		options = fallbackReleases(candidates);
	}

	options = decorateOptions(options, candidates);
	if(!options.length) options = decorateOptions(fallbackReleases(candidates), candidates);

	let result = { title: meta.title, year: meta.year, tmdbId, mediaType, options };

	// Quality-aware dedup: flag options you already own at equal-or-better quality.
	if(mediaType !== 'tv'){
		let lib = await library.movieInLibrary(tmdbId);
		if(lib.owned){
			for(let o of options) o.alreadyOwned = library.rankFromLabel(o.resolution) <= lib.rank;
			result.library = { owned: true, quality: lib.quality };
		}
	}

	return result;
}

// --- TV: library-aware season fill -----------------------------------------

function buildSeasonPrompt(title, season, today){
	return [
		`You are a torrent selector. Today's date is ${today}.`,
		`Pick the single best COMPLETE season pack for "${title}" Season ${season}.`,
		`It MUST be the full season (an S${String(season).padStart(2, '0')} / "Season ${season}"`,
		`complete pack), NOT a single episode. Prefer HEVC/x265, 1080p, English audio +`,
		`English subtitles, reasonable size. Reject single episodes, wrong seasons,`,
		`CAM/TS/telesync, and non-English-only.`,
		`Return STRICT JSON {"info_hash":"<from candidates>","why":"one short sentence"};`,
		`if nothing suitable, {"info_hash":null}.`,
	].join('\n');
}

async function bestSeasonPack(candidates, meta){
	let compact = candidates.map(t => ({ name: t.name, sizeHuman: humanSize(t.size), seeders: Number(t.seeders), info_hash: t.info_hash }));
	let parsed = await ollama.chatJSON(buildSeasonPrompt(meta.title, meta.season, meta.today), JSON.stringify(compact));
	return parsed && parsed.info_hash ? parsed : null;
}

// Compare TMDB's season/episode counts against the Emby inventory to classify each season.
async function seasonPlan(tmdbId){
	let [tv, inv] = await Promise.all([ tmdb.tmdbSeasons(tmdbId), library.seriesInventory(tmdbId) ]);
	let seasons = tv.seasons.map(s => {
		let have = Number(inv.have && inv.have[s.season]) || 0;
		let status = have === 0 ? 'missing' : (have >= s.episodeCount ? 'complete' : 'partial');
		return { season: s.season, episodeCount: s.episodeCount, haveCount: have, status };
	});
	return { title: tv.title, year: tv.year, seasons, missing: seasons.filter(s => s.status !== 'complete').map(s => s.season) };
}

// For every missing/incomplete season, pick the best complete-season pack. Returns the
// chosen packs( the route adds them to Transmission) plus any seasons it couldn't fill.
async function fillMissing(tmdbId){
	let plan = await seasonPlan(tmdbId);
	let details = await tmdbDetails(tmdbId, 'tv');
	let title = details.title || plan.title;
	let today = new Date().toISOString().slice(0, 10);

	let picks = [], skipped = [];
	for(let season of plan.missing){
		let cands = prefilter(await tpbSearch(`${title} S${String(season).padStart(2, '0')}`), null);
		if(!cands.length){ skipped.push({ season, reason: 'no results' }); continue; }

		let pick = null;
		try{ pick = await bestSeasonPack(cands, { title, season, today }); }catch(error){ pick = null; }
		let src = pick && cands.find(c => c.info_hash.toLowerCase() === String(pick.info_hash).toLowerCase());
		if(!src){ skipped.push({ season, reason: 'no suitable season pack' }); continue; }

		picks.push({ season, name: src.name, info_hash: src.info_hash, category: Number(src.category), magnetLink: buildMagnet(src.info_hash, src.name), why: pick.why });
	}
	return { title, year: plan.year, picks, skipped };
}

module.exports = {
	tmdbSearch,
	findReleases,
	seasonPlan,
	fillMissing,
	// exported for unit tests
	prefilter,
	buildMagnet,
	humanSize,
	videoKind,
	extractJSON,
	decorateOptions,
	fallbackReleases,
	tpbSearch,
	parseTPBHtml,
	sizeToBytes,
	isJunk,
};