perfect-postcode/server-rs/src/routes/postcodes.rs

use std::sync::Arc;

use axum::extract::{Path, Query, State};
use axum::http::StatusCode;
use axum::response::{IntoResponse, Json};
use axum::Extension;
use metrics::histogram;
use rustc_hash::FxHashMap;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use tracing::info;

use crate::aggregation::{Aggregator, EnumDistConfig};
use crate::auth::OptionalUser;
use crate::consts::MAX_CELLS_PER_REQUEST;
use crate::data::travel_time::TravelData;
use crate::licensing::check_license_bounds;
use crate::parsing::{
    bounds_intersect, parse_enum_dist, parse_field_indices, parse_filters, require_bounds,
    row_passes_filters,
};
use crate::pocketbase::log_user_location;
use crate::routes::travel_time::{parse_optional_travel, TravelTimeAgg};
use crate::state::SharedState;
use crate::utils::normalize_postcode;

#[derive(Serialize)]
pub struct PostcodesResponse {
    r#type: &'static str,
    features: Vec<Map<String, Value>>,
}

#[derive(Deserialize)]
pub struct NearestPostcodeParams {
    lat: f64,
    lng: f64,
}

#[derive(Deserialize)]
pub struct PostcodeParams {
    bounds: Option<String>,
    /// `;;`-separated filters: `name:min:max;;...`
    filters: Option<String>,
    /// Comma-separated feature names to include in min/max aggregation.
    fields: Option<String>,
    /// Pipe-separated travel time entries: `mode:slug|mode:slug:min:max`
    travel: Option<String>,
    /// Feature name for enum distribution counting (pie chart visualization).
    enum_dist: Option<String>,
}

