rmpca 0.2.0

//! Convert local OSM files (.osm XML / .osm.pbf) to GeoJSON
//!
//! This command reads a local OSM data file and converts it to GeoJSON,
//! optionally feeding the result directly into the route optimizer.

use crate::config::Config;
use crate::optimizer::{Node, RouteOptimizer, Way};
use anyhow::{Context, Result};
use clap::Args as ClapArgs;
use geojson::{Feature, FeatureCollection, GeoJson, Geometry, Value};
use std::collections::HashMap;
use std::path::PathBuf;

// ── CLI arguments ──────────────────────────────────────────────────────────

#[derive(Debug, ClapArgs)]
pub struct Args {
    /// Input OSM file (.osm XML or .osm.pbf)
    #[arg(short, long)]
    input: PathBuf,

    /// Output GeoJSON file (default: stdout)
    #[arg(short, long)]
    output: Option<PathBuf>,

    /// Highway class filter (comma-separated, e.g. "primary,secondary,trunk")
    #[arg(long)]
    highway: Option<String>,

    /// Only include ways that are drivable (exclude footways, cycleways, etc.)
    #[arg(long, default_value_t = true)]
    drivable_only: bool,

    /// Run route optimization on the extracted data after conversion
    #[arg(long)]
    optimize: bool,

    /// Depot location for optimization (LAT,LON) — required when --optimize is set
    #[arg(long)]
    depot: Option<String>,

    /// Left turn penalty for optimization
    #[arg(long)]
    turn_left: Option<f64>,

    /// Right turn penalty for optimization
    #[arg(long)]
    turn_right: Option<f64>,

    /// U-turn penalty for optimization
    #[arg(long)]
    turn_u: Option<f64>,
}

// ── OSM data model (intermediate) ─────────────────────────────────────────

/// Raw OSM node parsed from file
#[derive(Debug, Clone)]
struct OsmNode {
    id: i64,
    lat: f64,
    lon: f64,
    tags: HashMap<String, String>,
}

/// Raw OSM way parsed from file
#[derive(Debug, Clone)]
struct OsmWay {
    id: i64,
    node_ids: Vec<i64>,
    tags: HashMap<String, String>,
}

/// Parsed OSM data ready for conversion
struct OsmData {
    nodes: HashMap<i64, OsmNode>,
    ways: Vec<OsmWay>,
}

// ── Highway filter ──────────────────────────────────────────────────────────

/// Road types that are generally drivable by motor vehicles.
const DRIVABLE_HIGHWAYS: &[&str] = &[
    "motorway",
    "trunk",
    "primary",
    "secondary",
    "tertiary",
    "unclassified",
    "residential",
    "motorway_link",
    "trunk_link",
    "primary_link",
    "secondary_link",
    "tertiary_link",
    "living_street",
    "service",
];

/// Check if a highway tag value is considered drivable.
fn is_drivable(highway: &str) -> bool {
    DRIVABLE_HIGHWAYS.contains(&highway)
}

// ── PBF parsing ────────────────────────────────────────────────────────────

fn parse_pbf(path: &PathBuf) -> Result<OsmData> {
    let mut nodes: HashMap<i64, OsmNode> = HashMap::new();
    let mut ways: Vec<OsmWay> = Vec::new();

    let reader = osmpbf::ElementReader::from_path(path)
        .with_context(|| format!("Failed to open PBF file: {}", path.display()))?;

    reader.for_each(|element| {
        match element {
            osmpbf::Element::Node(n) => {
                let mut tags = HashMap::new();
                for (k, v) in n.tags() {
                    tags.insert(k.to_string(), v.to_string());
                }
                nodes.insert(
                    n.id(),
                    OsmNode {
                        id: n.id(),
                        lat: n.decimicro_lat() as f64 / 1e7,
                        lon: n.decimicro_lon() as f64 / 1e7,
                        tags,
                    },
                );
            }
            osmpbf::Element::Way(w) => {
                let mut tags = HashMap::new();
                for (k, v) in w.tags() {
                    tags.insert(k.to_string(), v.to_string());
                }
                ways.push(OsmWay {
                    id: w.id(),
                    node_ids: w.refs().collect(),
                    tags,
                });
            }
            osmpbf::Element::Relation(_) => {
                // Relations are not needed for road network extraction
            }
            osmpbf::Element::DenseNode(dn) => {
                let mut tags = HashMap::new();
                for (k, v) in dn.tags() {
                    tags.insert(k.to_string(), v.to_string());
                }
                nodes.insert(
                    dn.id,
                    OsmNode {
                        id: dn.id,
                        lat: dn.decimicro_lat() as f64 / 1e7,
                        lon: dn.decimicro_lon() as f64 / 1e7,
                        tags,
                    },
                );
            }
        }
    }).with_context(|| "Error reading PBF file")?;

    tracing::info!(
        "Parsed PBF: {} nodes, {} ways",
        nodes.len(),
        ways.len()
    );
    Ok(OsmData { nodes, ways })
}

