BlueArchiveToolkit/infrastructure/examples/official_pull_plan.rs

use bat_adapters::official::inventory::{
    YostarJpPlatformCatalogInventory, YostarJpPlatformDownloadInventory,
};
use bat_adapters::official::yostar_jp::{
    server_info_url, PatchPlatform, YostarJpResourceDiscoveryPlan, YostarJpResourceEndpoint,
    YostarJpResourceEndpointKind, YostarJpServerInfo,
};
use bat_infrastructure::{
    build_official_pull_plan_for_platform_inventory,
    build_official_pull_plan_from_platform_inventory, default_official_platforms,
    OfficialGameMainConfigBootstrapService, OfficialResourcePullService,
};
use std::collections::HashMap;
use std::env;
use std::fs;
use std::path::PathBuf;

#[derive(Debug)]
struct CliArgs {
    server_info_source: Option<ServerInfoSource>,
    connection_group: Option<String>,
    app_version: Option<String>,
    launcher_version: String,
    launcher_bootstrap: bool,
    auto_discover: bool,
    platforms: Option<Vec<PatchPlatform>>,
    output_root: PathBuf,
    curl_command: PathBuf,
    dry_run: bool,
}

#[derive(Debug)]
enum ServerInfoSource {
    LocalPath(PathBuf),
    OfficialFile(String),
    OfficialUrl(String),
}

