use anyhow::Result;
use clap::Parser;
use flatgeobuf::{ColumnType, FgbCrs, FgbWriter, FgbWriterOptions, GeometryType};
use geo::{point, Contains, Geometry, Polygon};
use geozero::{ColumnValue, PropertyProcessor};
use hashbrown::{HashMap, HashSet};
use log::{info, LevelFilter};
use once_cell::sync::Lazy;
use osmpbfreader::{NodeId, OsmObj, OsmPbfReader, Way};
use serde::Deserialize;
use simplelog::{ColorChoice, Config, TermLogger, TerminalMode};
use smartstring::{LazyCompact, SmartString};
use std::{
    fs::File,
    io::{BufReader, BufWriter, Read, Write},
    path::PathBuf,
};

static VALID_ACCESS_TAGS: Lazy<HashSet<&'static str>> = Lazy::new(|| {
    let mut s = HashSet::new();
    s.insert("yes");
    s.insert("permisive");
    s.insert("destination");
    s
});

/// Convert an OpenStreetMap file into a road-network file readable by Metropolis.
#[derive(Parser, Debug)]
#[command(author, version, about, long_about)]
struct Args {
    /// Path to the input .osm.pbf file.
    #[arg(long)]
    osm_file: PathBuf,
    /// Path to the JSON file with the parameters.
    #[arg(long)]
    parameters: PathBuf,
}

#[derive(Clone, Debug, Default, Deserialize)]
struct Parameters {
    /// Polygon representing the area to clip (in EPSG:4326).
    boundary: Option<Polygon>,
    /// Set of highway tags to filter.
    valid_highways: HashSet<SmartString<LazyCompact>>,
}

fn is_valid_highway(way: &Way, parameters: &Parameters) -> bool {
    let has_access =
        VALID_ACCESS_TAGS.contains(way.tags.get("access").map(|s| s.as_str()).unwrap_or("yes"));
    let is_valid_highway = way
        .tags
        .get("highway")
        .map(|t| parameters.valid_highways.contains(t))
        .unwrap_or(false);
    way.nodes.len() >= 2 && has_access && is_valid_highway
}

fn main() -> Result<()> {
    let args = Args::parse();

    TermLogger::init(
        LevelFilter::Info,
        Config::default(),
        TerminalMode::Mixed,
        ColorChoice::Auto,
    )
    .expect("Failed to initialize logging");

    info!("Reading OSM file");
    let file = File::open(&args.osm_file)
        .unwrap_or_else(|_| panic!("Cannot open file: {:?}", args.osm_file));
    let buf = BufReader::new(file);
    let mut pbf = OsmPbfReader::new(buf);

    info!("Reading parameters");
    let mut bytes = Vec::new();
    File::open(&args.parameters)
        .unwrap_or_else(|_| panic!("Cannot open file: {:?}", args.parameters))
        .read_to_end(&mut bytes)
        .unwrap_or_else(|_| panic!("Cannot read file: {:?}", args.parameters));
    let parameters: Parameters =
        serde_json::from_slice(&bytes).expect("Unable to parse parameters");

    // Find all the nodes and ways that we need to keep.
    //
    // A node is valid if all the following conditions are met:
    // - The node is part of at least two valid ways OR it is the first or last node of a valid
    // way.
    // - The node is inside the boundary to clip (if any).
    let mut valid_ways = HashSet::new();
    let mut source_target_nodes: HashSet<NodeId> = HashSet::new();
    let mut all_nodes_count: HashMap<NodeId, usize> = HashMap::new();
    for obj in pbf.par_iter().filter_map(|r| r.ok()) {
        if let OsmObj::Way(way) = obj {
            if is_valid_highway(&way, &parameters) {
                valid_ways.insert(way.id);
                source_target_nodes.insert(way.nodes[0]);
                source_target_nodes.insert(way.nodes[way.nodes.len() - 1]);
                for &node_id in way.nodes.iter() {
                    *all_nodes_count.entry(node_id).or_default() += 1;
                }
            }
        }
    }

    let nodes_to_keep: HashSet<NodeId> = source_target_nodes
        .into_iter()
        .chain(
            all_nodes_count
                .into_iter()
                .filter_map(|(k, v)| if v >= 2 { Some(k) } else { None }),
        )
        .collect();

    // Find the nodes inside the boundary.
    let mut valid_nodes: HashSet<NodeId> = HashSet::new();
    pbf.rewind().unwrap();
    for obj in pbf.par_iter().filter_map(|r| r.ok()) {
        if let OsmObj::Node(node) = obj {
            if nodes_to_keep.contains(&node.id) {
                let point = point!(x: node.lon(), y: node.lat());
                if parameters
                    .boundary
                    .as_ref()
                    .map(|b| b.contains(&point))
                    .unwrap_or(true)
                {
                    valid_nodes.insert(node.id);
                }
            }
        }
    }

    println!("Number of nodes to keep: {}", valid_nodes.len());

    // Store all the nodes that we keep in a FlatGeoBuf with their coordinates.
    let options = FgbWriterOptions {
        write_index: true,
        crs: FgbCrs {
            code: 4326,
            ..Default::default()
        },
        ..Default::default()
    };
    let mut fgb = FgbWriter::create_with_options("nodes", GeometryType::Point, options).unwrap();
    fgb.add_column("fid", ColumnType::Long, |_fbb, col| {
        col.nullable = false;
        col.unique = true;
    });

    pbf.rewind().unwrap();
    for obj in pbf.par_iter().filter_map(|r| r.ok()) {
        if let OsmObj::Node(node) = obj {
            if nodes_to_keep.contains(&node.id) {
                fgb.add_feature_geom(
                    Geometry::Point(point!(x: node.lon(), y: node.lat())),
                    |feat| {
                        feat.property(0, "fid", &ColumnValue::Long(node.id.0))
                            .unwrap();
                    },
                )
                .unwrap();
            }
        }
    }

    let output_path = "../../../output/road_network/france-nodes.fgb";
    let mut fout = BufWriter::new(File::create(output_path).expect("Cannot create output file"));
    fgb.write(&mut fout).unwrap();

    pbf.rewind().unwrap();
    let mut tags: HashMap<SmartString<LazyCompact>, usize> = HashMap::new();
    for obj in pbf.par_iter().filter_map(|r| r.ok()) {
        if let OsmObj::Way(way) = obj {
            if valid_ways.contains(&way.id) {
                for (key, _value) in way.tags.iter() {
                    *tags.entry(key.clone()).or_default() += 1;
                }
            }
        }
    }

    let buf = serde_json::to_vec(&tags)?;
    let mut file = BufWriter::new(File::open("tmp.json")?);
    file.write_all(&buf)?;

    Ok(())
}