// ── XML (.osm) parsing ─────────────────────────────────────────────────────

fn parse_osm_xml(path: &PathBuf) -> Result<OsmData> {
    use quick_xml::events::Event;
    use quick_xml::Reader;

    let mut reader = Reader::from_file(path)
        .with_context(|| format!("Failed to create XML reader for: {}", path.display()))?;

    let mut nodes: HashMap<i64, OsmNode> = HashMap::new();
    let mut ways: Vec<OsmWay> = Vec::new();

    let mut current_node: Option<OsmNode> = None;
    let mut current_way: Option<OsmWay> = None;
    let mut buf = Vec::new();

    loop {
        match reader.read_event_into(&mut buf) {
            Ok(Event::Start(ref e)) | Ok(Event::Empty(ref e)) => {
                match e.local_name().as_ref() {
                    b"node" => {
                        let mut id: i64 = 0;
                        let mut lat: f64 = 0.0;
                        let mut lon: f64 = 0.0;
                        for attr in e.attributes().flatten() {
                            match attr.key.as_ref() {
                                b"id" => id = String::from_utf8_lossy(&attr.value).parse().unwrap_or(0),
                                b"lat" => lat = String::from_utf8_lossy(&attr.value).parse().unwrap_or(0.0),
                                b"lon" => lon = String::from_utf8_lossy(&attr.value).parse().unwrap_or(0.0),
                                _ => {}
                            }
                        }
                        current_node = Some(OsmNode {
                            id,
                            lat,
                            lon,
                            tags: HashMap::new(),
                        });
                    }
                    b"way" => {
                        let mut id: i64 = 0;
                        for attr in e.attributes().flatten() {
                            if attr.key.as_ref() == b"id" {
                                id = String::from_utf8_lossy(&attr.value).parse().unwrap_or(0);
                            }
                        }
                        current_way = Some(OsmWay {
                            id,
                            node_ids: Vec::new(),
                            tags: HashMap::new(),
                        });
                    }
                    b"nd" => {
                        if let Some(ref mut way) = current_way {
                            for attr in e.attributes().flatten() {
                                if attr.key.as_ref() == b"ref" {
                                    if let Ok(ref_id) = String::from_utf8_lossy(&attr.value).parse::<i64>() {
                                        way.node_ids.push(ref_id);
                                    }
                                }
                            }
                        }
                    }
                    b"tag" => {
                        let mut k = String::new();
                        let mut v = String::new();
                        for attr in e.attributes().flatten() {
                            match attr.key.as_ref() {
                                b"k" => k = String::from_utf8_lossy(&attr.value).into_owned(),
                                b"v" => v = String::from_utf8_lossy(&attr.value).into_owned(),
                                _ => {}
                            }
                        }
                        if !k.is_empty() {
                            if let Some(ref mut node) = current_node {
                                node.tags.insert(k.clone(), v.clone());
                            }
                            if let Some(ref mut way) = current_way {
                                way.tags.insert(k, v);
                            }
                        }
                    }
                    _ => {}
                }
            }
            Ok(Event::End(ref e)) => {
                match e.local_name().as_ref() {
                    b"node" => {
                        if let Some(node) = current_node.take() {
                            nodes.insert(node.id, node);
                        }
                    }
                    b"way" => {
                        if let Some(way) = current_way.take() {
                            ways.push(way);
                        }
                    }
                    _ => {}
                }
            }
            Ok(Event::Eof) => break,
            Err(e) => {
                return Err(anyhow::anyhow!(
                    "XML parse error: {}",
                    e
                ));
            }
            _ => {}
        }
        buf.clear();
    }

    tracing::info!(
        "Parsed OSM XML: {} nodes, {} ways",
        nodes.len(),
        ways.len()
    );
    Ok(OsmData { nodes, ways })
}