fn main() -> anyhow::Result<()> {
    let args = parse_args()?;
    let default_platforms = default_official_platforms();
    let platforms = args
        .platforms
        .as_deref()
        .unwrap_or(default_platforms.as_slice());
    let metadata_bootstrap_enabled = args.auto_discover || args.launcher_bootstrap;
    let official_bootstrap = if metadata_bootstrap_enabled {
        Some(
            OfficialGameMainConfigBootstrapService::with_commands(
                &args.launcher_version,
                args.curl_command.to_string_lossy().to_string(),
                "unzip",
            )
            .fetch_bootstrap()
            .map_err(anyhow::Error::msg)?,
        )
    } else {
        None
    };
    let app_version = args
        .app_version
        .clone()
        .or_else(|| {
            official_bootstrap
                .as_ref()
                .map(|bootstrap| bootstrap.game_config.game_latest_version.clone())
        })
        .ok_or_else(|| anyhow::anyhow!("missing --app-version; production mode does not infer the game version from official metadata unless --auto-discover is set"))?;
    let connection_group = args
        .connection_group
        .clone()
        .or_else(|| {
            official_bootstrap
                .as_ref()
                .and_then(|bootstrap| bootstrap.game_main_config.default_connection_group.clone())
        })
        .ok_or_else(|| {
            anyhow::anyhow!(
                "missing --connection-group; production mode does not infer the connection group from GameMainConfig unless --auto-discover is set"
            )
        })?;

    let fetcher =
        OfficialResourcePullService::with_curl_command(&args.output_root, &args.curl_command);

    if metadata_bootstrap_enabled {
        let bootstrap = official_bootstrap.as_ref().ok_or_else(|| {
            anyhow::anyhow!("metadata bootstrap requested but official bootstrap is unavailable")
        })?;
        println!(
            "metadata_bootstrap_game_latest_version={}",
            bootstrap.game_config.game_latest_version
        );
        println!(
            "metadata_bootstrap_game_latest_file_path={}",
            bootstrap.game_config.game_latest_file_path
        );
        println!(
            "metadata_bootstrap_game_lowest_version={}",
            bootstrap
                .game_config
                .game_lowest_version
                .as_deref()
                .unwrap_or("<none>")
        );
        println!(
            "metadata_bootstrap_game_start_exe_name={}",
            bootstrap
                .game_config
                .game_start_exe_name
                .as_deref()
                .unwrap_or("<none>")
        );
        println!(
            "metadata_bootstrap_game_start_params={:?}",
            bootstrap.game_config.game_start_params
        );
        println!(
            "metadata_bootstrap_primary_cdn={}",
            bootstrap.cdn_config.primary_cdn
        );
        println!(
            "metadata_bootstrap_backup_cdn={}",
            bootstrap.cdn_config.back_up_cdn
        );
        println!("metadata_bootstrap_manifest_url={}", bootstrap.manifest_url);
        println!(
            "metadata_bootstrap_manifest_source={}",
            bootstrap.manifest_source.as_deref().unwrap_or("<none>")
        );
        println!(
            "metadata_bootstrap_manifest_file_count={}",
            bootstrap.manifest_file_count
        );
        println!(
            "metadata_bootstrap_game_main_config_server_info_url={}",
            bootstrap
                .game_main_config
                .server_info_data_url
                .as_deref()
                .unwrap_or("<none>")
        );
        println!(
            "metadata_bootstrap_game_main_config_default_connection_group={}",
            bootstrap
                .game_main_config
                .default_connection_group
                .as_deref()
                .unwrap_or("<none>")
        );
        println!();
    }

    let server_info_source = args.server_info_source.as_ref();

    let server_info_bytes = if let Some(source) = server_info_source {
        load_server_info(source, &fetcher)?
    } else {
        let bootstrap = official_bootstrap.as_ref().ok_or_else(|| {
            anyhow::anyhow!("missing server-info source; pass --server-info-url, --server-info-file, or --server-info-path, or explicitly opt in with --auto-discover")
        })?;
        let url = bootstrap
            .game_main_config
            .server_info_data_url
            .as_deref()
            .ok_or_else(|| anyhow::anyhow!("official GameMainConfig has no ServerInfoDataUrl"))?;
        fetcher.fetch_bytes(url).map_err(anyhow::Error::msg)?
    };
    let server_info = YostarJpServerInfo::from_slice(&server_info_bytes)?;
    let discovery = server_info
        .discovery_plan(&connection_group, &app_version, platforms)
        .map_err(anyhow::Error::msg)?;
    let seed_catalogs = fetch_seed_catalogs(&fetcher, &discovery)?;
    let inventory = build_inventory_from_seed_catalogs(&seed_catalogs, platforms)?;
    let plan = if args.platforms.is_some() {
        build_official_pull_plan_for_platform_inventory(discovery, inventory, platforms)
    } else {
        build_official_pull_plan_from_platform_inventory(discovery, inventory)
    };
    let all_urls = plan.all_urls().map_err(anyhow::Error::msg)?;

    println!("connection_group={}", plan.discovery.connection_group_name);
    println!("app_version={}", plan.discovery.app_version);
    println!(
        "bundle_version={}",
        plan.discovery.bundle_version.as_deref().unwrap_or("<none>")
    );
    println!("addressables_root={}", plan.discovery.addressables_root);
    println!("platforms={:?}", plan.platforms);
    println!("discovery_url_count={}", plan.discovery_urls().len());
    println!(
        "content_url_count={}",
        plan.content_urls().map_err(anyhow::Error::msg)?.len()
    );
    println!("total_url_count={}", all_urls.len());
    println!("output_root={}", args.output_root.display());
    println!("dry_run={}", args.dry_run);
    println!();

    for url in &all_urls {
        println!("{url}");
    }

    if args.dry_run {
        return Ok(());
    }

    let downloader =
        OfficialResourcePullService::with_curl_command(args.output_root, args.curl_command);
    let report = downloader.pull(&plan).map_err(anyhow::Error::msg)?;
    eprintln!("resource_count={}", report.items.len());
    eprintln!("downloaded_count={}", report.downloaded_count());
    eprintln!("resumed_count={}", report.resumed_count());
    eprintln!("skipped_count={}", report.skipped_count());
    eprintln!("final_bytes={}", report.total_bytes());
    eprintln!("transferred_bytes={}", report.transferred_bytes());
    eprintln!(
        "official_hash_verified_count={}",
        report.official_hash_verified_count()
    );
    eprintln!(
        "download_manifest={}",
        downloader.download_manifest_path().display()
    );

    Ok(())
}

fn load_server_info(
    source: &ServerInfoSource,
    fetcher: &OfficialResourcePullService,
) -> anyhow::Result<Vec<u8>> {
    match source {
        ServerInfoSource::LocalPath(path) => Ok(fs::read(path)?),
        ServerInfoSource::OfficialFile(file_name) => {
            let url = server_info_url(file_name).map_err(anyhow::Error::msg)?;
            fetcher.fetch_bytes(&url).map_err(anyhow::Error::msg)
        }
        ServerInfoSource::OfficialUrl(url) => fetcher.fetch_bytes(url).map_err(anyhow::Error::msg),
    }
}