pub async fn get_postcodes(
    State(shared): State<Arc<SharedState>>,
    Extension(user): Extension<OptionalUser>,
    Query(params): Query<PostcodeParams>,
) -> Result<Json<PostcodesResponse>, axum::response::Response> {
    let state = shared.load_state();
    let (south, west, north, east) =
        require_bounds(params.bounds).map_err(IntoResponse::into_response)?;

    check_license_bounds(&user.0, (south, west, north, east))?;

    let quant = state.data.quant_ref();
    let (parsed_filters, parsed_enum_filters) = parse_filters(
        params.filters.as_deref(),
        &state.feature_name_to_index,
        &state.data.enum_values,
        &quant,
    )
    .map_err(|err| (StatusCode::BAD_REQUEST, err).into_response())?;
    let num_filters = parsed_filters.len() + parsed_enum_filters.len();
    let filters_str = params.filters;

    let field_indices = parse_field_indices(params.fields.as_deref(), &state.feature_name_to_index)
        .map_err(|err| (err.0, err.1).into_response())?;

    let travel_entries = parse_optional_travel(params.travel.as_deref())
        .map_err(|err| (StatusCode::BAD_REQUEST, err).into_response())?;

    let enum_dist_config: EnumDistConfig = parse_enum_dist(
        params.enum_dist.as_deref(),
        &state.feature_name_to_index,
        &state.data.enum_values,
    )
    .map_err(|err| (err.0, err.1).into_response())?;

    let enum_dist_key: Option<String> = params
        .enum_dist
        .as_ref()
        .map(|name| format!("dist_{}", name.trim()));

    let response = tokio::task::spawn_blocking(move || -> Result<PostcodesResponse, String> {
        let postcode_data = &state.postcode_data;
        let t0 = std::time::Instant::now();

        // Load travel time data from precomputed parquet files
        let travel_data: Vec<TravelData> = if !travel_entries.is_empty() {
            let store = &state.travel_time_store;
            travel_entries
                .iter()
                .map(|entry| {
                    store
                        .get(&entry.mode, &entry.slug)
                        .map_err(|err| format!("Failed to load travel data: {}", err))
                })
                .collect::<Result<Vec<_>, _>>()?
        } else {
            Vec::new()
        };

        let has_travel = !travel_entries.is_empty();
        let travel_field_keys: Vec<String> = travel_entries
            .iter()
            .map(|te| format!("tt_{}_{}", te.mode, te.slug))
            .collect();

        let num_features = state.data.num_features;
        let feature_data = &state.data.feature_data;
        let quant = state.data.quant_ref();
        let min_keys = &state.min_keys;
        let max_keys = &state.max_keys;
        let avg_keys = &state.avg_keys;

        let has_selective = field_indices.is_some();
        let sel_indices = field_indices.as_deref().unwrap_or(&[]);

        // Single-pass: aggregate directly into postcode_aggs while iterating properties in bounds
        let mut postcode_aggs: FxHashMap<usize, Aggregator> = FxHashMap::default();

        state
            .grid
            .for_each_in_bounds(south, west, north, east, |row_idx| {
                let row = row_idx as usize;
                if !row_passes_filters(
                    row,
                    &parsed_filters,
                    &parsed_enum_filters,
                    feature_data,
                    num_features,
                ) {
                    return;
                }

                let postcode = state.data.postcode(row);
                if let Some(&pc_idx) = postcode_data.postcode_to_idx.get(postcode) {
                    let agg = postcode_aggs
                        .entry(pc_idx)
                        .or_insert_with(|| Aggregator::new(num_features, enum_dist_config));
                    if has_selective {
                        agg.add_row_selective(feature_data, row, num_features, sel_indices, &quant);
                    } else {
                        agg.add_row(feature_data, row, num_features, &quant);
                    }
                }
            });

        // Filter postcodes by travel time range (if specified)
        if has_travel {
            postcode_aggs.retain(|&pc_idx, _agg| {
                let postcode = &postcode_data.postcodes[pc_idx];
                for (ti, entry) in travel_entries.iter().enumerate() {
                    if let (Some(fmin), Some(fmax)) = (entry.filter_min, entry.filter_max) {
                        let minutes = travel_data[ti].get(postcode.as_str()).map(|r| {
                            if entry.use_best {
                                r.best_minutes.unwrap_or(r.minutes)
                            } else {
                                r.minutes
                            }
                        });
                        match minutes {
                            Some(mins) if (mins as f32) >= fmin && (mins as f32) <= fmax => {}
                            _ => return false,
                        }
                    }
                }
                true
            });
        }

        // Travel time aggregation per postcode
        let mut travel_aggs: FxHashMap<usize, Vec<TravelTimeAgg>> = FxHashMap::default();
        if has_travel {
            for &pc_idx in postcode_aggs.keys() {
                let postcode = &postcode_data.postcodes[pc_idx];
                let tt_aggs = travel_aggs.entry(pc_idx).or_insert_with(|| {
                    (0..travel_entries.len())
                        .map(|_| TravelTimeAgg::new())
                        .collect()
                });
                for (ti, entry) in travel_entries.iter().enumerate() {
                    if let Some(row_data) = travel_data[ti].get(postcode.as_str()) {
                        let minutes = if entry.use_best {
                            row_data.best_minutes.unwrap_or(row_data.minutes)
                        } else {
                            row_data.minutes
                        };
                        tt_aggs[ti].add(minutes as f32);
                    }
                }
            }
        }

        let t_agg = t0.elapsed();

        // Build response, filtering postcodes to only those whose polygon intersects query bounds
        let mut features = Vec::with_capacity(postcode_aggs.len());
        let postcodes_before_filter = postcode_aggs.len();
        let mut filtered_out = 0usize;

        for (pc_idx, aggregation) in postcode_aggs {
            if aggregation.count == 0 {
                continue;
            }

            // Use precomputed AABB for bounds intersection check
            let (pc_south, pc_west, pc_north, pc_east) = postcode_data.aabbs[pc_idx];

            if !bounds_intersect(
                pc_south as f64,
                pc_west as f64,
                pc_north as f64,
                pc_east as f64,
                south,
                west,
                north,
                east,
            ) {
                filtered_out += 1;
                continue;
            }

            let geometry = postcode_data.geometries[pc_idx].clone();

            // Build properties
            let centroid = postcode_data.centroids[pc_idx];
            let mut props = Map::new();
            props.insert(
                "postcode".into(),
                Value::String(postcode_data.postcodes[pc_idx].clone()),
            );
            props.insert("count".into(), Value::Number(aggregation.count.into()));
            props.insert(
                "centroid".into(),
                Value::Array(vec![
                    Value::from(centroid.1 as f64), // lon
                    Value::from(centroid.0 as f64), // lat
                ]),
            );

            let iter: Box<dyn Iterator<Item = usize>> = if let Some(idx) = field_indices.as_ref() {
                Box::new(idx.iter().copied())
            } else {
                Box::new(0..num_features)
            };

            for feat_index in iter {
                if aggregation.feat_counts[feat_index] > 0 {
                    let avg =
                        aggregation.sums[feat_index] / aggregation.feat_counts[feat_index] as f64;
                    if let (Some(min_num), Some(max_num), Some(avg_num)) = (
                        serde_json::Number::from_f64(aggregation.mins[feat_index] as f64),
                        serde_json::Number::from_f64(aggregation.maxs[feat_index] as f64),
                        serde_json::Number::from_f64(avg),
                    ) {
                        props.insert(min_keys[feat_index].clone(), Value::Number(min_num));
                        props.insert(max_keys[feat_index].clone(), Value::Number(max_num));
                        props.insert(avg_keys[feat_index].clone(), Value::Number(avg_num));
                    }
                }
            }

            // Add travel time aggregation fields
            if let Some(tt_aggs) = travel_aggs.get(&pc_idx) {
                for (ti, agg) in tt_aggs.iter().enumerate() {
                    if agg.count > 0 {
                        let key = &travel_field_keys[ti];
                        let avg = agg.sum / agg.count as f64;
                        if let Some(nm) = serde_json::Number::from_f64(agg.min as f64) {
                            props.insert(format!("min_{key}"), Value::Number(nm));
                        }
                        if let Some(nm) = serde_json::Number::from_f64(agg.max as f64) {
                            props.insert(format!("max_{key}"), Value::Number(nm));
                        }
                        if let Some(nm) = serde_json::Number::from_f64(avg) {
                            props.insert(format!("avg_{key}"), Value::Number(nm));
                        }
                    }
                }
            }

            // Add enum distribution array (for pie chart visualization)
            if let (Some(ref key), Some(ref ed)) = (&enum_dist_key, &aggregation.enum_dist) {
                let arr: Vec<Value> = ed.counts.iter().map(|&c| Value::from(c)).collect();
                props.insert(key.clone(), Value::Array(arr));
            }

            // Build GeoJSON Feature
            let mut feature = Map::new();
            feature.insert("type".into(), Value::String("Feature".into()));
            feature.insert("geometry".into(), geometry);
            feature.insert("properties".into(), Value::Object(props));

            features.push(feature);

            if features.len() >= MAX_CELLS_PER_REQUEST {
                break;
            }
        }

        histogram!("postcodes_response_count").record(features.len() as f64);

        let truncated = features.len() >= MAX_CELLS_PER_REQUEST;
        let t_total = t0.elapsed();
        info!(
            postcodes_before_filter,
            postcodes_after_filter = features.len(),
            filtered_out,
            truncated,
            bounds = format_args!("{:.6},{:.6},{:.6},{:.6}", south, west, north, east),
            filters = num_filters,
            filters_raw = filters_str.as_deref().unwrap_or("-"),
            fields = field_indices.as_ref().map(|v| v.len() as i32).unwrap_or(-1),
            travel_entries = travel_entries.len(),
            agg_ms = format_args!("{:.1}", t_agg.as_secs_f64() * 1000.0),
            json_ms = format_args!("{:.1}", (t_total - t_agg).as_secs_f64() * 1000.0),
            total_ms = format_args!("{:.1}", t_total.as_secs_f64() * 1000.0),
            "GET /api/postcodes"
        );

        Ok(PostcodesResponse {
            r#type: "FeatureCollection",
            features,
        })
    })
    .await
    .map_err(|error| (StatusCode::INTERNAL_SERVER_ERROR, error.to_string()).into_response())?
    .map_err(|error| (StatusCode::INTERNAL_SERVER_ERROR, error).into_response())?;

    Ok(Json(response))
}

