screencap-bot/src/shows/mod.rs

use std::{
    collections::HashMap,
    fs,
    path::{Path, PathBuf},
};

use anyhow::{anyhow, Context};
use log::{debug, error};
use regex::Regex;
use serde::Deserialize;
use serde_with::{serde_as, DisplayFromStr, KeyValueMap};

mod enumeration;

#[derive(Deserialize)]
pub struct Show {
    #[serde(rename = "$key$")]
    pub title: String,
    pub path: PathBuf,
    #[serde(default)]
    pub custom_episodes: Option<CustomEpisodes>,
    #[serde(default)]
    pub tags: Vec<String>,
    #[serde(default)]
    pub parts: HashMap<u32, String>,
    #[serde(default = "default_weight")]
    pub weight: f32,
    #[serde(default, rename = "18+")]
    pub eighteen_plus: Option<bool>,
}

#[serde_as]
#[derive(Deserialize)]
pub struct CustomEpisodes {
    #[serde(default)]
    pub prefix: String,
    #[serde_as(as = "DisplayFromStr")]
    pub regex: Regex,
}

fn default_weight() -> f32 {
    1.0
}

pub type Shows = Vec<Show>;

#[serde_as]
#[derive(Deserialize)]
struct ShowsWrapper(#[serde_as(as = "KeyValueMap<_>")] Shows);

#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub enum EpisodeNumber {
    Standalone,
    SingleSeason(u32),
    MultiSeason(u32, u32),
    Custom(String),
}

type Episodes = HashMap<EpisodeNumber, PathBuf>;

pub fn load<P: AsRef<Path>>(shows_file: P) -> anyhow::Result<Shows> {
    Ok(serde_yaml::from_reader::<_, ShowsWrapper>(
        fs::File::open(shows_file).context("Failed to open shows file")?,
    )
    .context("Failed to parse YAML from shows file")?
    .0)
}

pub fn display_show_episode(show: &Show, episode: &EpisodeNumber) -> String {
    match episode {
        EpisodeNumber::Standalone => show.title.to_string(),
        EpisodeNumber::SingleSeason(n) => format!("{} episode {}", show.title, n),
        EpisodeNumber::MultiSeason(season, ep) => format!(
            "{} {} episode {}",
            show.title,
            show.parts
                .get(season)
                .unwrap_or(&format!("season {}", season)),
            ep
        ),
        EpisodeNumber::Custom(s) => {
            show.custom_episodes
                .as_ref()
                .map(|c| c.prefix.clone())
                .unwrap_or_default()
                + s
        }
    }
}

impl Show {
    pub fn episodes(&self) -> anyhow::Result<Episodes> {
        let path = &self.path;
        let metadata =
            fs::metadata(path).context("Failed to stat the show's path to determine file type")?;
        if metadata.is_file() {
            debug!("{} is a file, standalone", path.display());
            Ok(HashMap::from([(
                EpisodeNumber::Standalone,
                path.to_path_buf(),
            )]))
        } else if metadata.is_dir() {
            debug!("{} is a directory, enumerating episodes", path.display());
            let files: Vec<PathBuf> = fs::read_dir(path)
                .context("Failed to get a directory listing for the show's path")?
                .map(|entry| {
                    let entry =
                        entry.context("Failed to read a directory entry under the show's path")?;
                    if !entry
                        .file_type()
                        .with_context(|| {
                            format!(
                                "Failed to get file type for directory entry {}",
                                entry.path().display()
                            )
                        })?
                        .is_file()
                    {
                        debug!("Skipping {}, not a file", entry.path().display());
                        return Ok(None);
                    }
                    if entry.file_name().into_string().is_err() {
                        error!(
                            "Path {} contains invalid unicode, skipping",
                            entry.path().display()
                        );
                        return Ok(None);
                    }
                    Ok(Some(entry.path()))
                })
                .filter_map(|r| r.transpose())
                .collect::<anyhow::Result<Vec<PathBuf>>>()?;
            if let Some(CustomEpisodes { ref regex, .. }) = self.custom_episodes {
                files
                    .into_iter()
                    .filter_map(|f| -> Option<anyhow::Result<(EpisodeNumber, PathBuf)>> {
                        let episode_name = regex
                            .captures(f.file_name().unwrap().to_str().unwrap())?
                            .name("episode")
                            .map(|m| m.as_str().to_string())
                            .ok_or(anyhow!(
                                "Failed to find capture group `episode` in episode regex"
                            ));
                        Some(episode_name.map(|n| (EpisodeNumber::Custom(n), f)))
                    })
                    .collect::<anyhow::Result<Episodes>>()
            } else {
                enumeration::enumerate_episodes(files)
                    .ok_or(anyhow!("Could not detect any episode numbering scheme"))
            }
        } else {
            Err(anyhow!("The show's path is not a file or a directory"))
        }
    }
}