diff --git a/workflow.py b/workflow.py index 87c35f0..2c2bc0e 100644 --- a/workflow.py +++ b/workflow.py @@ -83,6 +83,162 @@ def wait_for_download(qb, t_hash: str): print("Download complete!") +def check_local_filesystem(dl_dir: str, imdb_id: str): + """ + Checks the local download directory for any existing files or folders + that match the given IMDb ID. Returns the path if found. + """ + dl_path = Path(dl_dir) + if dl_path.exists(): + for item in dl_path.iterdir(): + if f"[{imdb_id}]" in item.name: + return item + return None + +def check_qbittorrent(qb, imdb_id: str): + """ + Checks qBittorrent for any existing torrent that has the IMDb ID in its name. + Returns the torrent hash if found, or None. + """ + existing_torrents = qb.torrents_info() + for t in existing_torrents: + if f"[{imdb_id}]" in t.name: + return t.hash + return None + +def search_and_download_mteam(qb, imdb_id: str, new_name: str, dl_dir: str) -> list: + """ + Searches M-Team for the IMDb ID, uses an LLM to select the best torrents, + downloads them, adds them to qBittorrent, and waits for them to complete. + Returns a list of tuples containing (torrent_hash, torrent_id). + """ + print(f"\n=== [1] Searching Torrents for {imdb_id} ===") + imdb_url = f"https://www.imdb.com/title/{imdb_id}/" + + # Extract the torrent list + torrents = search_mteam_torrents(imdb_url)["data"]["data"] + + if not torrents: + raise ValueError(f"No torrents found on M-Team for IMDb ID: {imdb_id}") + + # Format the torrents text + formatted_torrents = [] + for t in torrents: + if isinstance(t, dict): + formatted_torrents.append(format_mteam_torrent(t)) + torrents_text = "\n\n".join(formatted_torrents) + + print(f"\n=== [2] Selecting best torrents using LLM ===") + selected_ids_str = select_best_torrents(torrents_text) + selected_ids = [tid.strip() for tid in selected_ids_str.split() if tid.strip()] + print(f"Selected torrent IDs: {selected_ids}") + + if not selected_ids: + raise ValueError(f"LLM did not select any torrents for IMDb ID: {imdb_id}") + + hashes_to_process = [] + for tid in selected_ids: + print(f"\n=== [3] Downloading .torrent for ID: {tid} ===") + torrent_bytes = generate_mteam_download_token(tid) + + print(f"\n=== [4] Adding torrent to qBittorrent ===") + download_torrent(qb, torrent_bytes, dl_dir) + + # Parse local hash directly instead of hoping qB orders correctly + t_hash = get_torrent_hash(torrent_bytes) + if not t_hash: + print(f"Could not compute hash for {tid}, skipping!") + continue + + print(f"\n=== [5] Waiting for download to finish ===") + # Wait slightly for qB to process the adding request + time.sleep(3) + + print(f"Tracking torrent Hash: {t_hash}") + + rename_torrent_and_folder(qb, t_hash, new_name) + wait_for_download(qb, t_hash) + + hashes_to_process.append((t_hash, tid)) + + return hashes_to_process + +def process_qb_torrent(qb, t_hash: str, tid: str, new_name: str, dl_dir: str, title_dir: str, imdb_id: str, jellyfin_base_dir: str): + """ + Uses an LLM to generate a rename mapping for the files inside a completed qBittorrent download. + Applies tags based on TV/Movie type and creates symbolic links to the Jellyfin library. + """ + print(f"\n=== [6] Generating rename mapping using LLM ===") + file_tree = get_torrent_file_tree(qb, t_hash) + + src_dir_for_mapping = prepare_file_tree_paths(file_tree, new_name, dl_dir) + + file_tree_str = format_file_tree(file_tree) + + prompt_text = f"Base directory: `{title_dir}`\n\n{file_tree_str}" + print(f"Sending paths to LLM...") + mapping = generate_rename_mapping(prompt_text) + print("Generated Mapping:") + + is_tv = False + for src, dst in mapping.items(): + print(f" {src} -->> {dst}") + if "Season " in dst or "Series " in dst: + is_tv = True + + tag = "Jellyfin TV" if is_tv else "Jellyfin Movie" + opposite_tag = "Jellyfin Movie" if is_tv else "Jellyfin TV" + + jellyfin_dir = f"{jellyfin_base_dir}/TV" if is_tv else f"{jellyfin_base_dir}/Movie" + jellyfin_base = Path(jellyfin_dir) / f"{title_dir} [{imdb_id}]" + + print(f"\n=== [6.5] Adding '{tag}' tag to torrent ===") + qb.torrents_add_tags(tags=tag, torrent_hashes=t_hash) + remove_tag_if_exists(qb, t_hash, opposite_tag) + + print(f"\n=== [7] Creating symbolic links ===") + apply_rename_mapping(mapping, base_src_dir=src_dir_for_mapping, base_dst_dir=jellyfin_base) + print(f"Finished processing torrent: {tid}") + +def process_local_file(fs_path: Path, dl_dir: str, title_dir: str, imdb_id: str, jellyfin_base_dir: str): + """ + Generates an LLM rename mapping for purely local files/folders (skipping qBittorrent) + and creates symbolic links to the Jellyfin library. + """ + print(f"\n=== [6] Generating rename mapping using LLM for local path ===") + + # Mock file tree logic for local files + file_tree = [] + if fs_path.is_file(): + file_tree.append({"name": fs_path.name}) + else: + for p in fs_path.rglob('*'): + if p.is_file(): + # relative to fs_path's parent so it starts with fs_path.name + rel_path = p.relative_to(fs_path.parent) + file_tree.append({"name": str(rel_path.as_posix())}) + + src_dir_for_mapping = prepare_file_tree_paths(file_tree, fs_path.name, dl_dir) + file_tree_str = format_file_tree(file_tree) + + prompt_text = f"Base directory: `{title_dir}`\n\n{file_tree_str}" + print(f"Sending paths to LLM...") + mapping = generate_rename_mapping(prompt_text) + print("Generated Mapping:") + + is_tv = False + for src, dst in mapping.items(): + print(f" {src} -->> {dst}") + if "Season " in dst or "Series " in dst: + is_tv = True + + jellyfin_dir = f"{jellyfin_base_dir}/TV" if is_tv else f"{jellyfin_base_dir}/Movie" + jellyfin_base = Path(jellyfin_dir) / f"{title_dir} [{imdb_id}]" + + print(f"\n=== [7] Creating symbolic links ===") + apply_rename_mapping(mapping, base_src_dir=src_dir_for_mapping, base_dst_dir=jellyfin_base) + print(f"Finished processing local file: {fs_path.name}") + def process_imdb_workflow(imdb_id: str, dl_dir: str = DEFAULT_DL_DIR, jellyfin_base_dir: str = DEFAULT_JELLYFIN_DIR): """ Workflow to automatically find, download, and map torrents for an IMDb ID into a Jellyfin library. @@ -97,16 +253,18 @@ def process_imdb_workflow(imdb_id: str, dl_dir: str = DEFAULT_DL_DIR, jellyfin_b title_dir = f"{title} ({year})" print(f"Found Title: {title_dir}") - print(f"\n=== [0.5] Checking if torrent already exists in qBittorrent ===") - qb = get_qb_client() + print(f"\n=== [0.2] Checking if already exists in file system ===") + fs_match_dir = check_local_filesystem(dl_dir, imdb_id) new_name = f"{year} {title} [{imdb_id}]".strip() - existing_torrents = qb.torrents_info() - existing_t_hash = None - for t in existing_torrents: - if f"[{imdb_id}]" in t.name: - existing_t_hash = t.hash - break + if fs_match_dir: + print(f"Found existing file/directory in file system: {fs_match_dir.name}, skipping qBit check, search, and download.") + process_local_file(fs_match_dir, dl_dir, title_dir, imdb_id, jellyfin_base_dir) + return + + print(f"\n=== [0.5] Checking if torrent already exists in qBittorrent ===") + qb = get_qb_client() + existing_t_hash = check_qbittorrent(qb, imdb_id) hashes_to_process = [] @@ -119,94 +277,11 @@ def process_imdb_workflow(imdb_id: str, dl_dir: str = DEFAULT_DL_DIR, jellyfin_b hashes_to_process.append((existing_t_hash, "existing")) else: - print(f"\n=== [1] Searching Torrents for {imdb_id} ===") - imdb_url = f"https://www.imdb.com/title/{imdb_id}/" - search_res = search_mteam_torrents(imdb_url) - - # Extract the torrent list - if "data" in search_res and isinstance(search_res["data"], dict) and "data" in search_res["data"]: - torrents = search_res["data"]["data"] - elif "data" in search_res and isinstance(search_res["data"], list): - torrents = search_res["data"] - elif isinstance(search_res, list): - torrents = search_res - else: - torrents = [] - - if not torrents: - raise ValueError(f"No torrents found on M-Team for IMDb ID: {imdb_id}") - - # Format the torrents text - formatted_torrents = [] - for t in torrents: - if isinstance(t, dict): - formatted_torrents.append(format_mteam_torrent(t)) - torrents_text = "\n\n".join(formatted_torrents) - - print(f"\n=== [2] Selecting best torrents using LLM ===") - selected_ids_str = select_best_torrents(torrents_text) - selected_ids = [tid.strip() for tid in selected_ids_str.split() if tid.strip()] - print(f"Selected torrent IDs: {selected_ids}") - - if not selected_ids: - raise ValueError(f"LLM did not select any torrents for IMDb ID: {imdb_id}") - - for tid in selected_ids: - print(f"\n=== [3] Downloading .torrent for ID: {tid} ===") - torrent_bytes = generate_mteam_download_token(tid) - - print(f"\n=== [4] Adding torrent to qBittorrent ===") - download_torrent(qb, torrent_bytes, dl_dir) - - # Parse local hash directly instead of hoping qB orders correctly - t_hash = get_torrent_hash(torrent_bytes) - if not t_hash: - print(f"Could not compute hash for {tid}, skipping!") - continue - - print(f"\n=== [5] Waiting for download to finish ===") - # Wait slightly for qB to process the adding request - time.sleep(3) - - print(f"Tracking torrent Hash: {t_hash}") - - rename_torrent_and_folder(qb, t_hash, new_name) - wait_for_download(qb, t_hash) - - hashes_to_process.append((t_hash, tid)) + hashes_to_process = search_and_download_mteam(qb, imdb_id, new_name, dl_dir) + # Process qB torrents for t_hash, tid in hashes_to_process: - print(f"\n=== [6] Generating rename mapping using LLM ===") - file_tree = get_torrent_file_tree(qb, t_hash) - - src_dir_for_mapping = prepare_file_tree_paths(file_tree, new_name, dl_dir) - - file_tree_str = format_file_tree(file_tree) - - prompt_text = f"Base directory: `{title_dir}`\n\n{file_tree_str}" - print(f"Sending paths to LLM...") - mapping = generate_rename_mapping(prompt_text) - print("Generated Mapping:") - - is_tv = False - for src, dst in mapping.items(): - print(f" {src} -->> {dst}") - if "Season " in dst or "Series " in dst: - is_tv = True - - tag = "Jellyfin TV" if is_tv else "Jellyfin Movie" - opposite_tag = "Jellyfin Movie" if is_tv else "Jellyfin TV" - - jellyfin_dir = f"{jellyfin_base_dir}/TV" if is_tv else f"{jellyfin_base_dir}/Movie" - jellyfin_base = Path(jellyfin_dir) / f"{title_dir} [{imdb_id}]" - - print(f"\n=== [6.5] Adding '{tag}' tag to torrent ===") - qb.torrents_add_tags(tags=tag, torrent_hashes=t_hash) - remove_tag_if_exists(qb, t_hash, opposite_tag) - - print(f"\n=== [7] Creating symbolic links ===") - apply_rename_mapping(mapping, base_src_dir=src_dir_for_mapping, base_dst_dir=jellyfin_base) - print(f"Finished processing torrent: {tid}") + process_qb_torrent(qb, t_hash, tid, new_name, dl_dir, title_dir, imdb_id, jellyfin_base_dir) if __name__ == "__main__": import argparse