path: root/src/importer.rs

use chrono::Utc;
use diesel::PgConnection;
use serde::Deserialize;
use serde::Serialize;
use serde_json::to_value;
use serde_json::Value;
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::RwLock;
use std::thread;
use std::time::Duration;
use std::time::Instant;
use threadpool::ThreadPool;

use crate::db;
use crate::diesel::Connection;
use crate::error::Error;
use crate::models;
use crate::schema;
use crate::strava;

pub const WORKERS: usize = 10;
pub const EMPTY_PARAMS: &[(&str, &str)] = &[];

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum Command {
    ImportStravaUser { username: String },
    ProcessAllRawData,
    ProcessRawData(models::RawDataKey),
}

pub struct ImporterSharedData<StravaApi: strava::StravaApi + 'static> {
    strava: RwLock<StravaApi>,
    pool: Mutex<ThreadPool>,
    conn: Mutex<PgConnection>,
    running: Mutex<bool>,
}

pub struct Importer<StravaApi: strava::StravaApi + 'static> {
    shared: Arc<ImporterSharedData<StravaApi>>,
}

fn run_periodically<S: strava::StravaApi>(shared: Arc<ImporterSharedData<S>>, period: Duration) {
    let sleep_time = Duration::from_millis(1000);
    let mut now = Instant::now();
    loop {
        while now.elapsed() < period {
            if !*shared.running.lock().unwrap() {
                return;
            }
            thread::sleep(sleep_time);
        }
        now = Instant::now();

        info!("run_periodically: wakeup");
        handle_tasks(shared.clone())
    }
}

fn handle_one_task<S: strava::StravaApi>(
    shared: Arc<ImporterSharedData<S>>,
    task: models::Task,
) -> Result<(), Error> {
    let command = serde_json::from_value(task.payload.clone())?;

    match command {
        Command::ImportStravaUser { username } => import_strava_user(shared, username.as_str())?,
        Command::ProcessAllRawData => process_all_raw_data(shared, task.clone())?,
        Command::ProcessRawData(ref d) => process_raw_data(shared, d, task.clone())?,
    }

    Ok(())
}

fn handle_tasks<S: strava::StravaApi>(shared: Arc<ImporterSharedData<S>>) {
    let mut done = false;
    while !done {
        let task = (|| {
            let conn = shared.conn.lock().unwrap();
            let now = Utc::now();
            let eta = now + chrono::Duration::seconds(5);

            db::take_task(&conn, models::TaskState::NEW, now, eta)
        })();

        match task {
            Err(Error::NotFound) => {
                info!("No more tasks");
                done = true;
            }
            Err(e) => {
                error!("Failed to get task: {:?}", e);
                done = true;
            }
            Ok(t) => match handle_one_task(shared.clone(), t) {
                Err(e) => error!("Task failed: {:?}", e),
                _ => (),
            },
        }

        thread::sleep(Duration::from_secs(10));
    }
}

impl<StravaApi: strava::StravaApi> Importer<StravaApi> {
    pub fn new(conn: PgConnection, strava: StravaApi) -> Importer<StravaApi> {
        let shared = Arc::new(ImporterSharedData {
            pool: Mutex::new(ThreadPool::with_name("importer".to_string(), WORKERS)),
            conn: Mutex::new(conn),
            strava: RwLock::new(strava),
            running: Mutex::new(false),
        });
        Importer { shared: shared }
    }

    pub fn run(&self) {
        info!("run()");
        let pool = self.shared.pool.lock().unwrap();
        let mut running = self.shared.running.lock().unwrap();
        if !*running {
            *running = true;
            pool.execute({
                let shared = self.shared.clone();
                move || run_periodically(shared, Duration::from_secs(5))
            });
        }
    }

    pub fn join(&self) {
        *self.shared.running.lock().unwrap() = false;
        self.shared.pool.lock().unwrap().join();
    }
}

fn to_run(data: &Value) -> Result<Value, Error> {
    info!("to_run");
    macro_rules! get {
        ($id:ident str $e:expr) => {
            $id[$e]
                .as_str()
                .ok_or_else(|| Error::UnexpectedJson($id.clone()))
        };

        ($id:ident f64 $e:expr) => {
            $id[$e]
                .as_f64()
                .ok_or_else(|| Error::UnexpectedJson($id.clone()))
        };
    };

    Ok(json!({
        "type": "run",
        "distance": get!(data f64 "distance")?,
        "elapsed_time": get!(data f64 "elapsed_time")?,
    }))
}