/// Find the nearest postcode to a given lat/lng coordinate.
/// If the user is authenticated, logs their location to PocketBase in the background.
pub async fn get_nearest_postcode(
    State(shared): State<Arc<SharedState>>,
    Extension(user): Extension<OptionalUser>,
    Query(params): Query<NearestPostcodeParams>,
) -> Result<Json<Value>, StatusCode> {
    let state = shared.load_state();
    let postcode_data = &state.postcode_data;

    let query_lat = params.lat as f32;
    let query_lng = params.lng as f32;
    let cos_lat = (query_lat as f64).to_radians().cos() as f32;

    let mut best_idx: Option<usize> = None;
    let mut best_dist_sq = f32::MAX;

    for (idx, &(pc_lat, pc_lon)) in postcode_data.centroids.iter().enumerate() {
        let dlat = pc_lat - query_lat;
        let dlon = (pc_lon - query_lng) * cos_lat;
        let dist_sq = dlat * dlat + dlon * dlon;
        if dist_sq < best_dist_sq {
            best_dist_sq = dist_sq;
            best_idx = Some(idx);
        }
    }

    let idx = best_idx.ok_or(StatusCode::NOT_FOUND)?;
    let (lat, lon) = postcode_data.centroids[idx];
    let geometry = postcode_data.geometries[idx].clone();
    let postcode = &postcode_data.postcodes[idx];

    // Log location for authenticated users (best-effort, non-blocking)
    if let Some(ref pb_user) = user.0 {
        let state = state.clone();
        let user_id = pb_user.id.clone();
        let lat_f64 = params.lat;
        let lng_f64 = params.lng;
        let pc = postcode.clone();
        tokio::spawn(async move {
            log_user_location(&state, &user_id, lat_f64, lng_f64, &pc).await;
        });
    }

    info!(postcode = %postcode, "GET /api/nearest-postcode");
    Ok(Json(serde_json::json!({
        "postcode": postcode,
        "latitude": lat as f64,
        "longitude": lon as f64,
        "geometry": geometry,
    })))
}

/// Look up a single postcode and return its centroid coordinates and geometry.
pub async fn get_postcode_lookup(
    State(shared): State<Arc<SharedState>>,
    Path(postcode): Path<String>,
) -> Result<Json<Value>, StatusCode> {
    let state = shared.load_state();
    let normalized = normalize_postcode(&postcode);

    let postcode_data = &state.postcode_data;

    if let Some(&idx) = postcode_data.postcode_to_idx.get(&normalized) {
        let (lat, lon) = postcode_data.centroids[idx];
        let geometry = postcode_data.geometries[idx].clone();

        info!(postcode = %normalized, "GET /api/postcode/{postcode}");
        Ok(Json(serde_json::json!({
            "postcode": normalized,
            "latitude": lat as f64,
            "longitude": lon as f64,
            "geometry": geometry,
        })))
    } else {
        Err(StatusCode::NOT_FOUND)
    }
}