sdgb-utils-rs/sdgb-cli/src/main.rs

use std::sync::{
    Arc,
    atomic::{AtomicBool, Ordering},
};

use nyquest_preset::nyquest::ClientBuilder;
use palc::Parser;
use spdlog::{Level, LevelFilter::MoreSevereEqual, sink::StdStreamSink, terminal_style::StyleMode};

use sdgb_api::{
    all_net::QRCode,
    auth_lite::{SDGB, SDHJ, delivery_raw},
    helper::MUSIC_DB,
    title::{
        MaiVersionExt, Sdgb1_50,
        methods::APIMethod,
        model::{
            DataVersion, DxCalculatedEntries, DxRatingNet, GetUserDataApi, GetUserDataApiResp,
            GetUserMusicApi, GetUserMusicApiResp, GetUserPreviewApi, GetUserPreviewApiResp,
            GetUserRatingApi, GetUserRatingApiResp, Ping, PingResp, UserLogoutApi,
            UserLogoutApiResp,
        },
    },
};
use spdlog::{error, info, warn};

use crate::{
    commands::{Cli, RatingFormat},
    utils::{human_readable_display, json_display, login_action},
};

#[cfg(feature = "fetchall")]
mod cache;
mod commands;
mod utils;

static EARLY_QUIT: AtomicBool = AtomicBool::new(false);

