kensho/src/main.rs

use anyhow::{bail, format_err, Result};
use chrono::{DateTime, Local, LocalResult, NaiveDate, TimeZone, Utc};
use clap::{arg, command, Parser};
use html2md::parse_html;
use log::{debug, trace};
use megalodon::{
    entities::Status, generator, megalodon::GetLocalTimelineInputOptions, response::Response,
    Megalodon,
};
use tokio_stream::{iter, StreamExt};

use std::env;

#[derive(Debug)]
struct Range {
    start: DateTime<Local>,
    end: DateTime<Local>,
}

#[derive(Debug)]
struct Page<'a> {
    oldest_id: Option<String>,
    oldest: Option<&'a DateTime<Utc>>,
    newest: Option<&'a DateTime<Utc>>,
}

#[derive(Debug, Parser)]
#[command()]
struct Config {
    #[arg(required = true)]
    date: String,
    #[arg(short, long)]
    verbose: bool,
}

#[tokio::main]
async fn main() -> Result<()> {
    let Config { date, verbose } = Config::parse();
    if verbose {
        env::set_var("RUST_LOG", format!("{}=debug", module_path!()));
    } else {
        env::set_var("RUST_LOG", format!("{}=none", module_path!()));
    }
    env_logger::init();

    let day = try_create_range(date)?;

    debug!("Date {}", day.end.format("%Y-%m-%d"));

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

    let mut last_id_on_page: Option<String> = None;
    debug!("Fetching posts");
    let mut reversed = Vec::new();
    loop {
        let json = fetch_page(&client, &last_id_on_page).await?;
        let page = Page {
            newest: json.first().map(|s| &s.created_at),
            oldest_id: json.last().map(|s| s.id.clone()),
            oldest: json.last().map(|s| &s.created_at),
        };
        trace!("Page bounds {:?}", page);

        if last_id_on_page.is_some() && page_start_older_than(&page, &day) {
            break;
        }

        if let Some(oldest_id) = page_newer_than(&page, &day) {
            last_id_on_page.replace(oldest_id);
            continue;
        }

        let json = json
            .clone()
            .into_iter()
            .filter(|json| {
                json.account.id == account.id
                    && day.start <= json.created_at
                    && json.created_at <= day.end
            })
            .collect::<Vec<Status>>();
        trace!("Filtered to {} post(s)", json.len());

        let mut stream = iter(json);

        while let Some(status) = stream.next().await {
            let ancestor = format!(
                "{}
> {}",
                status.created_at.with_timezone(&Local).format("%H:%M"),
                parse_html(&status.content).trim()
            );
            let Response { json, .. } = client.get_status_context(status.id, None).await?;
            let thread = json
                .descendants
                .into_iter()
                .filter(|s| {
                    s.in_reply_to_account_id == Some(account.id.clone())
                        && s.account.id == account.id
                })
                .map(|status| {
                    format!(
                        ">
> {}
>> {}",
                        status.created_at.with_timezone(&Local).format("%H:%M"),
                        parse_html(&status.content).trim()
                    )
                })
                .collect::<Vec<String>>()
                .join("\n");
            reversed.push(format!(
                "{}{}",
                ancestor,
                if !thread.is_empty() {
                    thread
                } else {
                    String::default()
                }
            ));
        }

        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();
    println!("{}", reversed.join("\n\n"));
    Ok(())
}

fn create_client() -> Result<Box<dyn Megalodon + Send + Sync>> {
    let url = env::var("MASTODON_URL")?;
    let token = env::var("MASTODON_ACCESS_TOKEN")?;
    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 try_create_range<S: AsRef<str>>(date: S) -> Result<Range> {
    Ok(Range {
        start: create_day_bound(&date, 0, 0, 0)?,
        end: create_day_bound(date, 23, 59, 59)?,
    })
}

fn create_day_bound<S: AsRef<str>>(
    day: S,
    hour: u32,
    minute: u32,
    second: u32,
) -> Result<DateTime<Local>> {
    let ts: Vec<&str> = day.as_ref().split("-").collect();
    if ts.len() != 3 {
        bail!("Invalid date format! {}", day.as_ref())
    }
    let (year, month, day) = if let [year, month, day, ..] = &ts[..] {
        (year, month, day)
    } else {
        bail!("Invalid date format! {}", day.as_ref())
    };
    let b = Local.from_local_datetime(
        &NaiveDate::from_ymd_opt(year.parse()?, month.parse()?, day.parse()?)
            .ok_or_else(|| format_err!("Invalid date!"))?
            .and_hms_opt(hour, minute, second)
            .ok_or_else(|| format_err!("Invalid time!"))?,
    );
    if let LocalResult::Single(b) = b {
        Ok(b)
    } else {
        bail!("Cannot construct day boundary!")
    }
}

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()
}