fn fetch_seed_catalogs(
    fetcher: &OfficialResourcePullService,
    discovery: &YostarJpResourceDiscoveryPlan,
) -> anyhow::Result<SeedCatalogs> {
    let mut catalogs = SeedCatalogs::default();

    for endpoint in &discovery.endpoints {
        if !matches!(
            endpoint.kind,
            YostarJpResourceEndpointKind::TableCatalog
                | YostarJpResourceEndpointKind::BundlePackingInfo
                | YostarJpResourceEndpointKind::MediaCatalog
        ) {
            continue;
        }

        let bytes = fetcher
            .fetch_bytes(&endpoint.url)
            .map_err(anyhow::Error::msg)?;
        catalogs.insert(endpoint, bytes)?;
    }

    Ok(catalogs)
}

fn build_inventory_from_seed_catalogs(
    catalogs: &SeedCatalogs,
    platforms: &[PatchPlatform],
) -> anyhow::Result<YostarJpPlatformDownloadInventory> {
    let table_catalog = catalogs
        .table_catalog
        .as_deref()
        .ok_or_else(|| anyhow::anyhow!("official discovery did not include TableCatalog.bytes"))?;
    let mut platform_catalogs = Vec::new();

    for platform in platforms {
        let bundle_packing_info = catalogs.bundle_packing_infos.get(platform).ok_or_else(|| {
            anyhow::anyhow!(
                "official discovery did not include {} BundlePackingInfo.bytes",
                platform.as_str()
            )
        })?;
        let media_catalog = catalogs.media_catalogs.get(platform).ok_or_else(|| {
            anyhow::anyhow!(
                "official discovery did not include {} MediaCatalog.bytes",
                platform.as_str()
            )
        })?;

        platform_catalogs.push(YostarJpPlatformCatalogInventory::from_catalog_bytes(
            *platform,
            bundle_packing_info,
            media_catalog,
        ));
    }

    Ok(YostarJpPlatformDownloadInventory::from_catalog_bytes(
        table_catalog,
        platform_catalogs,
    ))
}

#[derive(Debug, Default)]
struct SeedCatalogs {
    table_catalog: Option<Vec<u8>>,
    bundle_packing_infos: HashMap<PatchPlatform, Vec<u8>>,
    media_catalogs: HashMap<PatchPlatform, Vec<u8>>,
}

impl SeedCatalogs {
    fn insert(
        &mut self,
        endpoint: &YostarJpResourceEndpoint,
        bytes: Vec<u8>,
    ) -> anyhow::Result<()> {
        match endpoint.kind {
            YostarJpResourceEndpointKind::TableCatalog => {
                self.table_catalog = Some(bytes);
            }
            YostarJpResourceEndpointKind::BundlePackingInfo => {
                let platform = required_platform(endpoint)?;
                self.bundle_packing_infos.insert(platform, bytes);
            }
            YostarJpResourceEndpointKind::MediaCatalog => {
                let platform = required_platform(endpoint)?;
                self.media_catalogs.insert(platform, bytes);
            }
            _ => {}
        }

        Ok(())
    }
}

fn required_platform(endpoint: &YostarJpResourceEndpoint) -> anyhow::Result<PatchPlatform> {
    endpoint.platform.ok_or_else(|| {
        anyhow::anyhow!(
            "discovery endpoint {:?} has no platform: {}",
            endpoint.kind,
            endpoint.url
        )
    })
}

fn parse_args() -> anyhow::Result<CliArgs> {
    parse_args_from(env::args())
}

