/* * Copyright (C) 2024 Jonni Liljamo * * This file is licensed under GPL-3.0-or-later, see NOTICE and LICENSE for * more information. */ use std::collections::HashMap; use fancy_regex::Regex; pub struct MkvInfo { pub title: String, pub year: usize, pub file_name: String, } impl MkvInfo { pub fn new(mkv: &matroska::Matroska, file_name: &str) -> MkvInfo { let metadata_title = mkv.info.title.clone(); let (guessed_title, guessed_year) = guess_title_and_release_year(file_name, &metadata_title.clone().unwrap_or_default()); let title = inquire::Text::new("Enter the title:") .with_initial_value(&guessed_title) .prompt() .unwrap(); let year = inquire::CustomType::::new("Please enter the release year:") .with_starting_input(&guessed_year) .prompt() .unwrap(); // Find the default video and audio tracks, which are needed for parts of the filename. let mut default_video_track = mkv.video_tracks().find(|track| track.default); if default_video_track.is_none() { default_video_track = mkv.video_tracks().nth(0); } let default_video_track_settings = match &default_video_track.unwrap().settings { matroska::Settings::Video(video_settings) => video_settings, _ => panic!("default video track did not have video settings"), }; let mut default_audio_track = mkv.audio_tracks().find(|track| track.default); if default_audio_track.is_none() { default_audio_track = mkv.audio_tracks().nth(0); } let default_audio_track_settings = match &default_audio_track.unwrap().settings { matroska::Settings::Audio(audio_settings) => audio_settings, _ => panic!("default video track did not have audio settings"), }; // "{T}.{Y}.{S}.{R}.{SRC}.{VC}.{AC}.mkv" let mut recommended_file_name = String::new(); // T: Title, with spaces replaced with periods and take out various special characters. recommended_file_name += &title .replace(" ", ".") .replace("'", "") .replace(":", "") .replace("-", ""); // Y: Release year. recommended_file_name += &format!(".{}", year); // S: Special, e.g. directors cut or extended. if let Some(special) = get_special(file_name) { recommended_file_name += &format!(".{}", special); } // R: Resolution, e.g. 360p, 720p, 1080p, 2160p. recommended_file_name += &format!(".{}p", get_resolution(default_video_track_settings)); // SRC: Source, one of: BluRay, DVD, WebDL, Unknown. recommended_file_name += &format!(".{}", get_source(file_name)); // VC: Video codec, the default or first one. let default_video_track_codec_id = default_video_track.unwrap().codec_id.clone(); match default_video_track_codec_id.as_str() { "V_MPEG4/ISO/AVC" => recommended_file_name += ".x264", "V_MPEGH/ISO/HEVC" => recommended_file_name += ".x265", _ => panic!("unhandled video codec: '{}'", default_video_track_codec_id), } // AC: Audio codec, the default or first one. let default_audio_track_codec_id = default_audio_track.unwrap().codec_id.clone(); match default_audio_track_codec_id.as_str() { "A_AAC" => recommended_file_name += ".AAC", "A_AC3" => recommended_file_name += ".DD", "A_EAC3" => recommended_file_name += ".DD+", "A_DTS" => recommended_file_name += ".DTS", _ => panic!("unhandled audio codec: '{}'", default_audio_track_codec_id), } match default_audio_track_settings.channels { 0 => {} // Unknown 2 => {} // Stereo, no need to mention 6 => recommended_file_name += ".5.1", 8 => recommended_file_name += ".7.1", _ => panic!( "unhandled amount of audio channels: '{}'", default_audio_track_settings.channels ), } // And last but not least, the file extension. recommended_file_name += ".mkv"; // Last minute cleanup, e.g. double periods. recommended_file_name = recommended_file_name.replace("..", "."); let file_name = inquire::Text::new("File name:") .with_initial_value(&recommended_file_name) .prompt() .unwrap(); MkvInfo { title, year, file_name, } } } /// Try to guess the title and release year from the file name and metadata title (if any). fn guess_title_and_release_year(file_name: &str, metadata_title: &str) -> (String, String) { // !p at the end, because we don't want to match the resolution accidentally. let year_regex = Regex::new(r"\d{4}+(?!p)").unwrap(); // First, let's find them from the file name. let mut fn_year_guess = String::new(); let mut fn_title_guess = String::new(); if let Some(year_match) = year_regex.find(file_name).unwrap() { // A possible year was found, great! fn_year_guess = year_match.as_str().to_string(); // Now let's take all the text before it, and use that as the title guess. fn_title_guess = file_name[..year_match.start()] .replace(".", " ") .trim_end_matches("(") // Year might've been in parentheses .trim() .to_string(); } // Second, let's find guesses from the metadata title, if one is set. let mut md_year_guess = String::new(); let mut md_title_guess = String::new(); if !metadata_title.is_empty() { // This is just slimmed down from the file name guessing. if let Some(year_match) = year_regex.find(metadata_title).unwrap() { md_year_guess = year_match.as_str().to_string(); md_title_guess = metadata_title[..year_match.start()] .trim_end_matches("(") // Year might've been in parentheses .trim() .to_string(); } } // Then let's decide which of each of these is correct. let mut year_guess = String::new(); if !fn_year_guess.is_empty() && !md_year_guess.is_empty() { // Both were found, prefer metadata. year_guess = md_year_guess; } else if fn_year_guess.is_empty() && !md_year_guess.is_empty() { // Only metadata was found. year_guess = md_year_guess; } else if !fn_year_guess.is_empty() && md_year_guess.is_empty() { // Only file name was found. year_guess = fn_year_guess; } let mut title_guess = String::new(); if !fn_title_guess.is_empty() && !md_title_guess.is_empty() { if fn_title_guess == md_title_guess { title_guess = md_title_guess; } else { let title_options = vec![md_title_guess, fn_title_guess]; title_guess = inquire::Select::new("Which of these titles is better?", title_options) .prompt() .unwrap(); } } else if fn_title_guess.is_empty() && !md_title_guess.is_empty() { // Only metadata was found. title_guess = md_title_guess; } else if !fn_title_guess.is_empty() && md_title_guess.is_empty() { // Only file name was found. title_guess = fn_title_guess; } (title_guess, year_guess) } /// Get possible special kind for the file. /// /// Tries to guess it and place the cursor on the right option. fn get_special(file_name: &str) -> Option { // Map with special feature kind mapped to ways it could be written in the file name. let mut special_option_pairs = HashMap::new(); special_option_pairs.insert("DIRECTORS.CUT", vec!["directors.cut", "directors cut"]); special_option_pairs.insert("EXTENDED", vec!["extended"]); special_option_pairs.insert("OPEN.MATTE", vec!["open.matte", "open matte"]); special_option_pairs.insert("REMASTERED", vec!["remastered"]); special_option_pairs.insert("RESTORED", vec!["restored"]); special_option_pairs.insert("THEATER", vec!["theater"]); special_option_pairs.insert("UNRATED", vec!["unrated"]); let mut special_matched = None; 'outer: for (special_kind, options) in &special_option_pairs { for o in options { if file_name.to_lowercase().contains(o) { special_matched = Some(special_kind); break 'outer; } } } if inquire::Confirm::new("Is this a special feature?") .with_default(special_matched.is_some()) .prompt() .unwrap() { let mut special_options = special_option_pairs .keys() .map(|key| key.to_string()) .collect::>(); special_options.sort(); let mut match_pos = 0; if let Some(special_kind) = special_matched { match_pos = special_options .iter() .position(|s| s == special_kind) .unwrap(); } let special_ans = inquire::Select::new("What kind of special feature?", special_options) .with_starting_cursor(match_pos) .prompt(); match special_ans { Ok(special) => Some(special), Err(_) => panic!("something went wrong while asking for special value"), } } else { None } } /// Get the resolution of the file. fn get_resolution(video_settings: &matroska::Video) -> u64 { // NOTE: Kinda dumb logic to figure out a sane value for the resolution field. // There's a lot of "theater aspect" content which is normal 1080p in width // but has a height of like 1769. I still want that labeled 1080p. match video_settings.pixel_height { 360 => 360, 720 => 720, 1080 => 1080, 2160 => 2160, _ => match video_settings.pixel_width { 480 => 360, 1280 => 720, 1920 => 1080, 3840 => 2160, _ => video_settings.pixel_height, }, } } /// Get the source of the file. /// /// Tries to guess it and place the cursor on the right option. fn get_source(file_name: &str) -> String { // Map with source kind mapped to ways it could be written in the file name. let mut source_option_pairs = HashMap::new(); source_option_pairs.insert("BluRay", vec!["bluray", "blu.ray", "blu-ray"]); source_option_pairs.insert("DVD", vec!["dvd"]); source_option_pairs.insert("WebDL", vec!["webdl", "web.dl", "web-dl"]); source_option_pairs.insert("Unknown", vec![]); let mut source_matched = None; 'outer: for (source_kind, options) in &source_option_pairs { for o in options { if file_name.to_lowercase().contains(o) { source_matched = Some(source_kind); break 'outer; } } } let mut source_options = source_option_pairs .keys() .map(|key| key.to_string()) .collect::>(); source_options.sort(); let mut match_pos = 0; if let Some(source_kind) = source_matched { match_pos = source_options .iter() .position(|s| s == source_kind) .unwrap(); } inquire::Select::new("What is the source of this file?", source_options) .with_starting_cursor(match_pos) .prompt() .unwrap() }