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

use std::sync::Arc;
use std::time::Instant;

use axum::extract::State;
use axum::http::StatusCode;
use axum::response::{IntoResponse, Json, Response};
use serde_json::json;
use tracing::{info, warn};

use crate::consts::GRID_CELL_SIZE;
use crate::data::{self, PropertyData};
use crate::metrics::record_data_stats;
use crate::routes::{build_features_response, build_system_prompt};
use crate::state::{AppState, SharedState};
use crate::utils::GridIndex;

pub async fn post_reload(State(shared): State<Arc<SharedState>>) -> Response {
    if !shared.try_start_reload() {
        return (StatusCode::CONFLICT, "Reload already in progress").into_response();
    }

    info!("Reload triggered — rebuilding property data");
    let start = Instant::now();

    // shared is cloned so we retain a reference after spawn_blocking
    let sh = Arc::clone(&shared);
    let result = tokio::task::spawn_blocking(move || rebuild_data(&sh, start)).await;

    // Always clear the reload flag
    shared.finish_reload();

    match result {
        Ok(Ok((rows, features, elapsed_ms))) => Json(json!({
            "status": "ok",
            "rows": rows,
            "features": features,
            "elapsed_ms": elapsed_ms,
        }))
        .into_response(),
        Ok(Err(err)) => {
            warn!("Reload failed: {err:#}");
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                Json(json!({ "error": format!("{err:#}") })),
            )
                .into_response()
        }
        Err(err) => {
            warn!("Reload task panicked: {err}");
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                Json(json!({ "error": format!("Reload task panicked: {err}") })),
            )
                .into_response()
        }
    }
}

fn rebuild_data(shared: &SharedState, start: Instant) -> anyhow::Result<(usize, usize, u128)> {
    let old = shared.load_state();

    // 1. Load PropertyData from parquet files
    let property_data = PropertyData::load(
        &shared.properties_path,
        &shared.postcode_features_path,
        &shared.listings_buy_path,
        &shared.listings_rent_path,
    )?;
    let row_count = property_data.lat.len();
    let feature_count = property_data.num_features;

    // 2. Build spatial grid index
    info!("Reload: building spatial grid index");
    let grid = GridIndex::build(&property_data.lat, &property_data.lon, GRID_CELL_SIZE);

    // 3. Precompute H3 cells
    info!("Reload: precomputing H3 cells");
    let h3_cells = data::precompute_h3(&property_data.lat, &property_data.lon)?;

    // 4. Build feature lookup tables
    let feature_name_to_index = property_data
        .feature_names
        .iter()
        .enumerate()
        .map(|(idx, name)| (name.clone(), idx))
        .collect();

    let min_keys = property_data
        .feature_names
        .iter()
        .map(|n| format!("min_{n}"))
        .collect();
    let max_keys = property_data
        .feature_names
        .iter()
        .map(|n| format!("max_{n}"))
        .collect();
    let avg_keys = property_data
        .feature_names
        .iter()
        .map(|n| format!("avg_{n}"))
        .collect();

    // 5. Build features response and AI prompt
    let features_response = build_features_response(&property_data);
    let mode_destinations: Vec<(String, usize)> = old
        .travel_time_store
        .available_modes
        .iter()
        .map(|mode| {
            let count = old
                .travel_time_store
                .destinations
                .get(mode.as_str())
                .map(|slugs| slugs.len())
                .unwrap_or(0);
            (mode.clone(), count)
        })
        .filter(|(_, count)| *count > 0)
        .collect();
    let ai_filters_system_prompt = build_system_prompt(&features_response, &mode_destinations);

    // 6. Update data metrics
    record_data_stats(
        row_count,
        old.poi_data.lat.len(),
        old.postcode_data.postcodes.len(),
    );

    // 7. Build new AppState, sharing unchanged fields via Arc
    let new_state = AppState {
        data: property_data,
        grid,
        h3_cells,
        feature_name_to_index,
        min_keys,
        max_keys,
        avg_keys,
        features_response,
        ai_filters_system_prompt,

        // Shared across reloads (Arc clone is cheap)
        poi_data: Arc::clone(&old.poi_data),
        poi_grid: Arc::clone(&old.poi_grid),
        place_data: Arc::clone(&old.place_data),
        postcode_data: Arc::clone(&old.postcode_data),
        outcode_data: Arc::clone(&old.outcode_data),
        poi_category_groups: Arc::clone(&old.poi_category_groups),
        travel_time_store: Arc::clone(&old.travel_time_store),
        token_cache: Arc::clone(&old.token_cache),
        superuser_token_cache: Arc::clone(&old.superuser_token_cache),

        // Config (cheap clone)
        screenshot_url: old.screenshot_url.clone(),
        public_url: old.public_url.clone(),
        is_dev: old.is_dev,
        index_html: old.index_html.clone(),
        http_client: old.http_client.clone(),
        pocketbase_url: old.pocketbase_url.clone(),
        pocketbase_admin_email: old.pocketbase_admin_email.clone(),
        pocketbase_admin_password: old.pocketbase_admin_password.clone(),
        gemini_api_key: old.gemini_api_key.clone(),
        gemini_model: old.gemini_model.clone(),
        google_maps_api_key: old.google_maps_api_key.clone(),
        stripe_secret_key: old.stripe_secret_key.clone(),
        stripe_webhook_secret: old.stripe_webhook_secret.clone(),
        stripe_referral_coupon_id: old.stripe_referral_coupon_id.clone(),
    };

    // 8. Atomic swap
    shared.swap_state(new_state);

    let elapsed = start.elapsed();
    info!(
        rows = row_count,
        features = feature_count,
        elapsed_ms = elapsed.as_millis(),
        "Reload complete"
    );

    Ok((row_count, feature_count, elapsed.as_millis()))
}