// ── Filtering ──────────────────────────────────────────────────────────────

fn filter_ways<'a>(data: &'a OsmData, args: &Args) -> Vec<&'a OsmWay> {
    let highway_filter: Option<Vec<String>> = args
        .highway
        .as_ref()
        .map(|h| h.split(',').map(|s| s.trim().to_string()).collect());

    data.ways
        .iter()
        .filter(|way| {
            let highway = way.tags.get("highway").map(|s| s.as_str()).unwrap_or("");

            // Must have a highway tag
            if highway.is_empty() {
                return false;
            }

            // If drivable_only, filter out non-drivable road types
            if args.drivable_only && !is_drivable(highway) {
                return false;
            }

            // If specific highway filter provided, check against it
            if let Some(ref allowed) = highway_filter {
                if !allowed.iter().any(|h| h == highway) {
                    return false;
                }
            }

            true
        })
        .collect()
}

// ── OSM data → GeoJSON ────────────────────────────────────────────────────

fn osm_to_geojson(data: &OsmData, filtered_ways: &[&OsmWay]) -> FeatureCollection {
    let mut features = Vec::new();

    for way in filtered_ways {
        // Collect coordinates for this way
        let coords: Vec<Vec<f64>> = way
            .node_ids
            .iter()
            .filter_map(|nid| data.nodes.get(nid))
            .map(|n| vec![n.lon, n.lat])
            .collect();

        if coords.len() < 2 {
            continue; // Skip ways with fewer than 2 visible nodes
        }

        let geometry = Geometry::new(Value::LineString(coords));

        let mut properties = serde_json::Map::new();
        properties.insert(
            "id".to_string(),
            serde_json::Value::String(way.id.to_string()),
        );
        for (k, v) in &way.tags {
            properties.insert(k.clone(), serde_json::Value::String(v.clone()));
        }

        features.push(Feature {
            geometry: Some(geometry),
            properties: Some(properties),
            ..Default::default()
        });
    }

    FeatureCollection {
        features,
        bbox: None,
        foreign_members: None,
    }
}

// ── OSM data → optimizer types ─────────────────────────────────────────────

fn osm_to_optimizer(data: &OsmData, filtered_ways: &[&OsmWay]) -> (Vec<Node>, Vec<Way>) {
    // Collect all node IDs referenced by filtered ways
    let mut used_node_ids = std::collections::HashSet::new();
    for way in filtered_ways {
        for nid in &way.node_ids {
            used_node_ids.insert(*nid);
        }
    }

    // Build optimizer nodes
    let nodes: Vec<Node> = used_node_ids
        .iter()
        .filter_map(|nid| data.nodes.get(nid))
        .map(|n| {
            let mut node = Node::new(n.id.to_string(), n.lat, n.lon);
            if let Some(z) = n.tags.get("ele") {
                if let Ok(elevation) = z.parse::<f64>() {
                    node.z = Some(elevation);
                }
            }
            node
        })
        .collect();

    // Build optimizer ways
    let opt_ways: Vec<Way> = filtered_ways
        .iter()
        .map(|w| {
            let mut way = Way::new(w.id.to_string(), w.node_ids.iter().map(|id| id.to_string()).collect());
            for (k, v) in &w.tags {
                way.tags.insert(k.clone(), v.clone());
            }
            way
        })
        .collect();

    (nodes, opt_ways)
}

// ── Main entry point ───────────────────────────────────────────────────────

