artcraft/src/main.rs

//! A Twitter bot for Vo.codes
//! Copyright 2020 Brandon Thomas

mod common;
mod movie_maker;
mod vocodes_api;
mod workers;

use anyhow::bail;
use crate::workers::{spawn_workers, Workload};
use crossbeam::Receiver;
use easyenv::{init_env_logger, get_env_num, get_env_string};
use egg_mode::error::Error;
use egg_mode::media::{media_types, upload_media, set_metadata, get_status, ProgressInfo, MediaHandle};
use egg_mode::stream::StreamMessage;
use egg_mode::tweet::DraftTweet;
use futures::TryStreamExt;
use log::{warn, info, error};
use std::io::{stdout, Write};
use std::path::PathBuf;
use std::time::Duration;
use std::{env, thread};
use tokio::runtime::Builder;
use tokio::time::delay_for;

pub type AnyhowResult<T> = anyhow::Result<T>;

const ENV_ASSET_IMAGE_PATH : &'static str = "ASSET_IMAGE_PATH";
const ENV_QUEUE_LENGTH : &'static str = "QUEUE_LENGTH";
const ENV_TEMP_DIRECTORY : &'static str = "TEMP_DIRECTORY";
const ENV_THREAD_COUNT : &'static str = "THREAD_COUNT";
const ENV_TWEET_TEXT : &'static str = "TWEET_TEXT";
const ENV_WORKER_THREAD_COUNT : &'static str = "WORKER_THREAD_COUNT";

const VOCODES_BOT_USER_ID : u64 = 1310037441441599488;

#[derive(Clone)]
pub struct Secrets {
  pub api_key: String,
  pub api_secret_key: String,
  pub access_token: String,
  pub access_token_secret: String,
  pub username: String,
  // Not sure where to get the user id, or what it's used for.
  pub user_id: u64,
}

impl Secrets {
  pub fn load_from_env() -> AnyhowResult<Self> {
    Ok(Secrets {
      api_key: env::var("API_KEY")?,
      api_secret_key: env::var("API_SECRET_KEY")?,
      access_token: env::var("ACCESS_TOKEN")?,
      access_token_secret: env::var("ACCESS_TOKEN_SECRET")?,
      username: env::var("USERNAME")?,
      user_id: get_env_num::<u64>("USER_ID", 0)?,
    })
  }
}

#[derive(Clone)]
pub struct ProgramArgs {
  pub tweet_text: String,
  pub png_filename: String,
  pub temp_directory: String,
  pub core_thread_count: usize,
  pub queue_length: usize,
  pub worker_count: usize,
}

//#[tokio::main]
fn main() -> AnyhowResult<()> {
  init_env_logger(None);

  let program_args = ProgramArgs {
    tweet_text: get_env_string(ENV_TWEET_TEXT, "This is a deepfake response generated by vo.codes. Use responsibly."),
    temp_directory: get_env_string(ENV_TEMP_DIRECTORY, "/tmp"),
    png_filename: get_env_string(ENV_ASSET_IMAGE_PATH, "./assets/sonic-full.png"),
    core_thread_count: get_env_num(ENV_THREAD_COUNT, 10)?,
    queue_length: get_env_num(ENV_QUEUE_LENGTH, 10)?,
    worker_count: get_env_num(ENV_WORKER_THREAD_COUNT, 5)?,
  };

  let mut runtime = tokio::runtime::Builder::new()
    .threaded_scheduler()
    .enable_all()
    .core_threads(program_args.core_thread_count)
    .on_thread_start(|| {
      println!("thread started");
    })
    .build()
    .unwrap();

  runtime.block_on(async {
    start_bot(program_args).await
  });

  Ok(())
}

