artcraft/old_code/twitter-vocodes-bot/src/workers.rs

use anyhow::Error;
use anyhow::bail;
use crate::movie_maker::create_movie;
use crate::vocodes_api::vocodes_fetch;
use crate::{common, AnyhowResult, ProgramArgs};
use crossbeam::Receiver;
use egg_mode::media::{MediaHandle, media_types, upload_media, set_metadata, get_status, ProgressInfo};
use egg_mode::tweet::{DraftTweet, Tweet};
use log::{info, error, warn};
use std::path::PathBuf;
use std::thread;
use tokio::runtime::Handle;
use tokio::time::{Duration, delay_for};

pub struct Workload {
  pub tweet: Tweet,
}

pub async fn spawn_workers(config: common::Config, receiver: Receiver<Workload>, program_args: ProgramArgs) {
  for i in 0 .. program_args.worker_count {
    let conf = config.clone();
    let recv = receiver.clone();
    let args = program_args.clone();

    info!("Spawning worker {}", i);

    let handle = Handle::current();

    handle.spawn(async move {
      start_worker(conf, recv, i, args).await;
    });
  }
}

async fn start_worker(config: common::Config, receiver: Receiver<Workload>, worker_id: usize, program_args: ProgramArgs) {
  loop {
    let workload = match receiver.recv() {
      Err(e) => {
        warn!("Error receiving from queue: {} (worker {})", e, worker_id);
        thread::sleep(Duration::from_millis(500));
        continue;
      },
      Ok(workload) => workload,
    };

    info!("Worker {} got tweet: {}", worker_id, &workload.tweet.text);

    let wav_bytes = match vocodes_fetch(&workload.tweet.text, "sonic").await {
      Ok(wav) => wav,
      Err(e) => {
        warn!("Couldn't fetch wav: {} (worker {})", e, worker_id);
        continue;
      }
    };

    let movie_path = match create_movie(wav_bytes, &program_args) {
      Ok(file) => file,
      Err(e) => {
        warn!("Couldn't create movie: {} (worker {})", e, worker_id);
        continue;
      }
    };

    let path = PathBuf::from(movie_path.path().to_str().unwrap());

    let r = respond_with_media(
      &config,
      &program_args,
      &workload.tweet,
      path,
      "deepfake media generated by vo.codes").await;

    match r {
      Ok(_) => info!("Responded to tweet successfully"),
      Err(e) => error!("Error responding to tweet: {}", e),
    }
  }
}

async fn respond_with_media(
  config: &common::Config,
  program_args: &ProgramArgs,
  original_tweet: &Tweet,
  media_path: PathBuf,
  media_alt_text: &str) -> AnyhowResult<()>
{
  let user_id = original_tweet.user
    .as_ref()
    .map(|u| u.id)
    .unwrap_or(0);

  info!("From user id: {}", user_id);

  let mut response_tweet = DraftTweet::new(program_args.tweet_text.to_string())
    .in_reply_to(original_tweet.id)
    .auto_populate_reply_metadata(true)
    .coordinates(33.753746, -84.386330, true); // Atlanta

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

  response_tweet.send(&config.token).await;

  Ok(())
}

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

  let typ = media_types::video_mp4();

  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?;

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

  Ok(handle)
}