kensho/src/main.rs

use anyhow::Result;
use chrono::{DateTime, Local, Utc};
use clap::{arg, command, Parser};
use log::{debug, trace};
use megalodon::{
    entities::{Account, Status},
    generator,
    megalodon::GetLocalTimelineInputOptions,
    response::Response,
    Megalodon,
};
use tokio_stream::{iter, StreamExt};

use std::{env, fs::File, io::prelude::*};

use self::{
    format::format_status,
    page::{bounds_from, Page},
    range::try_create_range,
    range::Range,
};

mod format;
mod page;
mod range;

#[derive(Debug, Parser)]
#[command()]
struct Config {
    #[arg(short, long, env = "MASTODON_URL", required = true)]
    url: String,
    #[arg(short, long, env = "MASTODON_ACCESS_TOKEN", required = true)]
    access_token: String,
    #[arg(short, long)]
    output_dir: Option<String>,
    #[arg(required = true)]
    date: String,
    #[arg(short, long, action = clap::ArgAction::Count)]
    verbose: u8,
}

#[tokio::main]
async fn main() -> Result<()> {
    let Config {
        date,
        verbose,
        url,
        access_token,
        output_dir,
    } = Config::parse();

    let level = match verbose {
        0 => "off",
        1 => "debug",
        _ => "trace",
    };
    env::set_var("RUST_LOG", format!("{}={}", module_path!(), level));
    env_logger::init();

    let day = try_create_range(date.clone())?;

    let client = create_client(url, access_token)?;
    let Response { json: account, .. } = client.verify_account_credentials().await?;

    debug!("Fetching posts for date, {}.", day.end.format("%Y-%m-%d"));

    // the server only provides a page of results at a time, keep the oldest status from any page
    // to request the next older page of statuses
    let mut last_id_on_page: Option<String> = None;
    // store the formatted posts in server order, reversed chronologically, to reverse at the end
    // for regular chronological ordering
    let mut reversed = Vec::new();
    loop {
        let statuses = fetch_page(&client, &last_id_on_page).await?;
        if statuses.is_empty() {
            debug!("No more posts in range.");
            break;
        }
        let page = bounds_from(&statuses);

        trace!("Page bounds {:?}", page);

        // this age comparison only applies after the first page is fetched; the rest of the loop
        // body handles if the requested date is newer than any statuses on the first page
        if last_id_on_page.is_some() && page_start_older_than(&page, &day) {
            break;
        }

        // fetching returns 20 at a time, in reverse chronological order so may require skipping
        // pages after the requested date
        if let Some(oldest_id) = page_newer_than(&page, &day) {
            last_id_on_page.replace(oldest_id);
            continue;
        }

        // mapping the vector runs into thorny ownership issues and only produces futures, not
        // resolved values; a for in loop works with await but also runs into thorny ownership
        // issues; a stream resolves both because the stream takes ownership of the statuses and
        // can be iterated in a simple way that allows the use of await in the body
        let mut stream = iter(filter_statuses(&account, &day, &statuses));
        while let Some(status) = stream.next().await {
            reversed.push(format_status(&client, &account, status).await?);
        }

        if page_end_older_than(&page, &day) {
            debug!("No more posts in range.");
            break;
        }

        if let Some(id) = page.oldest_id {
            last_id_on_page.replace(id.clone());
        }
    }
    reversed.reverse();
    if let Some(output_dir) = output_dir {
        let output = format!("{}{}.md", output_dir, date);
        let mut f = File::options().append(true).open(&output)?;
        f.write_all(&reversed.join("\n\n").as_bytes())?;
        println!("Appended matching posts to {}.", output);
    } else {
        println!("{}", reversed.join("\n\n"));
    }
    Ok(())
}

// Only ones authored by the user, on the date requested, that aren't a reply to any other status
fn filter_statuses<'a>(account: &Account, day: &Range, json: &'a Vec<Status>) -> Vec<&'a Status> {
    json.iter()
        .filter(|status| {
            status.account.id == account.id
                && status.in_reply_to_id.is_none()
                && day.start <= status.created_at
                && status.created_at <= day.end
        })
        .collect::<Vec<&Status>>()
}

fn create_client(url: String, token: String) -> Result<Box<dyn Megalodon + Send + Sync>> {
    Ok(generator(megalodon::SNS::Mastodon, url, Some(token), None))
}

async fn fetch_page(
    client: &Box<dyn Megalodon + Send + Sync>,
    last_id_on_page: &Option<String>,
) -> Result<Vec<Status>> {
    let Response { json, .. } = if let Some(max_id) = last_id_on_page.as_ref() {
        debug!("Fetching next page");
        client
            .get_local_timeline(Some(&GetLocalTimelineInputOptions {
                max_id: Some(max_id.clone()),
                ..GetLocalTimelineInputOptions::default()
            }))
            .await?
    } else {
        debug!("Fetching first page");
        client.get_local_timeline(None).await?
    };
    Ok(json)
}

fn page_newer_than(page: &Page, range: &Range) -> Option<String> {
    page.oldest
        .filter(|oldest| *oldest > &range.end)
        .and_then(|_| page.oldest_id.clone())
}

fn page_end_older_than(page: &Page, range: &Range) -> bool {
    status_older_than(&page.oldest, &range.start)
}

fn page_start_older_than(page: &Page, range: &Range) -> bool {
    status_older_than(&page.newest, &range.start)
}

fn status_older_than(status: &Option<&DateTime<Utc>>, dt: &DateTime<Local>) -> bool {
    status.map(|status| status < dt).unwrap_or_default()
}