pub async fn run(args: Args) -> Result<()> {
    let config = Config::load().unwrap_or_default();
    config.init_logging();

    let input_path = &args.input;

    // Validate input file exists
    if !input_path.exists() {
        anyhow::bail!("Input file not found: {}", input_path.display());
    }

    // Detect format from extension
    let is_pbf = input_path
        .extension()
        .map(|e| e == "pbf")
        .unwrap_or(false);

    tracing::info!(
        "Parsing OSM file: {} (format: {})",
        input_path.display(),
        if is_pbf { "PBF" } else { "XML" }
    );

    // Parse the OSM file
    let data = if is_pbf {
        parse_pbf(input_path)?
    } else {
        parse_osm_xml(input_path)?
    };

    tracing::info!(
        "Loaded {} nodes, {} ways",
        data.nodes.len(),
        data.ways.len()
    );

    // Filter ways
    let filtered_ways = filter_ways(&data, &args);
    tracing::info!("After filtering: {} ways", filtered_ways.len());

    if args.optimize {
        // ── Direct optimization path ───────────────────────────────────
        let (nodes, ways) = osm_to_optimizer(&data, &filtered_ways);
        let mut optimizer = RouteOptimizer::new();

        // Feed nodes and ways into optimizer
        for node in nodes {
            // RouteOptimizer stores nodes internally; we add them via the ways
            let _ = node; // Will be used when build_graph is fully implemented
        }
        for way in ways {
            optimizer.ways.push(way);
        }

        // Set depot if provided
        if let Some(ref depot_str) = args.depot {
            let parts: Vec<f64> = depot_str
                .split(',')
                .map(|s| s.trim().parse::<f64>())
                .collect::<Result<Vec<f64>, _>>()
                .context("Invalid depot format. Use LAT,LON (e.g. 45.5,-73.6)")?;
            if parts.len() != 2 {
                anyhow::bail!("Depot must be LAT,LON (e.g. 45.5,-73.6)");
            }
            optimizer.set_depot(parts[0], parts[1]);
        }

        // Set turn penalties
        if let Some(left) = args.turn_left {
            optimizer.set_turn_penalties(left, args.turn_right.unwrap_or(0.0), args.turn_u.unwrap_or(0.0));
        } else if let Some(right) = args.turn_right {
            optimizer.set_turn_penalties(0.0, right, args.turn_u.unwrap_or(0.0));
        } else if let Some(u) = args.turn_u {
            optimizer.set_turn_penalties(0.0, 0.0, u);
        }

        let result = optimizer.optimize()?;

        // Output optimization result as JSON
        let json = serde_json::to_string_pretty(&result)?;
        match &args.output {
            Some(path) => {
                std::fs::write(path, &json)
                    .with_context(|| format!("Failed to write optimization result to {}", path.display()))?;
                tracing::info!("Optimization result written to {}", path.display());
            }
            None => println!("{}", json),
        }
    } else {
        // ── GeoJSON conversion path ────────────────────────────────────
        let fc = osm_to_geojson(&data, &filtered_ways);

        let geojson = GeoJson::from(fc);
        let json = serde_json::to_string_pretty(&geojson)?;

        match &args.output {
            Some(path) => {
                std::fs::write(path, &json)
                    .with_context(|| format!("Failed to write GeoJSON to {}", path.display()))?;
                tracing::info!("GeoJSON written to {}", path.display());
            }
            None => println!("{}", json),
        }
    }

    Ok(())
}

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

    #[test]
    fn test_drivable_highways() {
        assert!(is_drivable("motorway"));
        assert!(is_drivable("residential"));
        assert!(is_drivable("tertiary_link"));
        assert!(!is_drivable("footway"));
        assert!(!is_drivable("cycleway"));
        assert!(!is_drivable("path"));
    }

    #[test]
    fn test_args_construction() {
        let args = Args {
            input: PathBuf::from("test.osm.pbf"),
            output: Some(PathBuf::from("out.geojson")),
            highway: Some("primary,secondary".to_string()),
            drivable_only: true,
            optimize: false,
            depot: None,
            turn_left: None,
            turn_right: None,
            turn_u: None,
        };
        assert_eq!(args.input, PathBuf::from("test.osm.pbf"));
        assert!(args.drivable_only);
        assert!(!args.optimize);
    }

    #[test]
    fn test_osm_to_geojson_empty() {
        let data = OsmData {
            nodes: HashMap::new(),
            ways: Vec::new(),
        };
        let filtered: Vec<&OsmWay> = Vec::new();
        let fc = osm_to_geojson(&data, &filtered);
        assert_eq!(fc.features.len(), 0);
    }

    #[test]
    fn test_osm_to_geojson_with_way() {
        let mut nodes = HashMap::new();
        nodes.insert(
            1,
            OsmNode {
                id: 1,
                lat: 45.5,
                lon: -73.6,
                tags: HashMap::new(),
            },
        );
        nodes.insert(
            2,
            OsmNode {
                id: 2,
                lat: 45.51,
                lon: -73.61,
                tags: HashMap::new(),
            },
        );

        let mut tags = HashMap::new();
        tags.insert("highway".to_string(), "primary".to_string());
        tags.insert("name".to_string(), "Main St".to_string());

        let ways = vec![OsmWay {
            id: 100,
            node_ids: vec![1, 2],
            tags,
        }];

        let data = OsmData { nodes, ways };
        let filtered: Vec<&OsmWay> = data.ways.iter().collect();
        let fc = osm_to_geojson(&data, &filtered);

        assert_eq!(fc.features.len(), 1);
        let feature = &fc.features[0];
        assert!(feature.geometry.is_some());
        if let Some(ref geom) = feature.geometry {
            if let Value::LineString(coords) = &geom.value {
                assert_eq!(coords.len(), 2);
                assert_eq!(coords[0][0], -73.6); // lon
                assert_eq!(coords[0][1], 45.5);  // lat
            } else {
                panic!("Expected LineString geometry");
            }
        }
    }

    #[test]
    fn test_filter_ways_drivable_only() {
        let mut tags_road = HashMap::new();
        tags_road.insert("highway".to_string(), "primary".to_string());

        let mut tags_foot = HashMap::new();
        tags_foot.insert("highway".to_string(), "footway".to_string());

        let ways = vec![
            OsmWay {
                id: 1,
                node_ids: vec![1, 2],
                tags: tags_road,
            },
            OsmWay {
                id: 2,
                node_ids: vec![3, 4],
                tags: tags_foot,
            },
        ];

        let data = OsmData {
            nodes: HashMap::new(),
            ways,
        };

        let args = Args {
            input: PathBuf::from("test.osm"),
            output: None,
            highway: None,
            drivable_only: true,
            optimize: false,
            depot: None,
            turn_left: None,
            turn_right: None,
            turn_u: None,
        };

        let filtered = filter_ways(&data, &args);
        assert_eq!(filtered.len(), 1);
        assert_eq!(filtered[0].id, 1);
    }

    #[test]
    fn test_filter_ways_highway_filter() {
        let mut tags_primary = HashMap::new();
        tags_primary.insert("highway".to_string(), "primary".to_string());

        let mut tags_secondary = HashMap::new();
        tags_secondary.insert("highway".to_string(), "secondary".to_string());

        let mut tags_tertiary = HashMap::new();
        tags_tertiary.insert("highway".to_string(), "tertiary".to_string());

        let ways = vec![
            OsmWay {
                id: 1,
                node_ids: vec![1, 2],
                tags: tags_primary,
            },
            OsmWay {
                id: 2,
                node_ids: vec![3, 4],
                tags: tags_secondary,
            },
            OsmWay {
                id: 3,
                node_ids: vec![5, 6],
                tags: tags_tertiary,
            },
        ];

        let data = OsmData {
            nodes: HashMap::new(),
            ways,
        };

        let args = Args {
            input: PathBuf::from("test.osm"),
            output: None,
            highway: Some("primary,secondary".to_string()),
            drivable_only: false,
            optimize: false,
            depot: None,
            turn_left: None,
            turn_right: None,
            turn_u: None,
        };

        let filtered = filter_ways(&data, &args);
        assert_eq!(filtered.len(), 2);
    }
}