eorst 1.0.1

Earth Observation and Remote Sensing Toolkit - library for raster processing pipelines
//! STAC helpers for querying and processing STAC ItemCollections.
//!
//! This module provides utility functions for working with STAC (SpatioTemporal Asset Catalog)
//! data including extracting datetimes, asset names, and sourcing files.

use chrono::{DateTime, FixedOffset};
use gdal::vector::Geometry;
use itertools::Itertools;
use log::debug;
use std::path::PathBuf;
use stac::ItemCollection;

/// Returns unique datetimes from a vector, merging those within 6 hours of each other.
pub fn unique_datetimes_in_range(dates: Vec<DateTime<FixedOffset>>) -> Vec<DateTime<FixedOffset>> {
    let six_hours = chrono::Duration::hours(6);
    let mut result: Vec<DateTime<FixedOffset>> = vec![];
    let mut last_date = dates[0];

    for &date in dates[1..].iter() {
        if date - last_date <= six_hours {
        } else {
            result.push(last_date);
            last_date = date;
        }
    }
    result.push(last_date); // push the last date

    result
}

/// Extracts sorted datetimes from a STAC ItemCollection.
pub fn get_sorted_datetimes(feature_collection: &ItemCollection) -> Vec<DateTime<FixedOffset>> {
    let mut dates: Vec<DateTime<FixedOffset>> = feature_collection
        .items
        .iter()
        .map(|i| {
            DateTime::parse_from_rfc3339(&i.properties.datetime.to_owned().unwrap().to_rfc3339())
                .unwrap()
        })
        .collect();
    dates.sort();

    dates
}

/// Extracts asset names from a STAC ItemCollection.
pub fn get_asset_names(feature_collection: &ItemCollection) -> Vec<String> {
    let mut names = Vec::new();

    for item in &feature_collection.items {
        for asset in item.assets.values() {
            if let Some(eo_bands) = asset.additional_fields.get("eo:bands") {
                // normal EO band logic
                let eo_band = &eo_bands[0];
                let mut name = eo_band["common_name"]
                    .as_str()
                    .unwrap_or_else(|| eo_band["name"].as_str().unwrap())
                    .to_owned();
                if name == "None" {
                    name = eo_band["name"]
                        .as_str()
                        .unwrap_or("unknown")
                        .to_owned();
                }
                names.push(name);
            } else {
                // fallback: just use the asset key
                names.push(asset.title.clone().unwrap_or_else(|| "unknown".to_string()));
            }
        }
    }

    let mut names = names.into_iter().unique().collect::<Vec<_>>();
    names.sort();
    names
}

fn get_name_from_bands(bands: Option<&serde_json::value::Value>) -> Option<String> {
    if let Some(json_array) = bands {
        if let Some(json_object) = json_array.get(0) {
            // Some catalogs will store the layer name under common_name
            if let Some(name_field) = json_object.get("common_name") {
                if let Some(name) = name_field.as_str() {
                    return Some(name.to_owned());
                }
            }
            // and some others, simply under name; So if the above condition fails (no common_name) it falls back
            // to the name.
            if let Some(name_field) = json_object.get("name") {
                if let Some(name) = name_field.as_str() {
                    return Some(name.to_owned());
                }
            }
        }
    }
    None
}

fn canonical_asset_name(asset: &stac::Asset) -> String {
    // First, try eo:bands
    if let Some(eo_bands) = asset.additional_fields.get("eo:bands") {
        let eo_band = &eo_bands[0];
        let mut name = eo_band["common_name"]
            .as_str()
            .or_else(|| eo_band["name"].as_str())
            .unwrap_or("unknown")
            .to_string();
        if name == "None" {
            name = eo_band["name"].as_str().unwrap_or("unknown").to_string();
        }
        return name.to_lowercase();
    }

    // Fallback: check title
    if let Some(title) = &asset.title {
        // map known QA bands to standard names
        debug!("title {:?}", title);
        match title.as_str() {
            "Surface Temperature Band" => "surface temperature band".to_string(),
            "Pixel Quality Assessment Band" => "pixel quality assessment band".to_string(),
            "RADSAT QA Band" => "qa_radsat".to_string(),
            _ => title.to_lowercase().replace(' ', "_"),
        }
    } else {
        // fallback to some generic string
        "unknown".to_string()
    }
}

/// Gets file paths for a specific asset from STAC items.
pub fn get_sources_for_asset(items: &Vec<stac::Item>, asset_name: &str) -> Vec<PathBuf> {
    let mut sources = Vec::new();

    for item in items {
        for asset in item.assets.values() {
            let name = canonical_asset_name(asset);

            if name == asset_name.to_lowercase() {
                sources.push(PathBuf::from(&asset.href));
            }
        }
    }
    sources
}

/// Gets the asset href for a specific date and asset name from a STAC feature collection.
pub fn get_asset_href(
    feature_collection: &ItemCollection,
    date_time: &DateTime<FixedOffset>,
    asset_name: &str,
) -> PathBuf {
    let mut found_asset = stac::Asset::new("test");
    for item in &feature_collection.items {
        let date = DateTime::parse_from_rfc3339(
            &item.properties.datetime.to_owned().unwrap().to_rfc3339(),
        )
        .unwrap();

        for asset in item.assets.values() {
            let asset = asset.clone();
            let bands = asset.additional_fields.get("eo:bands");
            let mut names = Vec::new();
            if let Some(name) = get_name_from_bands(bands) {
                names.push(name);
            } else {
                panic!("Name not found.");
            }

            if (names[0] == asset_name) && (date == *date_time) {
                found_asset = asset.clone();
            }
        }
    }
    PathBuf::from(found_asset.href)
}

/// Gets STAC items matching a specific date (within 24 hours).
pub fn get_items_for_date(
    feature_collection: &ItemCollection,
    date_time: &DateTime<FixedOffset>,
) -> Vec<stac::Item> {
    let mut found_items = Vec::new();

    for item in &feature_collection.items {
        let date = DateTime::parse_from_rfc3339(
            &item.properties.datetime.to_owned().unwrap().to_rfc3339(),
        )
        .unwrap();

        let delta = *date_time - date;

        if delta.abs() <= chrono::Duration::hours(24) {
            found_items.push(item.to_owned());
        }
    }
    found_items
}

/// Swaps x and y coordinates in a geometry.
pub fn swap_coordinates(gdal_geometry: &Geometry) -> Geometry {
    let wkt = gdal_geometry.wkt().unwrap();

    let mut swapped_wkt = String::new();
    swapped_wkt.push_str("POLYGON ((");
    let tokens: Vec<&str> = wkt.split([',', '(', ')']).collect();

    // Split the WKT string by spaces and parenthesesi
    for token in tokens {
        if !token.starts_with("POLY") {
            // Split each token by comma to get the coordinate values
            let coordinates: Vec<&str> = token.split(' ').collect();
            if coordinates.len() == 2 {
                // Swap the X and Y coordinates
                let x = coordinates[0].trim();
                let y = coordinates[1].trim();

                // Append the swapped coordinates to the new WKT string
                let to_append = &format!("{} {} {}, ", y, x, 0);

                swapped_wkt.push_str(to_append);
            }
        }

        //swapped_wkt.push(','); // Add space between coordinates
    }
    swapped_wkt.push_str("))");
    let mut result = swapped_wkt;

    // Find the last occurrence of a comma
    if let Some(last_comma_index) = result.rfind(',') {
        // Remove the comma from the string
        result.remove(last_comma_index);
    }

    Geometry::from_wkt(&result).unwrap()
}