1use crate::band::VrtBand;
4use crate::error::{Result, VrtError};
5use crate::source::PixelRect;
6use oxigdal_core::types::{GeoTransform, RasterDataType};
7use serde::{Deserialize, Serialize};
8use std::path::PathBuf;
9
10#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
12pub struct VrtDataset {
13 pub raster_x_size: u64,
15 pub raster_y_size: u64,
17 pub geo_transform: Option<GeoTransform>,
19 pub srs: Option<String>,
21 pub bands: Vec<VrtBand>,
23 pub block_size: Option<(u32, u32)>,
25 pub subclass: Option<VrtSubclass>,
27 pub vrt_path: Option<PathBuf>,
29}
30
31impl VrtDataset {
32 pub fn new(raster_x_size: u64, raster_y_size: u64) -> Self {
34 Self {
35 raster_x_size,
36 raster_y_size,
37 geo_transform: None,
38 srs: None,
39 bands: Vec::new(),
40 block_size: None,
41 subclass: None,
42 vrt_path: None,
43 }
44 }
45
46 pub fn from_bands(bands: Vec<VrtBand>) -> Result<Self> {
51 if bands.is_empty() {
52 return Err(VrtError::invalid_structure(
53 "Dataset must have at least one band",
54 ));
55 }
56
57 let (width, height) = Self::calculate_extent_from_band(&bands[0])?;
59
60 Ok(Self {
61 raster_x_size: width,
62 raster_y_size: height,
63 geo_transform: None,
64 srs: None,
65 bands,
66 block_size: None,
67 subclass: None,
68 vrt_path: None,
69 })
70 }
71
72 pub fn add_band(&mut self, band: VrtBand) {
74 self.bands.push(band);
75 }
76
77 pub fn with_geo_transform(mut self, geo_transform: GeoTransform) -> Self {
79 self.geo_transform = Some(geo_transform);
80 self
81 }
82
83 pub fn with_srs<S: Into<String>>(mut self, srs: S) -> Self {
85 self.srs = Some(srs.into());
86 self
87 }
88
89 pub fn with_block_size(mut self, width: u32, height: u32) -> Self {
91 self.block_size = Some((width, height));
92 self
93 }
94
95 pub fn with_subclass(mut self, subclass: VrtSubclass) -> Self {
97 self.subclass = Some(subclass);
98 self
99 }
100
101 pub fn with_vrt_path<P: Into<PathBuf>>(mut self, path: P) -> Self {
103 self.vrt_path = Some(path.into());
104 self
105 }
106
107 pub fn validate(&self) -> Result<()> {
112 if self.raster_x_size == 0 || self.raster_y_size == 0 {
113 return Err(VrtError::invalid_structure(
114 "Dataset dimensions must be > 0",
115 ));
116 }
117
118 if self.bands.is_empty() {
119 return Err(VrtError::invalid_structure(
120 "Dataset must have at least one band",
121 ));
122 }
123
124 for (idx, band) in self.bands.iter().enumerate() {
126 band.validate().map_err(|e| {
127 VrtError::invalid_structure(format!("Band {} validation failed: {}", idx + 1, e))
128 })?;
129
130 if band.band != idx + 1 {
132 return Err(VrtError::invalid_structure(format!(
133 "Band number mismatch: expected {}, got {}",
134 idx + 1,
135 band.band
136 )));
137 }
138 }
139
140 if let Some((width, height)) = self.block_size {
142 if width == 0 || height == 0 {
143 return Err(VrtError::invalid_structure("Block size must be > 0"));
144 }
145 }
146
147 Ok(())
148 }
149
150 pub fn band_count(&self) -> usize {
152 self.bands.len()
153 }
154
155 pub fn get_band(&self, index: usize) -> Option<&VrtBand> {
157 self.bands.get(index)
158 }
159
160 pub fn get_band_mut(&mut self, index: usize) -> Option<&mut VrtBand> {
162 self.bands.get_mut(index)
163 }
164
165 pub fn extent(&self) -> PixelRect {
167 PixelRect::new(0, 0, self.raster_x_size, self.raster_y_size)
168 }
169
170 pub fn effective_block_size(&self) -> (u32, u32) {
172 self.block_size.unwrap_or((256, 256))
173 }
174
175 fn calculate_extent_from_band(band: &VrtBand) -> Result<(u64, u64)> {
177 if band.sources.is_empty() {
178 return Err(VrtError::invalid_structure(
179 "Band has no sources to calculate extent from",
180 ));
181 }
182
183 let mut min_x = u64::MAX;
185 let mut min_y = u64::MAX;
186 let mut max_x = 0u64;
187 let mut max_y = 0u64;
188
189 for source in &band.sources {
190 if let Some(dst_rect) = source.dst_rect() {
191 min_x = min_x.min(dst_rect.x_off);
192 min_y = min_y.min(dst_rect.y_off);
193 max_x = max_x.max(dst_rect.x_off + dst_rect.x_size);
194 max_y = max_y.max(dst_rect.y_off + dst_rect.y_size);
195 } else if let Some(ref props) = source.properties {
196 max_x = max_x.max(props.width);
198 max_y = max_y.max(props.height);
199 }
200 }
201
202 if max_x == 0 || max_y == 0 {
203 return Err(VrtError::invalid_structure(
204 "Cannot calculate extent: no valid source windows",
205 ));
206 }
207
208 Ok((max_x - min_x, max_y - min_y))
209 }
210
211 pub fn merge_geo_transforms(&mut self) -> Result<()> {
216 if self.geo_transform.is_some() {
217 return Ok(()); }
219
220 let mut transforms = Vec::new();
221
222 for band in &self.bands {
224 for source in &band.sources {
225 if let Some(ref props) = source.properties {
226 if let Some(ref gt) = props.geo_transform {
227 transforms.push(*gt);
228 }
229 }
230 }
231 }
232
233 if transforms.is_empty() {
234 return Ok(()); }
236
237 self.geo_transform = Some(transforms[0]);
240
241 Ok(())
242 }
243
244 pub fn primary_data_type(&self) -> Option<RasterDataType> {
246 self.bands.first().map(|b| b.data_type)
247 }
248
249 pub fn has_uniform_data_type(&self) -> bool {
251 if self.bands.is_empty() {
252 return true;
253 }
254
255 let first_type = self.bands[0].data_type;
256 self.bands.iter().all(|b| b.data_type == first_type)
257 }
258}
259
260#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
262pub enum VrtSubclass {
263 #[default]
265 Standard,
266 Warped,
268 Pansharpened,
270 Processed,
272}
273
274#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
276pub struct VrtMetadata {
277 pub domain: Option<String>,
279 pub items: Vec<(String, String)>,
281}
282
283impl VrtMetadata {
284 pub fn new() -> Self {
286 Self {
287 domain: None,
288 items: Vec::new(),
289 }
290 }
291
292 pub fn with_domain<S: Into<String>>(domain: S) -> Self {
294 Self {
295 domain: Some(domain.into()),
296 items: Vec::new(),
297 }
298 }
299
300 pub fn add_item<K: Into<String>, V: Into<String>>(&mut self, key: K, value: V) {
302 self.items.push((key.into(), value.into()));
303 }
304
305 pub fn get(&self, key: &str) -> Option<&str> {
307 self.items
308 .iter()
309 .find(|(k, _)| k == key)
310 .map(|(_, v)| v.as_str())
311 }
312}
313
314impl Default for VrtMetadata {
315 fn default() -> Self {
316 Self::new()
317 }
318}
319
320#[cfg(test)]
321mod tests {
322 use super::*;
323 use crate::band::VrtBand;
324 use crate::source::{SourceFilename, VrtSource};
325 use oxigdal_core::types::RasterDataType;
326
327 #[test]
328 fn test_vrt_dataset_creation() {
329 let dataset = VrtDataset::new(512, 512);
330 assert_eq!(dataset.raster_x_size, 512);
331 assert_eq!(dataset.raster_y_size, 512);
332 assert_eq!(dataset.band_count(), 0);
333 }
334
335 #[test]
336 fn test_vrt_dataset_validation() {
337 let mut dataset = VrtDataset::new(512, 512);
338 let source = VrtSource::new(SourceFilename::absolute("/test.tif"), 1);
339 let band = VrtBand::simple(1, RasterDataType::UInt8, source);
340 dataset.add_band(band);
341
342 assert!(dataset.validate().is_ok());
343
344 let empty_dataset = VrtDataset::new(512, 512);
345 assert!(empty_dataset.validate().is_err());
346
347 let invalid_dataset = VrtDataset::new(0, 0);
348 assert!(invalid_dataset.validate().is_err());
349 }
350
351 #[test]
352 fn test_vrt_dataset_extent() {
353 let dataset = VrtDataset::new(1024, 768);
354 let extent = dataset.extent();
355 assert_eq!(extent.x_size, 1024);
356 assert_eq!(extent.y_size, 768);
357 }
358
359 #[test]
360 fn test_vrt_dataset_band_access() {
361 let mut dataset = VrtDataset::new(512, 512);
362 let source = VrtSource::new(SourceFilename::absolute("/test.tif"), 1);
363 let band = VrtBand::simple(1, RasterDataType::UInt8, source);
364 dataset.add_band(band);
365
366 assert_eq!(dataset.band_count(), 1);
367 assert!(dataset.get_band(0).is_some());
368 assert!(dataset.get_band(1).is_none());
369 }
370
371 #[test]
372 fn test_effective_block_size() {
373 let dataset = VrtDataset::new(512, 512);
374 assert_eq!(dataset.effective_block_size(), (256, 256));
375
376 let dataset_with_blocks = VrtDataset::new(512, 512).with_block_size(128, 128);
377 assert_eq!(dataset_with_blocks.effective_block_size(), (128, 128));
378 }
379
380 #[test]
381 fn test_vrt_metadata() {
382 let mut metadata = VrtMetadata::new();
383 metadata.add_item("author", "test");
384 metadata.add_item("version", "1.0");
385
386 assert_eq!(metadata.get("author"), Some("test"));
387 assert_eq!(metadata.get("version"), Some("1.0"));
388 assert_eq!(metadata.get("missing"), None);
389 }
390
391 #[test]
392 fn test_uniform_data_type() {
393 let mut dataset = VrtDataset::new(512, 512);
394
395 let source1 = VrtSource::new(SourceFilename::absolute("/test1.tif"), 1);
396 let band1 = VrtBand::simple(1, RasterDataType::UInt8, source1);
397 dataset.add_band(band1);
398
399 let source2 = VrtSource::new(SourceFilename::absolute("/test2.tif"), 1);
400 let band2 = VrtBand::simple(2, RasterDataType::UInt8, source2);
401 dataset.add_band(band2);
402
403 assert!(dataset.has_uniform_data_type());
404
405 let source3 = VrtSource::new(SourceFilename::absolute("/test3.tif"), 1);
406 let band3 = VrtBand::simple(3, RasterDataType::Float32, source3);
407 dataset.add_band(band3);
408
409 assert!(!dataset.has_uniform_data_type());
410 }
411}