fn parse_args_from(raw_args: impl IntoIterator<Item = String>) -> anyhow::Result<CliArgs> {
    let mut args = raw_args.into_iter();
    let binary = args
        .next()
        .unwrap_or_else(|| "official_pull_plan".to_string());
    let mut parsed = CliArgs {
        server_info_source: None,
        connection_group: None,
        app_version: None,
        launcher_version: "1.7.2".to_string(),
        launcher_bootstrap: false,
        auto_discover: false,
        platforms: None,
        output_root: PathBuf::from("official_pull_output"),
        curl_command: PathBuf::from("curl"),
        dry_run: false,
    };

    while let Some(flag) = args.next() {
        match flag.as_str() {
            "--server-info-url" => {
                parsed.server_info_source = Some(ServerInfoSource::OfficialUrl(next_option_value(
                    &mut args, &flag,
                )?));
            }
            "--server-info-file" => {
                parsed.server_info_source = Some(ServerInfoSource::OfficialFile(
                    next_option_value(&mut args, &flag)?,
                ));
            }
            "--server-info-path" => {
                parsed.server_info_source = Some(ServerInfoSource::LocalPath(PathBuf::from(
                    next_option_value(&mut args, &flag)?,
                )));
            }
            "--connection-group" => {
                parsed.connection_group = Some(next_option_value(&mut args, &flag)?);
            }
            "--app-version" => {
                parsed.app_version = Some(next_option_value(&mut args, &flag)?);
            }
            "--launcher-version" => {
                parsed.launcher_version = next_option_value(&mut args, &flag)?;
            }
            "--launcher-bootstrap" => {
                parsed.launcher_bootstrap = true;
            }
            "--auto-discover" => {
                parsed.auto_discover = true;
            }
            "--platforms" => {
                let value = next_option_value(&mut args, &flag)?;
                parsed.platforms = Some(parse_platforms(&value).map_err(anyhow::Error::msg)?);
            }
            "--output" => {
                parsed.output_root = PathBuf::from(next_option_value(&mut args, &flag)?);
            }
            "--curl" => {
                parsed.curl_command = PathBuf::from(next_option_value(&mut args, &flag)?);
            }
            "--dry-run" => {
                parsed.dry_run = true;
            }
            "--help" | "-h" => {
                print_usage(&binary);
                std::process::exit(0);
            }
            _ => return Err(anyhow::anyhow!("unknown argument: {flag}")),
        }
    }

    Ok(parsed)
}

fn next_option_value(
    args: &mut impl Iterator<Item = String>,
    flag: &str,
) -> anyhow::Result<String> {
    args.next()
        .ok_or_else(|| anyhow::anyhow!("missing value for {flag}"))
}

fn print_usage(binary: &str) {
    eprintln!(
        "usage: {binary} [--server-info-url <official-url> | --server-info-file <official-json-name> | --server-info-path <local-json>] \
         [--app-version 1.70.0] \
         [--connection-group <name>] [--auto-discover] [--launcher-bootstrap] [--launcher-version 1.7.2] \
         [--platforms Windows,Android] [--output official_pull_output] [--curl curl] [--dry-run]"
    );
}

fn parse_platforms(value: &str) -> Result<Vec<PatchPlatform>, String> {
    value
        .split(',')
        .map(str::trim)
        .filter(|part| !part.is_empty())
        .map(parse_platform)
        .collect()
}

fn parse_platform(value: &str) -> Result<PatchPlatform, String> {
    match value.to_ascii_lowercase().as_str() {
        "windows" | "win" => Ok(PatchPlatform::Windows),
        "android" => Ok(PatchPlatform::Android),
        _ => Err(format!("unsupported platform: {value}")),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn parse(values: &[&str]) -> anyhow::Result<CliArgs> {
        parse_args_from(values.iter().map(|value| value.to_string()))
    }

    #[test]
    fn parses_explicit_official_server_info_url_without_launcher_bootstrap() {
        let args = parse(&[
            "official_pull_plan",
            "--server-info-url",
            "https://yostar-serverinfo.bluearchiveyostar.com/r93_fixture.json",
            "--connection-group",
            "Prod-Audit",
            "--app-version",
            "1.70.0",
            "--dry-run",
        ])
        .unwrap();

        assert!(!args.launcher_bootstrap);
        assert_eq!(args.connection_group.as_deref(), Some("Prod-Audit"));
        assert_eq!(args.app_version.as_deref(), Some("1.70.0"));
        assert!(matches!(
            args.server_info_source,
            Some(ServerInfoSource::OfficialUrl(ref url))
                if url == "https://yostar-serverinfo.bluearchiveyostar.com/r93_fixture.json"
        ));
    }

    #[test]
    fn auto_discover_is_explicit_opt_in() {
        let args = parse(&["official_pull_plan", "--auto-discover", "--dry-run"]).unwrap();

        assert!(args.auto_discover);
        assert!(!args.launcher_bootstrap);
        assert!(args.server_info_source.is_none());
        assert!(args.connection_group.is_none());
        assert!(args.app_version.is_none());
    }

    #[test]
    fn launcher_bootstrap_alias_is_explicit_opt_in() {
        let args = parse(&["official_pull_plan", "--launcher-bootstrap", "--dry-run"]).unwrap();

        assert!(!args.auto_discover);
        assert!(args.launcher_bootstrap);
        assert!(args.server_info_source.is_none());
        assert!(args.connection_group.is_none());
        assert!(args.app_version.is_none());
    }
}