slack_message_count.rs

use std::fs::File;
use std::{
    collections::BTreeMap,
    path::PathBuf,
    time::{Duration, UNIX_EPOCH},
};

use anyhow::Result;
use chrono::prelude::*;
use clap::{arg, command, Parser};
use csv::Writer;
use serde::Deserialize;
use tokio::fs::File as AsyncFile;
use tokio::io::{AsyncReadExt, BufWriter};
use walkdir::{DirEntry, WalkDir};

#[derive(Parser, Debug, Clone)]
#[command(author, version, about, long_about = None)]
pub struct Args {
    /// Directory containing Slack export JSON files
    #[arg(long, short = 'd')]
    pub dir_path: PathBuf,
}

fn is_json_file(entry: &DirEntry) -> bool {
    entry
        .path()
        .extension()
        .map(|ext| ext == "json")
        .unwrap_or(false)
}

#[derive(Debug, Deserialize)]
struct SlackMessage {
    #[serde(rename = "type")]
    message_type: String,
    ts: String, // Slack timestamp
}

pub struct SlackExport {
    root_path: PathBuf,
}

impl SlackExport {
    pub fn new(root_path: &PathBuf) -> Self {
        Self {
            root_path: root_path.clone(),
        }
    }

    async fn process_directory(&self) -> Result<BTreeMap<String, usize>> {
        let mut date_counts: BTreeMap<String, usize> = BTreeMap::new();

        for entry in WalkDir::new(&self.root_path)
            .into_iter()
            .filter_map(Result::ok)
            .filter(is_json_file)
        {
            let file_path = entry.path();
            let mut file = AsyncFile::open(file_path).await?;

            let mut contents = String::new();
            file.read_to_string(&mut contents).await?;

            let messages: Vec<SlackMessage> = serde_json::from_str(&contents)?;

            for message in messages {
                if message.message_type == "message" {
                    let timestamp = message.ts.parse::<f64>().unwrap_or(0.0);
                    let datetime = UNIX_EPOCH + Duration::from_secs_f64(timestamp);
                    let formatted_date = Utc
                        .timestamp(datetime.duration_since(UNIX_EPOCH)?.as_secs() as i64, 0)
                        .format("%Y-%m-%d")
                        .to_string();
                    *date_counts.entry(formatted_date).or_insert(0) += 1;
                }
            }
        }

        Ok(date_counts)
    }
}

fn write_to_csv(
    data: &BTreeMap<String, usize>,
    file_path: &str,
) -> Result<(), Box<dyn std::error::Error>> {
    let file = File::create(file_path)?;
    let mut wtr = Writer::from_writer(file);

    // Write headers
    wtr.write_record(&["Date", "Count"])?;

    // Write data
    for (date, count) in data {
        wtr.write_record(&[date, &count.to_string()])?;
    }

    Ok(())
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    let args: Args = Args::parse();

    let slack_export = SlackExport::new(&args.dir_path);
    let date_counts = slack_export.process_directory().await?;

    // Write the results to a CSV file named "slack_messages_by_date.csv"
    write_to_csv(&date_counts, "slack_messages_by_date.csv").expect("Unable to write CSV file");

    println!("Data written to slack_messages_by_date.csv");

    Ok(())
}