async fn start_bot(program_args: ProgramArgs) -> AnyhowResult<()> {
  info!("Loading configs, secrets, and logging in...");

  let secrets = Secrets::load_from_env()?;
  let config = common::Config::load(secrets.clone()).await;

  info!("Configs loaded. Spawning workers...");

  let (sender, receiver) = crossbeam::bounded(program_args.queue_length);
  spawn_workers(config.clone(), receiver, program_args).await;

  info!("Streaming tweets...");

  let stream = egg_mode::stream::filter()
    .track(&["vocodesbot"])
    .language(&["en"])
    .start(&config.token)
    .try_for_each(|m| {
      if let StreamMessage::Tweet(tweet) = m {
        let user_id = tweet.user
          .as_ref()
          .map(|u| u.id)
          .unwrap_or(0);

        let user_name = tweet.user
          .as_ref()
          .map(|u| u.screen_name.clone())
          .unwrap_or("unnamed".to_string());

        if user_id == VOCODES_BOT_USER_ID {
          info!("Tweet originated from VocodesBot ({}); skipping.", &user_name);
          return futures::future::ok(());
        }

        if tweet.text.contains("deepfake response") {
          info!("Tweet contains \"deepfake response\"; skipping. Username: {}", &user_name);
          return futures::future::ok(());
        }

        if tweet.in_reply_to_user_id.is_some() {
          info!("Tweet has 'in_reply_to_user_id'; skipping. Username: {}", &user_name);
          return futures::future::ok(());
        }

        common::print_tweet(&tweet);

        sender.send(Workload { tweet: tweet.clone() });

      } else {
        info!("Not a message: {:?}", m);
      }

      futures::future::ok(())
    });

  if let Err(e) = stream.await {
    let e : egg_mode::error::Error = e;
    error!("Stream error: {}", &e);
    match e {
      Error::BadUrl => {},
      Error::InvalidResponse(_, _) => {},
      Error::MissingValue(_) => {},
      Error::FutureAlreadyCompleted => {},
      Error::TwitterError(_, _) => {},
      Error::RateLimit(rate_limit) => {
        error!("Rate limited. Will reopen at unix timestamp: {}", rate_limit);
      },
      Error::MediaError(_) => {},
      Error::BadStatus(_) => {},
      Error::NetError(_) => {},
      Error::TlsError(_) => {},
      Error::IOError(_) => {},
      Error::DeserializeError(_) => {},
      Error::TimestampParseError(_) => {},
      Error::TimerShutdownError(_) => {},
      Error::HeaderParseError(_) => {},
      Error::HeaderConvertError(_) => {},
    }
    error!("Disconnected")
  }

  Ok(())
}

async fn respond_with_media(
  config: &common::Config,
  tweet_text: &str,
  media_path: PathBuf,
  media_alt_text: &str) -> Result<(), Box<dyn std::error::Error>>
{
  let mut tweet = DraftTweet::new(tweet_text.to_string());

  let handle = upload_media_to_twitter(config, media_path, media_alt_text).await?;
  tweet.add_media(handle.id.clone());

  tweet.send(&config.token).await?;

  Ok(())
}

async fn upload_media_to_twitter(
  config: &common::Config,
  media_path: PathBuf,
  media_alt_text: &str) -> AnyhowResult<MediaHandle>
{
  println!("Uploading media from '{}'", media_path.display());

  let typ = match media_path.extension().and_then(|os| os.to_str()).unwrap_or("") {
    "jpg" | "jpeg" => media_types::image_jpg(),
    "gif" => media_types::image_gif(),
    "png" => media_types::image_png(),
    "webp" => media_types::image_webp(),
    "mp4" => media_types::video_mp4(),
    _ => {
      eprintln!("Format not recognized, must be one of [jpg, jpeg, gif, png, webp, mp4]");
      std::process::exit(1);
    }
  };

  let bytes = std::fs::read(media_path)?;
  let handle = upload_media(&bytes, &typ, &config.token).await?;

  set_metadata(&handle.id, media_alt_text, &config.token).await?;

  println!("Media uploaded");

  // Wait 60 seconds for processing
  println!("Waiting for media to finish processing...");

  for ct in 0..=60u32 {
    match get_status(handle.id.clone(), &config.token).await?.progress {
      None | Some(ProgressInfo::Success) => {
        println!("\nMedia sucessfully processed");
        break;
      }
      Some(ProgressInfo::Pending(_)) | Some(ProgressInfo::InProgress(_)) => {
        print!(".");
        stdout().flush()?;
        delay_for(Duration::from_secs(1)).await;
      }
      Some(ProgressInfo::Failed(err)) => Err(err)?,
    }
    if ct == 60 {
      bail!("Error: timeout");
    }
  }

  Ok(handle)
}