cli and full earth export

3 files changed, 221 insertions, 51 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 9d99b4d..cc06028 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -194,6 +194,46 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
 
 [[package]]
+name = "clap"
+version = "4.5.35"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d8aa86934b44c19c50f87cc2790e19f54f7a67aedb64101c2e1a2e5ecfb73944"
+dependencies = [
+ "clap_builder",
+ "clap_derive",
+]
+
+[[package]]
+name = "clap_builder"
+version = "4.5.35"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2414dbb2dd0695280da6ea9261e327479e9d37b0630f6b53ba2a11c60c679fd9"
+dependencies = [
+ "anstream",
+ "anstyle",
+ "clap_lex",
+ "strsim",
+]
+
+[[package]]
+name = "clap_derive"
+version = "4.5.32"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "09176aae279615badda0765c0c0b3f6ed53f4709118af73cf4655d85d1530cd7"
+dependencies = [
+ "heck",
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "clap_lex"
+version = "0.7.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f46ad14479a25103f283c0f10005961cf086d8dc42205bb44c46ac563475dca6"
+
+[[package]]
 name = "colorchoice"
 version = "1.0.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -344,6 +384,17 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "05f29059c0c2090612e8d742178b0580d2dc940c837851ad723096f87af6663e"
 
 [[package]]
+name = "futures-macro"
+version = "0.3.31"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "162ee34ebcb7c64a8abebc059ce0fee27c2262618d7b60ed8faf72fef13c3650"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
+[[package]]
 name = "futures-sink"
 version = "0.3.31"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -362,9 +413,11 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9fa08315bb612088cc391249efdc3bc77536f16c91f6cf495e6fbe85b20a4a81"
 dependencies = [
  "futures-core",
+ "futures-macro",
  "futures-task",
  "pin-project-lite",
  "pin-utils",
+ "slab",
 ]
 
 [[package]]
@@ -1408,6 +1461,12 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a8f112729512f8e442d81f95a8a7ddf2b7c6b8a1a6f509a95864142b30cab2d3"
 
 [[package]]
+name = "strsim"
+version = "0.11.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f"
+
+[[package]]
 name = "subtle"
 version = "2.6.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -1761,7 +1820,9 @@ name = "weareearth"
 version = "0.1.0"
 dependencies = [
  "anyhow",
+ "clap",
  "env_logger",
+ "futures-util",
  "glam",
  "log",
  "prost",
diff --git a/Cargo.toml b/Cargo.toml
index becffc1..887db88 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -13,6 +13,8 @@ prost-types = "0.13.5"
 prost = "0.13.5"
 xdg = "2.5.2"
 glam = "0.30.1"
+futures-util = "0.3.31"
+clap = { version = "4.5.35", features = ["derive"] }
 
 weareshared = { path = "../wearechat/shared" }
 
diff --git a/src/main.rs b/src/main.rs
index 51c30d7..23d840f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -2,8 +2,10 @@
 pub mod mesh;
 
 use anyhow::{Result, bail};
-use glam::{DAffine3, DMat4};
-use log::{debug, error};
+use clap::Parser;
+use futures_util::{StreamExt, lock::Mutex, stream::FuturesUnordered};
+use glam::DMat4;
+use log::{debug, error, info};
 use mesh::{convert_mesh, decode_normal_table};
 use prost::{Message, bytes::Bytes};
 use proto::{BulkMetadata, NodeData, NodeMetadata, PlanetoidMetadata};
@@ -11,81 +13,166 @@ use reqwest::{
     Client,
     header::{HeaderMap, HeaderName, HeaderValue},
 };
-use std::{f32::consts::PI, path::PathBuf};
+use std::{f32::consts::PI, path::PathBuf, pin::Pin, sync::Arc};
 use tokio::{
     fs::{File, create_dir_all},
     io::{AsyncReadExt, AsyncWriteExt},
+    sync::Semaphore,
 };
 use weareshared::{
     Affine3A,
-    resources::{Prefab, RespackEntry},
+    packets::Resource,
+    resources::{MeshPart, Prefab, RespackEntry},
     respack::save_full_respack,
     store::ResourceStore,
 };
 
+#[derive(Parser)]
+struct Args {
+    #[arg(short, long, default_value = "16")]
+    par_limit: usize,
+    #[clap(subcommand)]
+    action: Action,
+}
+#[derive(Parser)]
+enum Action {
+    Cache { level: usize },
+    Export { level: usize },
+}
+
 #[tokio::main]
 async fn main() -> Result<()> {
     env_logger::init_from_env("LOG");
 
-    let c = GeClient::new().await?;
-    let entry = c.planetoid_metdata().await?;
+    let args = Args::parse();
 
-    let bulk = c
-        .bulk_metadata("", entry.root_node_metadata.unwrap().bulk_metadata_epoch())
-        .await?;
+    match args.action {
+        Action::Cache { level } => {
+            let c = GeClient::new(16).await?;
+            let entry = c.planetoid_metdata().await?;
+            let epoch = entry.root_node_metadata.unwrap().bulk_metadata_epoch();
+            cache_all(Arc::new(c), "".to_string(), epoch, level).await?;
+        }
+        Action::Export { level } => {
+            let c = GeClient::new(16).await?;
+            let entry = c.planetoid_metdata().await?;
 
-    let store = ResourceStore::new_memory();
+            let store = Arc::new(ResourceStore::new_memory());
+            let meshes = Arc::new(Mutex::new(Vec::new()));
 
-    let mut prefabs = Vec::new();
-    for node_meta in &bulk.node_metadata {
-        if unpack_path_and_id(node_meta.path_and_flags()).0.len() != 4 {
-            // toplevel octs are not found
-            // also i want most subdivided
-            continue;
+            do_node(
+                Arc::new(c),
+                "".to_string(),
+                meshes.clone(),
+                store.clone(),
+                entry.root_node_metadata.unwrap().bulk_metadata_epoch(),
+                level,
+            )
+            .await?;
+
+            let prefab = store.set(&Prefab {
+                mesh: meshes.lock().await.clone(),
+                ..Default::default()
+            })?;
+            let entry = store.set(&RespackEntry {
+                c_prefab: vec![prefab],
+                ..Default::default()
+            })?;
+            let file = std::fs::File::create("/tmp/a.respack")?;
+            save_full_respack(file, &store, Some(entry))?;
         }
+    }
 
-        let Ok(node_data) = c.node_data(&bulk, &node_meta).await else {
-            continue;
-        };
+    Ok(())
+}
 
-        let transform = DMat4::from_cols_slice(&node_data.matrix_globe_from_mesh);
+fn do_node(
+    c: Arc<GeClient>,
+    path: String,
+    meshes: Arc<Mutex<Vec<(Affine3A, Resource<MeshPart>)>>>,
+    store: Arc<ResourceStore>,
+    epoch: u32,
+    level: usize,
+) -> Pin<Box<dyn Future<Output = Result<()>>>> {
+    Box::pin(async move {
+        let bulk = c.bulk_metadata(&path, epoch).await?;
 
-        let for_normals = decode_normal_table(node_data.for_normals());
+        let mut fu = FuturesUnordered::new();
 
-        let mut meshes = Vec::new();
-        for m in node_data.meshes {
-            let mesh = convert_mesh(m, &store, &for_normals)?;
-            meshes.push((
-                Affine3A::from_rotation_x(-PI / 2.)
-                    * Affine3A::from_mat4((transform / 3_000_000.).as_mat4()),
-                mesh,
-            ))
+        for node_meta in &bulk.node_metadata {
+            let (cpath, flags) = unpack_path_and_id(node_meta.path_and_flags());
+            // eprintln!("{path}+{cpath} {flags:?}");
+            let abspath = format!("{path}{cpath}");
+            if flags.has_node && abspath.len() == level {
+                let node = c.node_data(&abspath, flags, &bulk, node_meta).await?;
+                let transform = DMat4::from_cols_slice(&node.matrix_globe_from_mesh);
+                let for_normals = decode_normal_table(node.for_normals());
+                for m in node.meshes {
+                    let mesh = convert_mesh(m, &store, &for_normals)?;
+                    meshes.lock().await.push((
+                        Affine3A::from_rotation_x(-PI / 2.)
+                            * Affine3A::from_mat4((transform / 500_000.).as_mat4()),
+                        mesh,
+                    ))
+                }
+            }
+            if cpath.len() == 4 && flags.has_metadata && abspath.len() < level {
+                fu.push(do_node(
+                    c.clone(),
+                    abspath,
+                    meshes.clone(),
+                    store.clone(),
+                    epoch,
+                    level,
+                ));
+            }
         }
 
-        prefabs.push(store.set(&Prefab {
-            transform: Some(DAffine3::from_mat4(transform)),
-            mesh: meshes,
-            ..Default::default()
-        })?);
-    }
+        while let Some(res) = fu.next().await {
+            res?;
+        }
+
+        Ok(())
+    })
+}
 
-    let entry = store.set(&RespackEntry {
-        c_prefab: prefabs,
-        ..Default::default()
-    })?;
+fn cache_all(
+    c: Arc<GeClient>,
+    path: String,
+    epoch: u32,
+    level: usize,
+) -> Pin<Box<dyn Future<Output = Result<()>>>> {
+    Box::pin(async move {
+        let bulk = c.bulk_metadata(&path, epoch).await?;
 
-    let file = std::fs::File::create("/tmp/a.respack")?;
-    save_full_respack(file, &store, Some(entry))?;
+        let mut fu = FuturesUnordered::new();
 
-    Ok(())
+        for node_meta in &bulk.node_metadata {
+            let (cpath, flags) = unpack_path_and_id(node_meta.path_and_flags());
+            let abspath = format!("{path}{cpath}");
+            if flags.has_node {
+                c.node_data(&abspath, flags, &bulk, node_meta).await?;
+            }
+            if cpath.len() == 4 && flags.has_metadata && abspath.len() < level {
+                fu.push(cache_all(c.clone(), abspath, epoch, level));
+            }
+        }
+
+        while let Some(res) = fu.next().await {
+            res?;
+        }
+
+        Ok(())
+    })
 }
 
 struct GeClient {
     client: Client,
     cachedir: PathBuf,
+    par_limit: Semaphore,
 }
 impl GeClient {
-    pub async fn new() -> Result<Self> {
+    pub async fn new(par_limit: usize) -> Result<Self> {
         let cachedir = xdg::BaseDirectories::with_prefix("weareearth")
             .unwrap()
             .create_cache_directory("download")
@@ -94,6 +181,7 @@ impl GeClient {
         create_dir_all(cachedir.join("NodeData")).await?;
         Ok(Self {
             cachedir,
+            par_limit: Semaphore::new(par_limit),
             client: Client::builder().default_headers(HeaderMap::from_iter([
                 (HeaderName::from_static("user-agent"), HeaderValue::from_static("Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/135.0.0.0 Safari/537.36")),
                 (HeaderName::from_static("referer"), HeaderValue::from_static("https://earth.google.com/"))
@@ -101,6 +189,7 @@ impl GeClient {
         })
     }
     pub async fn download(&self, path: &str) -> Result<Bytes> {
+        let _permit = self.par_limit.acquire().await?;
         let cachepath = self.cachedir.join(path);
         if cachepath.exists() {
             debug!("cached {path:?}");
@@ -108,7 +197,7 @@ impl GeClient {
             File::open(cachepath).await?.read_to_end(&mut buf).await?;
             Ok(buf.into())
         } else {
-            debug!("download {path:?}");
+            info!("download {path:?}");
             let res = self
                 .client
                 .get(format!("https://kh.google.com/rt/earth/{path}"))
@@ -134,22 +223,26 @@ impl GeClient {
             .await?;
         Ok(BulkMetadata::decode(buf)?)
     }
-    pub async fn node_data(&self, bulk: &BulkMetadata, node: &NodeMetadata) -> Result<NodeData> {
-        let (path, flags) = unpack_path_and_id(node.path_and_flags());
-
+    pub async fn node_data(
+        &self,
+        abspath: &str,
+        flags: Flags,
+        bulk: &BulkMetadata,
+        node: &NodeMetadata,
+    ) -> Result<NodeData> {
         let texture_format = bulk.default_available_texture_formats();
         let imagery_epoch = node.imagery_epoch.unwrap_or(bulk.default_imagery_epoch());
         let node_epoch = node
             .epoch
             .unwrap_or(bulk.head_node_key.as_ref().unwrap().epoch.unwrap()); // ?
 
-        let image_epoch_part = if flags & 16 != 0 {
+        let image_epoch_part = if flags.use_image_epoch {
             format!("!3u{imagery_epoch}")
         } else {
             String::new()
         };
         let url = format!(
-            "NodeData/pb=!1m2!1s{path}!2u{node_epoch}!2e{texture_format}{image_epoch_part}!4b0"
+            "NodeData/pb=!1m2!1s{abspath}!2u{node_epoch}!2e{texture_format}{image_epoch_part}!4b0"
         );
 
         let buf = self.download(&url).await?;
@@ -157,7 +250,14 @@ impl GeClient {
     }
 }
 
-fn unpack_path_and_id(mut path_id: u32) -> (String, u32) {
+#[derive(Debug, Clone, Copy)]
+struct Flags {
+    has_node: bool,
+    has_metadata: bool,
+    use_image_epoch: bool,
+}
+
+fn unpack_path_and_id(mut path_id: u32) -> (String, Flags) {
     let level = 1 + (path_id & 3);
     path_id >>= 2;
     let mut path = String::new();
@@ -166,7 +266,14 @@ fn unpack_path_and_id(mut path_id: u32) -> (String, u32) {
         path_id >>= 3;
     }
     let flags = path_id;
-    (path, flags)
+    (
+        path,
+        Flags {
+            has_node: flags & 8 == 0,
+            has_metadata: flags & 4 == 0,
+            use_image_epoch: flags & 16 != 0,
+        },
+    )
 }
 
 pub mod proto {