use std::collections::{BTreeMap, BTreeSet};
use std::fs::File;
use std::io::{BufWriter, Write};
use std::path::Path;
use fst::{MapBuilder, SetBuilder};
use rkyv::rancor::Error as RkyvError;
use crate::Error;
pub use crate::data::{EngineRow, EuModelRow, LookupRow, MakeRow, ModelRow, SchemaRow, VinRuleRow};
use crate::data::{RkyvSer, Saveable};
pub fn write_grouped<T>(
sorted_pairs: &[(String, Vec<T>)],
fst_path: &Path,
values_path: &Path,
) -> crate::Result<()>
where
T: RkyvSer,
{
let fst_writer = BufWriter::new(File::create(fst_path)?);
let mut values_writer = BufWriter::new(File::create(values_path)?);
let mut builder = MapBuilder::new(fst_writer).map_err(|e| Error::MissingData(e.to_string()))?;
let mut offset: u64 = 0;
for (key, values) in sorted_pairs {
let bytes =
rkyv::to_bytes::<RkyvError>(values).map_err(|e| Error::MissingData(e.to_string()))?;
values_writer.write_all(&bytes)?;
let combined = (offset << 32) | (bytes.len() as u64);
builder
.insert(key, combined)
.map_err(|e| Error::MissingData(e.to_string()))?;
offset += bytes.len() as u64;
}
builder
.finish()
.map_err(|e| Error::MissingData(e.to_string()))?;
values_writer.flush()?;
Ok(())
}
pub fn write_set(sorted_keys: &[String], fst_path: &Path) -> crate::Result<()> {
let writer = BufWriter::new(File::create(fst_path)?);
let mut builder = SetBuilder::new(writer).map_err(|e| Error::MissingData(e.to_string()))?;
for key in sorted_keys {
builder
.insert(key)
.map_err(|e| Error::MissingData(e.to_string()))?;
}
builder
.finish()
.map_err(|e| Error::MissingData(e.to_string()))?;
Ok(())
}
pub fn derive_make_models(
wmi_make: &[(String, Vec<MakeRow>)],
wmi_schema: &[(String, Vec<SchemaRow>)],
schema_lookup: &[(String, Vec<LookupRow>)],
) -> Vec<(String, Vec<ModelRow>)> {
let make_by_wmi: BTreeMap<&str, &str> = wmi_make
.iter()
.filter_map(|(wmi, rows)| rows.first().map(|r| (wmi.as_str(), r.name.as_str())))
.collect();
let schemas_by_wmi: BTreeMap<&str, Vec<&str>> = wmi_schema
.iter()
.map(|(wmi, rows)| {
(
wmi.as_str(),
rows.iter().map(|r| r.id.as_str()).collect::<Vec<_>>(),
)
})
.collect();
let lookups_by_schema: BTreeMap<&str, &[LookupRow]> = schema_lookup
.iter()
.map(|(id, rows)| (id.as_str(), rows.as_slice()))
.collect();
let mut acc: BTreeMap<String, BTreeSet<String>> = BTreeMap::new();
for (wmi, make) in &make_by_wmi {
let Some(schemas) = schemas_by_wmi.get(wmi) else {
continue;
};
for sid in schemas {
let Some(lookups) = lookups_by_schema.get(sid) else {
continue;
};
for lk in lookups.iter() {
if lk.element == "Model" {
acc.entry(make.to_ascii_uppercase())
.or_default()
.insert(lk.value.clone());
}
}
}
}
acc.into_iter()
.map(|(make, models)| {
(
make,
models.into_iter().map(|name| ModelRow { name }).collect(),
)
})
.collect()
}
pub fn collect_makes(wmi_make: &[(String, Vec<MakeRow>)]) -> Vec<String> {
let mut set: BTreeSet<String> = BTreeSet::new();
for (_, rows) in wmi_make {
for r in rows {
set.insert(r.name.to_ascii_uppercase());
}
}
set.into_iter().collect()
}
pub const EXTRA_MAKES: &[&str] = &[
"DENZA", "JAC", "LEAPMOTOR", "AION", "AIWAYS", "ZEEKR", "OMODA", "JAECOO", "SERES", "VOYAH", "HONGQI", "MAXUS", "JETOUR", "AVATR", "BAIC", "GWM", "TANK", "ORA", "FORTHING", "BESTUNE", "SKYWELL", "DFSK", "XIAOMI", "IM", "LYNK & CO", "KGM", "AIXAM", "LIGIER", "MICROLINO", "TAZZARI", "XEV", "ESTRIMA", "CHATENET", "DR", "EVO", "SWM", ];
pub fn merge_makes(wmi_makes: &[String]) -> Vec<String> {
let mut set: BTreeSet<String> = wmi_makes.iter().map(|m| m.to_ascii_uppercase()).collect();
for extra in EXTRA_MAKES {
set.insert(extra.to_ascii_uppercase());
}
set.into_iter().collect()
}
pub fn build_all(
wmi_make: &[(String, Vec<MakeRow>)],
wmi_schema: &[(String, Vec<SchemaRow>)],
schema_lookup: &[(String, Vec<LookupRow>)],
out_dir: &Path,
) -> crate::Result<()> {
std::fs::create_dir_all(out_dir)?;
write_grouped(
wmi_make,
&out_dir.join(format!("{}.fst", MakeRow::base_name())),
&out_dir.join(format!("{}.bin", MakeRow::base_name())),
)?;
write_grouped(
wmi_schema,
&out_dir.join(format!("{}.fst", SchemaRow::base_name())),
&out_dir.join(format!("{}.bin", SchemaRow::base_name())),
)?;
write_grouped(
schema_lookup,
&out_dir.join(format!("{}.fst", LookupRow::base_name())),
&out_dir.join(format!("{}.bin", LookupRow::base_name())),
)?;
let makes = collect_makes(wmi_make);
write_set(&makes, &out_dir.join("makes.fst"))?;
let make_models = derive_make_models(wmi_make, wmi_schema, schema_lookup);
write_grouped(
&make_models,
&out_dir.join(format!("{}.fst", ModelRow::base_name())),
&out_dir.join(format!("{}.bin", ModelRow::base_name())),
)?;
Ok(())
}
pub fn write_csv_grouped<T, F>(
csv_path: &Path,
fst_path: &Path,
values_path: &Path,
mut row_fn: F,
) -> crate::Result<()>
where
T: RkyvSer,
F: FnMut(&csv::StringRecord) -> Option<T>,
{
let mut reader = csv::ReaderBuilder::new()
.has_headers(true)
.delimiter(b'\t')
.quoting(false)
.from_path(csv_path)?;
let fst_writer = BufWriter::new(File::create(fst_path)?);
let mut values_writer = BufWriter::new(File::create(values_path)?);
let mut builder = MapBuilder::new(fst_writer).map_err(|e| Error::MissingData(e.to_string()))?;
let mut offset: u64 = 0;
let mut current_key: Option<String> = None;
let mut group: Vec<T> = Vec::new();
let flush = |builder: &mut MapBuilder<BufWriter<File>>,
values_writer: &mut BufWriter<File>,
offset: &mut u64,
key: &str,
group: &Vec<T>|
-> crate::Result<()> {
let bytes =
rkyv::to_bytes::<RkyvError>(group).map_err(|e| Error::MissingData(e.to_string()))?;
values_writer.write_all(&bytes)?;
let combined = (*offset << 32) | (bytes.len() as u64);
builder
.insert(key, combined)
.map_err(|e| Error::MissingData(format!("FST insert failed at key {key:?}: {e}")))?;
*offset += bytes.len() as u64;
Ok(())
};
for rec in reader.records() {
let rec = rec.map_err(|e| Error::MissingData(e.to_string()))?;
let Some(key) = rec.get(0) else { continue };
let Some(item) = row_fn(&rec) else { continue };
match ¤t_key {
Some(k) if k == key => group.push(item),
Some(k) => {
flush(&mut builder, &mut values_writer, &mut offset, k, &group)?;
group.clear();
group.push(item);
current_key = Some(key.to_string());
}
None => {
current_key = Some(key.to_string());
group.push(item);
}
}
}
if let Some(k) = current_key {
flush(&mut builder, &mut values_writer, &mut offset, &k, &group)?;
}
builder
.finish()
.map_err(|e| Error::MissingData(e.to_string()))?;
values_writer.flush()?;
Ok(())
}
pub fn build_from_csv(csv_dir: &Path, out_dir: &Path) -> crate::Result<()> {
std::fs::create_dir_all(out_dir)?;
write_csv_grouped::<MakeRow, _>(
&csv_dir.join("wmi_make.csv"),
&out_dir.join(format!("{}.fst", MakeRow::base_name())),
&out_dir.join(format!("{}.bin", MakeRow::base_name())),
|rec| {
rec.get(1).map(|s| MakeRow {
name: s.to_string(),
country: rec.get(2).unwrap_or("").to_string(),
region: rec.get(3).unwrap_or("").to_string(),
})
},
)?;
write_csv_grouped::<SchemaRow, _>(
&csv_dir.join("wmi_schema_id.csv"),
&out_dir.join(format!("{}.fst", SchemaRow::base_name())),
&out_dir.join(format!("{}.bin", SchemaRow::base_name())),
|rec| rec.get(1).map(|s| SchemaRow { id: s.to_string() }),
)?;
write_csv_grouped::<LookupRow, _>(
&csv_dir.join("schema_id_lookup.csv"),
&out_dir.join(format!("{}.fst", LookupRow::base_name())),
&out_dir.join(format!("{}.bin", LookupRow::base_name())),
|rec| {
let pattern = rec.get(1)?.to_string();
let element = rec.get(2)?.to_string();
let value = rec.get(3)?.to_string();
let weight = rec.get(4).and_then(|s| s.parse().ok()).unwrap_or(0);
Some(LookupRow {
pattern,
element,
value,
weight,
})
},
)?;
let wmi_make = read_csv_grouped(&csv_dir.join("wmi_make.csv"), |rec| {
rec.get(1).map(|s| MakeRow {
name: s.to_string(),
country: rec.get(2).unwrap_or("").to_string(),
region: rec.get(3).unwrap_or("").to_string(),
})
})?;
let wmi_schema = read_csv_grouped(&csv_dir.join("wmi_schema_id.csv"), |rec| {
rec.get(1).map(|s| SchemaRow { id: s.to_string() })
})?;
let schema_lookup = read_csv_grouped(&csv_dir.join("schema_id_lookup.csv"), |rec| {
let pattern = rec.get(1)?.to_string();
let element = rec.get(2)?.to_string();
let value = rec.get(3)?.to_string();
let weight = rec.get(4).and_then(|s| s.parse().ok()).unwrap_or(0);
Some(LookupRow {
pattern,
element,
value,
weight,
})
})?;
let makes = merge_makes(&collect_makes(&wmi_make));
write_set(&makes, &out_dir.join("makes.fst"))?;
let make_models = derive_make_models(&wmi_make, &wmi_schema, &schema_lookup);
write_grouped(
&make_models,
&out_dir.join(format!("{}.fst", ModelRow::base_name())),
&out_dir.join(format!("{}.bin", ModelRow::base_name())),
)?;
Ok(())
}
fn read_csv_grouped<T, F>(csv_path: &Path, mut row_fn: F) -> crate::Result<Vec<(String, Vec<T>)>>
where
F: FnMut(&csv::StringRecord) -> Option<T>,
{
let mut reader = csv::ReaderBuilder::new()
.has_headers(true)
.delimiter(b'\t')
.quoting(false)
.from_path(csv_path)?;
let mut acc: BTreeMap<String, Vec<T>> = BTreeMap::new();
for rec in reader.records() {
let rec = rec.map_err(|e| Error::MissingData(e.to_string()))?;
let Some(key) = rec.get(0) else { continue };
let Some(item) = row_fn(&rec) else { continue };
acc.entry(key.to_string()).or_default().push(item);
}
Ok(acc.into_iter().collect())
}
impl From<csv::Error> for Error {
fn from(e: csv::Error) -> Self {
Error::MissingData(e.to_string())
}
}
pub fn build_from_rip(rip_dir: &Path, out_dir: &Path) -> crate::Result<()> {
std::fs::create_dir_all(out_dir)?;
write_csv_grouped::<MakeRow, _>(
&rip_dir.join("wmi_merged.tsv"),
&out_dir.join(format!("{}.fst", MakeRow::base_name())),
&out_dir.join(format!("{}.bin", MakeRow::base_name())),
|rec| {
rec.get(1).map(|s| MakeRow {
name: s.to_string(),
country: rec.get(2).unwrap_or("").to_string(),
region: rec.get(3).unwrap_or("").to_string(),
})
},
)?;
let mut brands: Vec<String> = Vec::new();
let mut brand_reader = csv::ReaderBuilder::new()
.has_headers(true)
.delimiter(b'\t')
.quoting(false)
.from_path(rip_dir.join("brands.tsv"))?;
for rec in brand_reader.records() {
let rec = rec.map_err(|e| Error::MissingData(e.to_string()))?;
if let Some(b) = rec.get(0) {
brands.push(b.to_string());
}
}
let brands = merge_makes(&brands);
write_set(&brands, &out_dir.join("makes.fst"))?;
write_csv_grouped::<EuModelRow, _>(
&rip_dir.join("brand_models.tsv"),
&out_dir.join(format!("{}.fst", EuModelRow::base_name())),
&out_dir.join(format!("{}.bin", EuModelRow::base_name())),
|rec| {
let name = rec.get(1)?.to_string();
let first_year: u16 = rec.get(2).and_then(|s| s.parse().ok()).unwrap_or(0);
let last_year: u16 = rec.get(3).and_then(|s| s.parse().ok()).unwrap_or(0);
Some(EuModelRow {
name,
first_year,
last_year,
})
},
)?;
write_csv_grouped::<EngineRow, _>(
&rip_dir.join("engines.tsv"),
&out_dir.join(format!("{}.fst", EngineRow::base_name())),
&out_dir.join(format!("{}.bin", EngineRow::base_name())),
engine_row_from_record,
)?;
let wmi_rules_path = rip_dir.join("wmi_rules.tsv");
if wmi_rules_path.exists() {
let rule_rows = read_csv_grouped(&wmi_rules_path, |rec| {
let remainder = rec.get(1).unwrap_or("").to_string();
let make = rec.get(2).unwrap_or("").to_string();
let model = rec.get(3).unwrap_or("").to_string();
Some(VinRuleRow {
remainder,
make,
model,
})
})?;
let mut sorted: Vec<(String, Vec<VinRuleRow>)> = rule_rows
.into_iter()
.map(|(wmi, mut rows)| {
rows.sort_by_key(|r| std::cmp::Reverse(r.remainder.len()));
(wmi, rows)
})
.collect();
sorted.sort_by(|a, b| a.0.cmp(&b.0));
write_grouped(
&sorted,
&out_dir.join(format!("{}.fst", VinRuleRow::base_name())),
&out_dir.join(format!("{}.bin", VinRuleRow::base_name())),
)?;
}
Ok(())
}
fn engine_row_from_record(rec: &csv::StringRecord) -> Option<EngineRow> {
let model = rec.get(1)?.to_string();
let year: u16 = rec.get(2).and_then(|s| s.parse().ok()).unwrap_or(0);
let name = rec.get(3).unwrap_or("").to_string();
let cylinders = rec.get(4).unwrap_or("").to_string();
let displacement_cm3: u32 = rec
.get(5)
.and_then(|s| s.parse::<f64>().ok())
.map(|v| v.round() as u32)
.unwrap_or(0);
let power_kw: u32 = rec
.get(6)
.and_then(|s| s.parse::<f64>().ok())
.map(|v| v.round() as u32)
.unwrap_or(0);
let power_hp: u32 = rec
.get(7)
.and_then(|s| s.parse::<f64>().ok())
.map(|v| v.round() as u32)
.unwrap_or(0);
let torque_nm: u32 = rec
.get(8)
.and_then(|s| s.parse::<f64>().ok())
.map(|v| v.round() as u32)
.unwrap_or(0);
let fuel = rec.get(9).unwrap_or("").to_string();
let drive = rec.get(11).unwrap_or("").to_string();
let gearbox = rec.get(12).unwrap_or("").to_string();
Some(EngineRow {
model,
year,
name,
cylinders,
displacement_cm3,
power_kw,
power_hp,
torque_nm,
fuel,
drive,
gearbox,
})
}