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

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

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},
    title::{
        MaiVersionExt, Sdgb1_50,
        helper::get_user_all_music,
        methods::APIMethod,
        model::{
            DataVersion, DxCalculatedEntries, DxMusicRecord, DxRatingNet, GetUserDataApi,
            GetUserDataApiResp, GetUserPreviewApi, GetUserPreviewApiResp, GetUserRatingApi,
            GetUserRatingApiResp, Ping, PingResp, UserLogoutApi, UserLogoutApiResp,
        },
    },
};
use spdlog::{error, info, warn};

use crate::{
    commands::{Cli, Commands, 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::Always)
            .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 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::MusicDetail { user_id, format } => {
            let music_detail = get_user_all_music(&client, user_id).await?;

            match (human_readable, format) {
                (true, _) => {
                    for detail in &music_detail {
                        println!("{detail}");
                        println!("----------");
                    }

                    println!("共查询到 {} 条记录！", music_detail.len());
                }
                (false, RatingFormat::Origin) => json_display(music_detail)?,
                (false, RatingFormat::DxRatingNet) => {
                    let dx_export = Vec::from_iter(
                        music_detail
                            .iter()
                            .map(|music| {
                                DxMusicRecord::try_from(music).inspect_err(|e| {
                                    warn!("failed to process {}: {e}", music.music_id)
                                })
                            })
                            .flatten(),
                    );
                    json_display(dx_export)?;
                }
                (_, format) => {
                    error!("{format:?} was not supported yet");
                    json_display(())?;
                }
            }
        }
        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::Logout { user_id, timestamp } => {
            let logout: UserLogoutApiResp = Sdgb1_50::request(
                &client,
                APIMethod::UserLogoutApi,
                user_id,
                UserLogoutApi {
                    user_id,
                    date_time: timestamp,
                    ..Default::default()
                },
            )
            .await?;

            if human_readable {
                println!("啥都木有");
            } else {
                json_display(logout)?;
            }
        }
        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::Ping => {
            let time = SystemTime::now();
            let decoded: PingResp =
                Sdgb1_50::request(&client, APIMethod::Ping, "", Ping {}).await?;
            info!(
                "sdgb 1.50 resp: {decoded}, {}ms",
                time.elapsed().unwrap_or_default().as_millis()
            );
        }
        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::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::ListAllUser { concurrency } => {
            use crate::{cache::PLAYERS, utils::helpers::cached_concurrent_fetch};
            use sdgb_api::title::methods::GetUserPreviewApiExt;
            use std::io::BufRead as _;

            let mut user_ids = Vec::new();
            {
                let mut stdin = std::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);
                }
            }

            cached_concurrent_fetch::<GetUserPreviewApiExt>(
                user_ids,
                &client,
                concurrency,
                PLAYERS,
            )
            .await?;
        }

        #[cfg(feature = "fetchall")]
        Commands::ScrapeAllB50 {
            concurrency,
            min_rating,
            max_rating,
        } => {
            use sdgb_api::title::methods::GetUserRatingApiExt;

            use crate::{
                cache::{PLAYER_B50, PLAYERS},
                utils::helpers::{cached_concurrent_fetch, read_cache},
            };

            let mut players: Vec<GetUserPreviewApiResp> = read_cache(PLAYERS)?;
            players.retain(|p| p.player_rating >= min_rating && p.player_rating <= max_rating);

            cached_concurrent_fetch::<GetUserRatingApiExt>(
                players.iter().map(|p| p.user_id).collect::<Vec<u32>>(),
                &client,
                concurrency,
                PLAYER_B50,
            )
            .await?;
        }

        #[cfg(feature = "fetchall")]
        Commands::ListAllUserDump {} => {
            use crate::{cache::PLAYERS, utils::helpers::dump_cache};

            dump_cache::<GetUserPreviewApiResp>("players.json", PLAYERS)?;
        }
        #[cfg(feature = "fetchall")]
        Commands::ScrapeAllB50Dump {} => {
            use crate::{cache::PLAYER_B50, utils::helpers::dump_cache};

            dump_cache::<GetUserRatingApiResp>("b50.json", PLAYER_B50)?;
        }

        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`");