1use std::fs::File;
4use std::io::{BufRead, BufReader};
5use std::path::Path;
6use std::path::PathBuf;
7
8use rusqlite::{Connection, OpenFlags};
9
10pub fn geonames_db_path() -> PathBuf {
17 let mut first_existing = None;
18 for root in candidate_geonames_roots() {
20 let p = root.join("data").join("geonames.db");
21 if p.exists() {
22 if geonames_db_is_current(&p) {
23 return p;
24 }
25 first_existing.get_or_insert(p);
26 }
27 }
28 let cwd = PathBuf::from("data").join("geonames.db");
30 if cwd.exists() {
31 if geonames_db_is_current(&cwd) {
32 return cwd;
33 }
34 first_existing.get_or_insert(cwd);
35 }
36 if let Some(existing) = first_existing {
37 return existing;
38 }
39 crate::bootstrap::default_asset_install_dir()
41 .join("data")
42 .join("geonames.db")
43}
44
45fn candidate_geonames_roots() -> Vec<PathBuf> {
49 let mut roots = Vec::new();
50 roots.push(crate::bootstrap::default_asset_install_dir());
51 if let Ok(from_env) = std::env::var("SMRITI_ASSET_DIR") {
52 roots.push(PathBuf::from(from_env));
53 }
54 if let Ok(from_env) = std::env::var("PHOTOVAULT_ASSET_DIR") {
55 roots.push(PathBuf::from(from_env));
56 }
57 if let Ok(exe) = std::env::current_exe() {
58 if let Some(dir) = exe.parent() {
59 roots.push(dir.to_path_buf());
60 roots.push(dir.join("..").join("lib").join("smriti"));
62 roots.push(dir.join("..").join("lib").join("photovault"));
63 roots.push(dir.join("..").join(".."));
65 roots.push(dir.join("..").join("..").join(".."));
66 }
67 }
68 let cwd = crate::bootstrap::project_root();
69 if let Some(parent) = cwd.parent() {
72 roots.push(parent.to_path_buf());
73 if let Some(grand) = parent.parent() {
74 roots.push(grand.to_path_buf());
75 }
76 }
77 roots.push(cwd);
78 roots.push(PathBuf::from("/usr/lib/smriti"));
79 roots.push(PathBuf::from("/usr/lib/photovault"));
80 roots
81}
82
83pub fn geonames_db_exists() -> bool {
85 geonames_schema_is_current()
86}
87
88pub const GEONAMES_SCHEMA_VERSION: i64 = 2;
97
98pub fn geonames_schema_is_current() -> bool {
103 let path = geonames_db_path();
104 geonames_db_is_current(&path)
105}
106
107pub fn geonames_db_is_current(path: &Path) -> bool {
109 if !path.exists() {
110 return false;
111 }
112 let Ok(conn) = Connection::open_with_flags(path, OpenFlags::SQLITE_OPEN_READ_ONLY) else {
113 return false;
114 };
115 conn.query_row(
118 "SELECT COUNT(*) FROM cities WHERE feature_code IS NOT NULL",
119 [],
120 |row| row.get::<_, i64>(0),
121 )
122 .is_ok_and(|count| count > 1_000)
123}
124
125pub fn build_geonames_db(project_root: &Path) -> Result<(), String> {
126 let data_dir = project_root.join("data");
127 let countries_path = data_dir.join("country_codes.txt");
128 let cities_path = data_dir.join("cities1000.txt");
129 let db_path = data_dir.join("geonames.db");
133
134 std::fs::create_dir_all(&data_dir).map_err(|e| format!("Failed to create db dir: {}", e))?;
135
136 let temp_path = data_dir.join(format!("geonames.db.tmp-{}", std::process::id()));
139 if temp_path.exists() {
140 std::fs::remove_file(&temp_path)
141 .map_err(|e| format!("Failed removing stale GeoNames temp DB: {}", e))?;
142 }
143
144 if let Err(error) = build_geonames_db_file(&countries_path, &cities_path, &temp_path) {
145 let _ = std::fs::remove_file(&temp_path);
146 return Err(error);
147 }
148 if !geonames_db_is_current(&temp_path) {
149 let _ = std::fs::remove_file(&temp_path);
150 return Err(
151 "Built GeoNames database failed validation; source data may be incomplete".to_string(),
152 );
153 }
154
155 let backup_path = data_dir.join(format!("geonames.db.backup-{}", std::process::id()));
156 if backup_path.exists() {
157 std::fs::remove_file(&backup_path)
158 .map_err(|e| format!("Failed removing stale GeoNames backup: {}", e))?;
159 }
160 let had_existing = db_path.exists();
161 if had_existing {
162 std::fs::rename(&db_path, &backup_path)
163 .map_err(|e| format!("Failed backing up old GeoNames DB: {}", e))?;
164 }
165 if let Err(error) = std::fs::rename(&temp_path, &db_path) {
166 if had_existing {
167 let _ = std::fs::rename(&backup_path, &db_path);
168 }
169 return Err(format!("Failed installing rebuilt GeoNames DB: {}", error));
170 }
171 if had_existing {
172 let _ = std::fs::remove_file(&backup_path);
173 }
174
175 Ok(())
176}
177
178fn build_geonames_db_file(
179 countries_path: &Path,
180 cities_path: &Path,
181 db_path: &Path,
182) -> Result<(), String> {
183 let mut conn = Connection::open(db_path).map_err(|e| format!("Failed to open DB: {}", e))?;
184
185 conn.execute_batch(
186 r#"
187 CREATE TABLE IF NOT EXISTS cities (
188 id INTEGER PRIMARY KEY,
189 name TEXT NOT NULL,
190 ascii_name TEXT NOT NULL,
191 latitude REAL NOT NULL,
192 longitude REAL NOT NULL,
193 country_code TEXT NOT NULL,
194 country_name TEXT NOT NULL,
195 population INTEGER,
196 feature_code TEXT,
197 timezone TEXT
198 );
199
200 CREATE TABLE IF NOT EXISTS countries (
201 code TEXT PRIMARY KEY,
202 name TEXT NOT NULL
203 );
204
205 CREATE INDEX IF NOT EXISTS idx_cities_coords ON cities(latitude, longitude);
206 CREATE INDEX IF NOT EXISTS idx_cities_feature ON cities(feature_code);
207 "#,
208 )
209 .map_err(|e| format!("Failed creating schema: {}", e))?;
210
211 let tx = conn
212 .transaction()
213 .map_err(|e| format!("Failed starting transaction: {}", e))?;
214
215 let countries = std::fs::read_to_string(countries_path)
216 .map_err(|e| format!("Failed reading {}: {}", countries_path.display(), e))?;
217 for line in countries.lines() {
218 let parts: Vec<&str> = line.split('\t').collect();
219 if parts.len() >= 2 {
220 let country_code = parts[0].trim_start_matches('\u{feff}');
224 if country_code.starts_with('#') {
225 continue;
226 }
227 let country_name = if parts.len() >= 5 { parts[4] } else { parts[1] };
230 tx.execute(
231 "INSERT OR IGNORE INTO countries (code, name) VALUES (?1, ?2)",
232 [country_code, country_name],
233 )
234 .map_err(|e| format!("Failed inserting country: {}", e))?;
235 }
236 }
237
238 let file = File::open(cities_path)
239 .map_err(|e| format!("Failed opening {}: {}", cities_path.display(), e))?;
240 let reader = BufReader::new(file);
241 let mut stmt = tx
242 .prepare(
243 r#"
244 INSERT OR IGNORE INTO cities
245 (id, name, ascii_name, latitude, longitude, country_code, country_name, population, feature_code, timezone)
246 VALUES
247 (?1, ?2, ?3, ?4, ?5, ?6, COALESCE((SELECT name FROM countries WHERE code = ?6), ?6), ?7, ?8, ?9)
248 "#,
249 )
250 .map_err(|e| format!("Failed preparing statement: {}", e))?;
251
252 for line in reader.lines() {
253 let line = line.map_err(|e| format!("Failed reading cities line: {}", e))?;
254 let parts: Vec<&str> = line.split('\t').collect();
255 if parts.len() >= 18 {
261 let id: i64 = parts[0].parse().unwrap_or(0);
262 let name = parts[1];
263 let ascii_name = parts[2];
264 let lat: f64 = parts[4].parse().unwrap_or(0.0);
265 let lon: f64 = parts[5].parse().unwrap_or(0.0);
266 let country_code = parts[8];
267 let population: i64 = parts[14].parse().unwrap_or(0);
268 let feature_code = parts[7];
269 let timezone = parts[17];
270
271 stmt.execute(rusqlite::params![
272 id,
273 name,
274 ascii_name,
275 lat,
276 lon,
277 country_code,
278 population,
279 feature_code,
280 timezone
281 ])
282 .map_err(|e| format!("Failed inserting city row: {}", e))?;
283 }
284 }
285
286 drop(stmt);
287 tx.commit()
288 .map_err(|e| format!("Failed committing GeoNames DB: {}", e))?;
289
290 Ok(())
291}
292
293#[cfg(test)]
294mod tests {
295 use super::*;
296
297 #[test]
298 fn builder_writes_beside_sources_and_accepts_raw_country_info() {
299 let temp = tempfile::tempdir().expect("tempdir");
300 let data = temp.path().join("data");
301 std::fs::create_dir_all(&data).expect("data dir");
302 std::fs::write(
303 data.join("country_codes.txt"),
304 "\u{feff}ZZ\tZZZ\t999\tZZ\tTest Country\n",
305 )
306 .expect("countries");
307
308 let mut cities = String::new();
309 for id in 1..=1_001 {
310 cities.push_str(&format!(
311 "{id}\tCity {id}\tCity {id}\t\t1.0\t2.0\tP\tPPL\tZZ\t\t\t\t\t\t{id}\t\t\tUTC\t2026-01-01\n"
312 ));
313 }
314 std::fs::write(data.join("cities1000.txt"), cities).expect("cities");
315
316 build_geonames_db(temp.path()).expect("build");
317
318 let db = data.join("geonames.db");
319 assert!(db.exists());
320 assert!(!temp.path().join("geonames.db").exists());
321 let conn = Connection::open(db).expect("open built db");
322 let country: String = conn
323 .query_row("SELECT country_name FROM cities WHERE id = 1", [], |row| {
324 row.get(0)
325 })
326 .expect("country name");
327 assert_eq!(country, "Test Country");
328 }
329}