fn process_strava_activity<S: strava::StravaApi>(
    shared: Arc<ImporterSharedData<S>>,
    data: models::RawData,
    mut task: models::Task,
) -> Result<(), Error> {
    let json_error = || Error::UnexpectedJson(data.payload.clone());
    let strava_type = data.payload["type"].as_str().ok_or_else(json_error)?;

    let entry_type = match strava_type {
        "Run" => Ok("run".to_string()),
        &_ => Err(Error::NotFound),
    }?;

    let conn = &shared.conn.lock().unwrap();
    conn.transaction::<(), Error, _>(|| {
        let entry = models::NewEntry {
            username: data.username.as_str(),
            entry_type: entry_type.as_str(),
            timestamp: None,
            payload: json!({}),
        };

        let id = insert!(entries <= conn, &entry)?;

        let entry_data = models::EntryData {
            username: entry.username.to_string(),
            entry_type: entry.entry_type.to_string(),
            entry_id: id,
            data_type: data.data_type,
            data_id: data.id,
        };
        insert!(entry_data <= conn, &entry_data)?;

        task.state = models::TaskState::SUCCESSFUL;
        db::update_task(conn, task)?;
        Ok(())
    })?;

    Ok(())
}

fn process_raw_data<S: strava::StravaApi>(
    shared: Arc<ImporterSharedData<S>>,
    key: &models::RawDataKey,
    mut task: models::Task,
) -> Result<(), Error> {
    let data = db::get_raw_data(&shared.conn.lock().unwrap(), key)?;
    println!("Process raw data: {:#?}", data);

    match data.data_type {
        models::DataType::StravaActivity => process_strava_activity(shared.clone(), data, task)?,
    };

    info!("Process finished.");
    Ok(())
}

fn process_all_raw_data<S: strava::StravaApi>(
    shared: Arc<ImporterSharedData<S>>,
    mut task: models::Task,
) -> Result<(), Error> {
    let now = Utc::now();
    let conn = &shared.conn.lock().unwrap();

    let mut tasks = 0;
    conn.transaction::<(), Error, _>(|| {
        let keys = db::get_raw_data_keys(conn)?;
        tasks = keys.len();
        for key in keys {
            db::insert_task(
                conn,
                &models::NewTask {
                    start_at: now,
                    state: models::TaskState::NEW,
                    username: key.username.as_str(),
                    payload: &to_value(Command::ProcessRawData(key.clone()))?,
                },
            )?;
        }
        task.state = models::TaskState::SUCCESSFUL;
        db::update_task(conn, task)?;
        Ok(())
    })?;
    info!("process_all_raw_data: Added {} tasks", tasks);
    Ok(())
}

fn import_strava_user<S: strava::StravaApi>(
    shared: Arc<ImporterSharedData<S>>,
    username: &str,
) -> Result<(), Error> {
    let strava = shared.strava.read().unwrap();
    let user = db::get_user(&shared.conn.lock().unwrap(), username)?;

    let token = {
        let conn = shared.conn.lock().unwrap();
        get_or_refresh_token(&*strava, &conn, &user)?
    };

    let per_page = 30;
    for page in 1.. {
        let params = [
            ("page", &format!("{}", page)[..]),
            ("per_page", &format!("{}", per_page)[..]),
        ];

        let result = strava.get("/athlete/activities", &token.access_token, &params[..])?;

        let json_error = || Error::UnexpectedJson(result.clone());
        let result = result.as_array().ok_or_else(json_error)?;

        for activity in result {
            let id = activity["id"].as_i64().ok_or_else(json_error)?;
            info!("activity id: {} start: {}", id, activity["start_date"]);

            db::insert_data(
                &shared.conn.lock().unwrap(),
                &models::RawData {
                    data_type: models::DataType::StravaActivity,
                    id: id,
                    username: username.to_string(),
                    payload: activity.clone(),
                },
            )?;
        }

        if result.len() < per_page {
            break;
        }
        thread::sleep(Duration::from_secs(1));
    }

    Err(Error::InternalError)
}

fn get_or_refresh_token<Strava: strava::StravaApi>(
    strava: &Strava,
    conn: &PgConnection,
    user: &models::User,
) -> Result<models::StravaToken, Error> {
    let mut token = db::get_strava_token(&conn, &user).expect("FIX");

    if token.expires_at < Utc::now() {
        info!("refresh expired token: {:?}", token.expires_at);
        let new_token = strava.refresh_token(&From::from(&token))?;
        new_token.update_model(&mut token);
    }

    Ok(token)
}