#[cfg_attr(feature = "compio", compio::main)]
#[cfg_attr(feature = "tokio", tokio::main)]
async fn main() -> Result<(), Box<dyn snafu::Error>> {
    nyquest_preset::register();

    let logger = spdlog::default_logger().fork_with(|log| {
        log.set_level_filter(MoreSevereEqual(if cfg!(debug_assertions) {
            Level::Debug
        } else {
            Level::Info
        }));
        let sink = StdStreamSink::builder()
            .stderr()
            .style_mode(StyleMode::Auto)
            .build()?;
        *log.sinks_mut() = vec![Arc::new(sink)];
        Ok(())
    })?;
    spdlog::swap_default_logger(logger);

    ctrlc::set_handler(|| {
        if EARLY_QUIT.load(Ordering::Relaxed) {
            error!("force-quit triggered!");
            std::process::exit(1);
        } else {
            warn!("received early-quit request! will abort soon");
            EARLY_QUIT.store(true, Ordering::Relaxed);
        }
    })?;

    let _ = &*MUSIC_DB;

    let Cli {
        command,
        machine_readable,
    } = <Cli as Parser>::parse();
    let human_readable = !machine_readable;

    let client = ClientBuilder::default().build_async().await?;

    // TODO: refactor via enum_dispatch
    match command {
        commands::Commands::MusicDetail { user_id } => {
            let mut music_detail = Vec::new();
            let mut index = None;

            loop {
                let GetUserMusicApiResp {
                    next_index,
                    mut user_music_list,
                    ..
                } = Sdgb1_50::request::<_, GetUserMusicApiResp>(
                    &client,
                    APIMethod::GetUserMusicApi,
                    user_id,
                    GetUserMusicApi {
                        user_id,
                        next_index: index.unwrap_or_default(),
                        max_count: 50,
                    },
                )
                .await?;

                for list in &mut user_music_list {
                    music_detail.append(&mut list.user_music_detail_list);
                }

                if next_index == 0 || user_music_list.is_empty() {
                    break;
                }
                index = Some(next_index);
            }

            // TODO: `Display` support for MusicDetail
            json_display(music_detail)?;
        }
        commands::Commands::Rating { user_id, format } => {
            let rating: GetUserRatingApiResp = Sdgb1_50::request(
                &client,
                APIMethod::GetUserRatingApi,
                user_id,
                GetUserRatingApi { user_id },
            )
            .await?;

            match (human_readable, format) {
                (true, _) => println!("{rating}"),
                (false, RatingFormat::Origin) => json_display(rating)?,
                (false, RatingFormat::DxRatingNet) => {
                    let mut data = DxCalculatedEntries::from_user_rating_lossy(&rating.user_rating);
                    let mut records = data.b35;
                    records.append(&mut data.b15);
                    json_display(records)?;
                }
                (false, RatingFormat::DxRatingPayload) => {
                    let data = DxCalculatedEntries::from_user_rating_lossy(&rating.user_rating);
                    let payload = DxRatingNet {
                        calculated_entries: data,
                        version: DataVersion::Prism,
                        region: "_generic",
                    };
                    json_display(payload)?;
                }
            }
        }
        commands::Commands::Logout { user_id, timestamp } => {
            let logout: UserLogoutApiResp = Sdgb1_50::request(
                &client,
                APIMethod::UserLogoutApi,
                user_id,
                UserLogoutApi {
                    user_id,
                    date_time: timestamp,
                    ..Default::default()
                },
            )
            .await?;
            println!("{logout:?}");
        }
        commands::Commands::Preview { user_id } => {
            let preview: GetUserPreviewApiResp = Sdgb1_50::request(
                &client,
                APIMethod::GetUserPreviewApi,
                user_id,
                GetUserPreviewApi { user_id },
            )
            .await?;

            human_readable_display(preview, human_readable)?;
        }
        commands::Commands::Ping => {
            let decoded: PingResp =
                Sdgb1_50::request(&client, APIMethod::Ping, "", Ping {}).await?;
            info!("sdgb 1.50 resp: {decoded}");
        }
        commands::Commands::QRLogin { ref qrcode_content } => {
            let qrcode = QRCode { qrcode_content };
            let resp = qrcode.login(&client).await;

            match &resp {
                Ok(user_id) => info!("login succeed: {user_id}"),
                Err(e) => error!("login failed: {e}"),
            }

            if !human_readable {
                json_display(resp)?;
            }
        }

        commands::Commands::AuthLite { title_ver, variant } => {
            let resp = match variant {
                commands::AuthLiteVariant::SDGB => delivery_raw::<SDGB>(&client, title_ver).await?,
                commands::AuthLiteVariant::SDHJ => delivery_raw::<SDHJ>(&client, title_ver).await?,
            };
            println!("{}", String::from_utf8_lossy(&resp));
        }

        #[cfg(feature = "fetchall")]
        commands::Commands::ListAllUser { concurrency } => {
            use futures_util::StreamExt;
            use sdgb_api::bincode::borrow_decode_from_slice;
            use std::io::{self, BufRead};

            let mut user_ids = Vec::new();
            {
                let mut stdin = io::stdin().lock();
                let mut buf = String::new();

                while stdin.read_line(&mut buf).is_ok_and(|size| size != 0) {
                    if buf.is_empty() {
                        continue;
                    }

                    let user_id: u32 = buf.trim().parse()?;
                    buf.clear();
                    user_ids.push(user_id);
                }
            }

            let _ = cache::init_db();
            let read = cache::read_txn()?;
            let write = cache::write_txn()?;
            let config = sdgb_api::bincode::config::Configuration::<
                sdgb_api::bincode::config::LittleEndian,
            >::default()
            .with_no_limit();

            info!("number of user_id: {}", user_ids.len());

            let collect = futures_util::stream::iter(user_ids)
                .map(async |user_id| {
                    {
                        let cache_table = cache::open_table_read(&read)?;
                        let data = cache_table.get(user_id)?;
                        if let Some(data) = data {
                            let decoded: (GetUserPreviewApiResp, _) =
                                borrow_decode_from_slice(&data.value(), config)?;

                            return Ok(decoded.0);
                        }
                    }

                    if EARLY_QUIT.load(Ordering::Relaxed) {
                        return Err("early skip due to ctrl-c")?;
                    }

                    let resp = Sdgb1_50::request::<_, GetUserPreviewApiResp>(
                        &client,
                        APIMethod::GetUserPreviewApi,
                        user_id,
                        GetUserPreviewApi { user_id },
                    )
                    .await;

                    match &resp {
                        Ok(resp) => {
                            use sdgb_api::bincode::encode_to_vec;

                            info!("found: {user_id}");

                            if let Ok(mut table) = cache::open_table(&write)
                                && let Ok(encoded) = encode_to_vec(resp, config)
                            {
                                _ = table.insert(resp.user_id, encoded);
                            }
                        }
                        Err(sdgb_api::ApiError::JSON { .. }) => {}
                        Err(e) => {
                            error!("preview failed: {e}");
                        }
                    }

                    Result::<_, Box<dyn snafu::Error>>::Ok(resp?)
                })
                .buffer_unordered(concurrency) // slower to avoid being banned
                .filter_map(async |r| r.ok())
                .collect::<Vec<_>>()
                .await;
            drop(collect);

            let _ = write.commit();
        }
        #[cfg(feature = "fetchall")]
        commands::Commands::ListAllUserDump { .. } => {
            use std::{fs::OpenOptions, io::BufWriter};

            use redb::ReadableTable;
            use sdgb_api::bincode::{self, borrow_decode_from_slice};

            use crate::cache::{open_table_read, read_txn};

            let file = OpenOptions::new()
                .create(true)
                .truncate(true)
                .write(true)
                .open("players.json")?;

            #[cfg(file_lock_ready)]
            file.try_lock()?;

            let txn = read_txn()?;
            let table = open_table_read(&txn)?;

            let config = bincode::config::Configuration::<bincode::config::LittleEndian>::default()
                .with_no_limit();

            let user_ids = table
                .iter()?
                .flatten()
                .map(|d| borrow_decode_from_slice(&d.1.value(), config))
                .flatten()
                .map(|(value, _)| value)
                .collect::<Vec<GetUserPreviewApiResp>>();

            let writer = BufWriter::new(file);
            serde_json::to_writer(writer, &user_ids)?;
            info!("dumped {} user id", user_ids.len());
        }

        commands::Commands::Userdata { user_id } => {
            let action = async |_| match Sdgb1_50::request::<_, GetUserDataApiResp>(
                &client,
                APIMethod::GetUserDataApi,
                user_id,
                GetUserDataApi { user_id },
            )
            .await
            {
                Ok(udata) => {
                    info!("{udata:#?}");
                }
                Err(e) => {
                    error!("failed to get userdata: {e}");
                }
            };
            login_action(&client, user_id, action).await?;
        }
    }

    Ok(())
}

#[cfg(all(feature = "compio", feature = "tokio"))]
compile_error!("you must not enable both `compio` and `tokio`");