v2rmp 0.4.9

use crate::core::overture::{BBox, Geometry, OvertureExtractor, OvertureSegment};
use anyhow::{Context, Result};
use geo_traits::{
    CoordTrait, GeometryTrait, GeometryType, LineStringTrait, MultiLineStringTrait, PointTrait,
};
use geojson::{Feature, FeatureCollection, Geometry as GeoJsonGeometry, Value as GeoJsonValue};
use serde::{Deserialize, Serialize};
use serde_json::json;
use sqlx::Row;
use std::collections::HashMap;
use std::fs::File;
use std::io::Write;
use wkb::reader::read_wkb;

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum ExtractSource {
    Osm,
    Overture,
    Postgres,
    R2,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BBoxRequest {
    pub min_lon: f64,
    pub min_lat: f64,
    pub max_lon: f64,
    pub max_lat: f64,
}

impl From<BBoxRequest> for BBox {
    fn from(r: BBoxRequest) -> Self {
        BBox {
            min_lon: r.min_lon,
            min_lat: r.min_lat,
            max_lon: r.max_lon,
            max_lat: r.max_lat,
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExtractRequest {
    pub source: ExtractSource,
    pub bbox: BBoxRequest,
    pub road_classes: Vec<RoadClass>,
    pub output_path: String,
    pub database_url: Option<String>,
    pub table_name: Option<String>,
    pub r2_bucket: Option<String>,
    pub r2_access_key_id: Option<String>,
    pub r2_secret_access_key: Option<String>,
    pub r2_endpoint: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
pub enum RoadClass {
    Residential,
    Tertiary,
    Secondary,
    Primary,
    Trunk,
    Motorway,
    Unclassified,
    LivingStreet,
    Service,
    SecondaryLink,
    PrimaryLink,
    TrunkLink,
    MotorwayLink,
}

impl RoadClass {
    pub fn all_vehicle() -> Vec<Self> {
        vec![
            Self::Residential,
            Self::Tertiary,
            Self::Secondary,
            Self::Primary,
            Self::Trunk,
            Self::Motorway,
            Self::Unclassified,
            Self::LivingStreet,
            Self::Service,
            Self::SecondaryLink,
            Self::PrimaryLink,
            Self::TrunkLink,
            Self::MotorwayLink,
        ]
    }

    pub fn as_str(&self) -> &'static str {
        match self {
            Self::Residential => "residential",
            Self::Tertiary => "tertiary",
            Self::Secondary => "secondary",
            Self::Primary => "primary",
            Self::Trunk => "trunk",
            Self::Motorway => "motorway",
            Self::Unclassified => "unclassified",
            Self::LivingStreet => "living_street",
            Self::Service => "service",
            Self::SecondaryLink => "secondary_link",
            Self::PrimaryLink => "primary_link",
            Self::TrunkLink => "trunk_link",
            Self::MotorwayLink => "motorway_link",
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExtractResult {
    pub nodes: usize,
    pub edges: usize,
    pub total_km: f64,
    pub output_path: String,
}

/// Extract road network from Overture Maps S3 Parquet files.
///
/// This is production implementation for Overture extraction.
/// For OSM PBF, this currently returns a stub result.
pub fn run_extract(req: &ExtractRequest) -> anyhow::Result<ExtractResult> {
    match req.source {
        ExtractSource::Overture => run_overture_extract(req),
        ExtractSource::Osm => {
            // TODO: Implement OSM PBF extraction using osmpbf crate
            anyhow::bail!("OSM PBF extraction not yet implemented")
        }
        ExtractSource::Postgres => run_postgres_extract(req),
        ExtractSource::R2 => run_r2_extract(req),
    }
}

/// Run extraction from Cloudflare R2
fn run_r2_extract(req: &ExtractRequest) -> anyhow::Result<ExtractResult> {
    let runtime = tokio::runtime::Runtime::new()
        .context("Failed to create Tokio runtime for R2 operations")?;

    runtime.block_on(async {
        let bucket = req.r2_bucket.as_deref().unwrap_or("v2rmp");
        
        let _storage = if let (Some(akid), Some(sak), Some(end)) = (
            &req.r2_access_key_id,
            &req.r2_secret_access_key,
            &req.r2_endpoint,
        ) {
            crate::core::r2::R2Storage::new(bucket, akid, sak, end)?
        } else {
            crate::core::r2::R2Storage::from_env(bucket)?
        };

        tracing::info!("Connected to R2 bucket: {}", bucket);
        
        // For now, R2 extraction is a stub that just confirms connection
        // In a real scenario, you'd list files, download, and process them
        // similar to Overture extraction.
        
        Ok(ExtractResult {
            nodes: 0,
            edges: 0,
            total_km: 0.0,
            output_path: req.output_path.clone(),
        })
    })
}

/// Run extraction from PostgreSQL/PostGIS
fn run_postgres_extract(req: &ExtractRequest) -> anyhow::Result<ExtractResult> {
    let runtime = tokio::runtime::Runtime::new()
        .context("Failed to create Tokio runtime for Postgres operations")?;

    runtime.block_on(async {
        let db_url = req
            .database_url
            .as_ref()
            .context("Postgres database URL is required for Postgres source")?;
        let table = req.table_name.as_deref().unwrap_or("roads");

        tracing::info!("Connecting to PostgreSQL: {}", db_url);
        let pool = sqlx::PgPool::connect(db_url)
            .await
            .context("Failed to connect to PostgreSQL")?;

        tracing::info!(
            "Querying table '{}' for bbox: [{:.4}, {:.4}, {:.4}, {:.4}]",
            table,
            req.bbox.min_lon,
            req.bbox.min_lat,
            req.bbox.max_lon,
            req.bbox.max_lat
        );

        // This query assumes a typical PostGIS schema
        // We look for common column names: geom/geometry, name, class/highway, oneway
        let query = format!(
            "SELECT 
                COALESCE(id::text, gen_random_uuid()::text) as id,
                name,
                COALESCE(class, highway, 'unclassified') as class,
                oneway,
                ST_AsBinary(geom) as geom_wkb
             FROM {}
             WHERE geom && ST_MakeEnvelope($1, $2, $3, $4, 4326)",
            table
        );

        let rows = sqlx::query(&query)
            .bind(req.bbox.min_lon)
            .bind(req.bbox.min_lat)
            .bind(req.bbox.max_lon)
            .bind(req.bbox.max_lat)
            .fetch_all(&pool)
            .await
            .context("Failed to execute Postgres query")?;

        tracing::info!("Extracted {} rows from Postgres", rows.len());

        let mut features = Vec::new();

        for row in rows {
            let id: String = row.get("id");
            let name: Option<String> = row.get("name");
            let class: Option<String> = row.get("class");
            let oneway: Option<String> = row.get("oneway");
            let wkb_data: Vec<u8> = row.get("geom_wkb");

            // Parse WKB
            let geometry = match read_wkb(&wkb_data) {
                Ok(geom) => geom,
                Err(e) => {
                    tracing::warn!("Failed to parse WKB for row {}: {}", id, e);
                    continue;
                }
            };

            // Convert to GeoJSON geometry
            let geojson_geom = match convert_wkb_to_geojson_geom(&geometry) {
                Some(g) => g,
                None => continue,
            };

            // Build properties
            let mut props = serde_json::Map::new();
            props.insert("id".to_string(), serde_json::Value::String(id));
            if let Some(n) = name {
                props.insert("name".to_string(), serde_json::Value::String(n));
            }
            if let Some(c) = class {
                props.insert("class".to_string(), serde_json::Value::String(c));
            }
            if let Some(o) = oneway {
                props.insert("oneway".to_string(), serde_json::Value::String(o));
            }

            features.push(Feature {
                id: None,
                bbox: None,
                geometry: Some(geojson_geom),
                properties: Some(props),
                foreign_members: None,
            });
        }

        tracing::info!("Converted {} features from Postgres", features.len());

        // Build graph statistics
        let (nodes, edges, total_km) = build_graph_stats(&features)?;

        // Write GeoJSON output
        let geojson = FeatureCollection {
            bbox: None,
            features,
            foreign_members: None,
        };

        let geojson_string = serde_json::to_string_pretty(&geojson)?;
        let output_path = &req.output_path;

        File::create(output_path)?
            .write_all(geojson_string.as_bytes())
            .context("Failed to write GeoJSON output")?;

        Ok(ExtractResult {
            nodes,
            edges,
            total_km,
            output_path: output_path.clone(),
        })
    })
}

/// Helper to convert wkb::reader::Geometry to geojson::Geometry
fn convert_wkb_to_geojson_geom(geom: &impl GeometryTrait<T = f64>) -> Option<GeoJsonGeometry> {
    let value = match geom.as_type() {
        GeometryType::LineString(ls) => {
            let coords = (0..ls.num_coords())
                .map(|i| {
                    let c = unsafe { ls.coord_unchecked(i) };
                    vec![c.x(), c.y()]
                })
                .collect();
            GeoJsonValue::LineString(coords)
        }
        GeometryType::Point(p) => {
            let c = p.coord()?;
            GeoJsonValue::Point(vec![c.x(), c.y()])
        }
        GeometryType::MultiLineString(mls) => {
            let coords = (0..mls.num_line_strings())
                .map(|i| {
                    let ls = unsafe { mls.line_string_unchecked(i) };
                    (0..ls.num_coords())
                        .map(|j| {
                            let c = unsafe { ls.coord_unchecked(j) };
                            vec![c.x(), c.y()]
                        })
                        .collect()
                })
                .collect();
            GeoJsonValue::MultiLineString(coords)
        }
        _ => return None, // Only supporting LineString, Point, MultiLineString for roads
    };

    Some(GeoJsonGeometry {
        bbox: None,
        value,
        foreign_members: None,
    })
}

/// Run Overture extraction from S3
fn run_overture_extract(req: &ExtractRequest) -> anyhow::Result<ExtractResult> {
    let runtime = tokio::runtime::Runtime::new()
        .context("Failed to create Tokio runtime for async S3 operations")?;

    runtime.block_on(async {
        let bbox: BBox = req.bbox.clone().into();
        let extractor = OvertureExtractor::new()?;

        tracing::info!(
            "Extracting Overture data for bbox: [{:.4}, {:.4}, {:.4}, {:.4}]",
            bbox.min_lon,
            bbox.min_lat,
            bbox.max_lon,
            bbox.max_lat
        );

        let segments = extractor.extract_bbox(&bbox).await?;

        tracing::info!("Extracted {} segments from Overture S3", segments.len());

        // Convert segments to GeoJSON features
        let features: Vec<Feature> = segments
            .into_iter()
            .filter(|seg| should_include_segment(seg, &req.road_classes))
            .map(|seg| segment_to_feature(seg))
            .collect();

        tracing::info!(
            "After road class filtering: {} road segments",
            features.len()
        );

        // Build graph statistics
        let (nodes, edges, total_km) = build_graph_stats(&features)?;

        // Write GeoJSON output
        let geojson = FeatureCollection {
            bbox: None,
            features,
            foreign_members: None,
        };

        let geojson_string = serde_json::to_string_pretty(&geojson)?;
        let output_path = &req.output_path;

        File::create(output_path)?
            .write_all(geojson_string.as_bytes())
            .context("Failed to write GeoJSON output")?;

        Ok(ExtractResult {
            nodes,
            edges,
            total_km,
            output_path: output_path.clone(),
        })
    })
}

/// Check if a segment should be included based on road classes
fn should_include_segment(seg: &OvertureSegment, classes: &[RoadClass]) -> bool {
    if classes.is_empty() {
        return true;
    }

    let class_str = seg.class.as_deref().unwrap_or("");

    classes.iter().any(|rc| rc.as_str() == class_str)
}

/// Convert an Overture segment to a GeoJSON feature
fn segment_to_feature(seg: OvertureSegment) -> Feature {
    let geometry_json = match seg.geometry {
        Geometry::LineString(coords) => {
            json!({
                "type": "LineString",
                "coordinates": coords
            })
        }
        Geometry::Point(lon, lat) => {
            json!({
                "type": "Point",
                "coordinates": [lon, lat]
            })
        }
    };

    // Convert JSON to geojson::Geometry
    let geometry = GeoJsonGeometry::from_json_value(geometry_json)
        .expect("Failed to convert JSON to geojson::Geometry");

    // Build properties map
    let mut props = serde_json::Map::new();
    props.insert("id".to_string(), serde_json::Value::String(seg.id.clone()));
    if let Some(ref name) = seg.name {
        props.insert("name".to_string(), serde_json::Value::String(name.clone()));
    }
    if let Some(ref class) = seg.class {
        props.insert("class".to_string(), serde_json::Value::String(class.clone()));
    }
    if let Some(ref subtype) = seg.subtype {
        props
            .insert("subtype".to_string(), serde_json::Value::String(subtype.clone()));
    }
    if let Some(ref surface) = seg.surface {
        props
            .insert("surface".to_string(), serde_json::Value::String(surface.clone()));
    }
    if let Some(ref oneway) = seg.oneway {
        props
            .insert("oneway".to_string(), serde_json::Value::String(oneway.clone()));
    }
    if let Some(ref junction) = seg.junction {
        props
            .insert("junction".to_string(), serde_json::Value::String(junction.clone()));
    }
    if let Some(ref osm_id) = seg.osm_id {
        props
            .insert("osm_id".to_string(), serde_json::Value::String(osm_id.clone()));
    }

    Feature {
        id: None,
        bbox: None,
        geometry: Some(geometry),
        properties: Some(props),
        foreign_members: None,
    }
}

/// Build graph statistics from features
fn build_graph_stats(features: &[Feature]) -> Result<(usize, usize, f64)> {
    // Node deduplication: snap coordinates to ~1m precision
    let mut node_map: HashMap<(i64, i64), usize> = HashMap::new();
    let mut next_node_id: usize = 0;
    let mut edge_count = 0;
    let mut total_km = 0.0;

    for f in features {
        if let Some(ref geom) = f.geometry {
            let line_strings: Vec<&Vec<Vec<f64>>> = match &geom.value {
                GeoJsonValue::LineString(coords) => vec![coords],
                GeoJsonValue::MultiLineString(multi) => multi.iter().collect(),
                _ => continue,
            };

            for coords in line_strings {
                if coords.len() < 2 {
                    continue;
                }

                // Extract coordinates as (lon, lat) pairs
                let coord_points: Vec<(f64, f64)> = coords
                    .iter()
                    .filter(|p| p.len() >= 2)
                    .map(|p| (p[0], p[1]))
                    .collect();

                // Calculate length for each segment
                for window in coord_points.windows(2) {
                    let (lon1, lat1) = window[0];
                    let (lon2, lat2) = window[1];

                    // Haversine distance in km
                    let d = haversine_distance_km(lat1, lon1, lat2, lon2);
                    total_km += d;

                    // Get/create node IDs
                    let _node1 = get_or_create_node(&mut node_map, &mut next_node_id, lon1, lat1);
                    let _node2 = get_or_create_node(&mut node_map, &mut next_node_id, lon2, lat2);

                    // Count edge
                    edge_count += 1;
                }
            }
        }
    }

    Ok((node_map.len(), edge_count, total_km))
}

/// Get or create a node ID for given coordinates
fn get_or_create_node(
    node_map: &mut HashMap<(i64, i64), usize>,
    next_node_id: &mut usize,
    lon: f64,
    lat: f64,
) -> usize {
    // Snap to ~1m precision (6 decimal places)
    let key = ((lon * 1e6) as i64, (lat * 1e6) as i64);
    *node_map.entry(key).or_insert_with(|| {
        let id = *next_node_id;
        *next_node_id += 1;
        id
    })
}

/// Calculate Haversine distance between two points in km
fn haversine_distance_km(lat1: f64, lon1: f64, lat2: f64, lon2: f64) -> f64 {
    const EARTH_RADIUS_KM: f64 = 6371.0;

    let lat1_rad = lat1.to_radians();
    let lat2_rad = lat2.to_radians();
    let delta_lat = (lat2 - lat1).to_radians();
    let delta_lon = (lon2 - lon1).to_radians();

    let a = (delta_lat / 2.0).sin().powi(2)
        + lat1_rad.cos() * lat2_rad.cos() * (delta_lon / 2.0).sin().powi(2);
    let c = 2.0 * a.sqrt().atan2((1.0 - a).sqrt());

    EARTH_RADIUS_KM * c
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_haversine_distance() {
        // Distance from New York to Los Angeles (approx 3935 km)
        let ny_lat = 40.7128;
        let ny_lon = -74.0060;
        let la_lat = 34.0522;
        let la_lon = -118.2437;

        let dist = haversine_distance_km(ny_lat, ny_lon, la_lat, la_lon);
        assert!((dist - 3935.0).abs() < 10.0);
    }

    #[test]
    fn test_road_class_conversion() {
        assert_eq!(RoadClass::Motorway.as_str(), "motorway");
        assert_eq!(RoadClass::Residential.as_str(), "residential");
        assert_eq!(RoadClass::Tertiary.as_str(), "tertiary");
    }
}