reddit-magnet/src/main.rs

use crate::actions::action::Action;
use crate::actions::transmission::TransmissionAction;
use crate::args::Args;
use crate::config::{get_db_path, load_config};
use crate::db::Database;
use crate::magnet::{extract_magnet_links, Magnet};
use chrono::{DateTime, Utc};
use clap::Parser;
use color_eyre::eyre::{eyre, Result, WrapErr};
use log::{debug, info, warn};
use multimap::MultiMap;
use reddit_client::RedditClient;
use regex::Regex;
use std::collections::{HashMap, HashSet};
use std::fs::create_dir_all;

mod actions;
mod args;
mod config;
mod db;
mod magnet;
mod models;
mod reddit_client;
mod report;
mod schema;

#[derive(Debug)]
struct PostInfo {
    title: String,
    submitter: String,
    magnet_links: Vec<Magnet>,
    subreddit: String,
    timestamp: DateTime<Utc>,
    imdb_id: Option<String>,
}

/// Filters posts based on a title filter pattern
fn filter_posts(title: &str, title_filter: Option<Regex>) -> bool {
    match title_filter {
        Some(pattern) => pattern.is_match(title),
        None => true,
    }
}

#[tokio::main]
async fn main() -> Result<()> {
    color_eyre::install()?;

    let args = Args::parse();

    pretty_env_logger::formatted_timed_builder()
        .filter_level(args.verbose.log_level_filter())
        .init();

    // Initialize database
    let db_path = get_db_path(&args)?;

    // Create parent directory if it doesn't exist
    if let Some(parent) = db_path.parent() {
        create_dir_all(parent)
            .wrap_err_with(|| format!("Failed to create directory: {:?}", parent))?;
    }

    let mut db = Database::new(&db_path)
        .wrap_err_with(|| format!("Failed to initialize database at {:?}", db_path))?;

    let conf = load_config(&args)?;

    if conf.sources.is_empty() {
        return Err(eyre!("No sources found in configuration. Please add at least one source to your configuration file.").into());
    }

    let mut unique_usernames = HashSet::new();
    for (_, source_config) in &conf.sources {
        unique_usernames.insert(source_config.username.clone());
    }

    let reddit_client = RedditClient::new();
    let mut user_posts = MultiMap::new();
    for username in unique_usernames {
        let submissions = reddit_client.fetch_user_submissions(&username).await?;
        user_posts.insert_many(username, submissions);
    }

    // Process sources and store magnet links
    let mut total_new_links = 0;
    for (source_name, source_config) in conf.sources {
        info!("Processing source [{}]", source_name);

        let username = source_config.username.clone();
        let title_filter = match source_config.title_filter {
            Some(filter) => Some(
                Regex::new(filter.as_str())
                    .context(format!("Invalid regex pattern: {}", filter))?,
            ),
            None => None,
        };

        if let Some(submissions) = user_posts.get_vec(&username) {
            for post in submissions
                .iter()
                .filter(|s| filter_posts(&*s.title, title_filter.clone()))
            {
                let title = &post.title;
                let body = &post.body;
                let subreddit = &post.subreddit;

                let magnet_links = extract_magnet_links(body);
                if !magnet_links.is_empty() {
                    let post_info = PostInfo {
                        title: title.to_string(),
                        submitter: username.clone(),
                        subreddit: subreddit.to_string(),
                        magnet_links,
                        timestamp: post.created,
                        imdb_id: source_config.imdb_id.clone(),
                    };

                    // Store the post info in the database
                    match db.store_magnets(&post_info) {
                        Ok(count) => {
                            total_new_links += count;
                        }
                        Err(e) => {
                            warn!("Failed to store post info in database: {}", e);
                        }
                    }
                }
            }
        }
    }

    // Initialize actions
    let mut actions: Vec<Box<dyn Action>> = Vec::new();

    // Add Transmission action if enabled
    if let Some(transmission_config) = conf.transmission {
        if transmission_config.enable {
            info!("Initializing Transmission action");
            match TransmissionAction::new(&transmission_config).await {
                Ok(transmission_action) => {
                    actions.push(Box::new(transmission_action));
                }
                Err(e) => {
                    warn!("Failed to initialize Transmission action: {}", e);
                }
            }
        } else {
            debug!("Transmission action is disabled");
        }
    } else {
        debug!("No Transmission configuration found");
    }

    // Process all actions and collect results
    let mut action_results = HashMap::new();

    for action in &mut actions {
        let action_name = action.name().to_string();
        debug!("Processing magnet links with {} action", action_name);

        match action.process_unprocessed_magnets(&mut db).await {
            Ok(processed_magnets) => {
                let count = processed_magnets.success.len();
                debug!(
                    "Successfully processed {} magnet links with {}",
                    count, action_name
                );
                action_results.insert(action_name, processed_magnets);
            }
            Err(e) => {
                warn!("Failed to process magnet links with {}: {}", action_name, e);
            }
        }
    }

    // Generate report
    if let Err(e) = report::generate_report(&action_results, total_new_links) {
        warn!("Failed to generate report: {}", e);
    }

    Ok(())
}