mod consts;
mod data;
mod features;
mod filter;
mod grid_index;
mod og_middleware;
mod routes;
mod state;
#[cfg(test)]
mod tests;

use std::path::PathBuf;
use std::sync::Arc;

use anyhow::{bail, Context};
use axum::middleware;
use axum::routing::get;
use axum::Router;
use clap::Parser;
use tower_http::compression::CompressionLayer;
use tower_http::cors::{Any, CorsLayer};
use tower_http::services::ServeDir;
use tower_http::trace::TraceLayer;
use tracing::info;
use tracing_subscriber::EnvFilter;

use state::AppState;

#[derive(Parser)]
#[command(name = "narrowit", about = "Narrowit property map server")]
struct Cli {
    /// Path to the wide property parquet file
    #[arg(long)]
    data: PathBuf,

    /// Path to the POI parquet file
    #[arg(long)]
    pois: PathBuf,

    /// Path to the frontend dist directory
    #[arg(long)]
    dist: Option<PathBuf>,

    /// URL of the OG screenshot sidecar (e.g. http://og-screenshot:8002)
    #[arg(long, env = "OG_SIDECAR_URL")]
    og_sidecar_url: Option<String>,

    /// Public-facing URL for absolute og:image URLs
    #[arg(
        long,
        env = "PUBLIC_URL",
        default_value = "https://narrowit.schmelczer.dev"
    )]
    public_url: String,
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    tracing_subscriber::fmt()
        .with_env_filter(
            EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info")),
        )
        .with_ansi(true)
        .init();

    let cli = Cli::parse();

    let parquet_path = &cli.data;
    if !parquet_path.exists() {
        bail!(
            "Property parquet file not found: {}",
            parquet_path.display()
        );
    }

    info!("Loading property data from {}", parquet_path.display());
    let property_data = data::PropertyData::load(parquet_path)?;
    info!(
        rows = property_data.lat.len(),
        features = property_data.num_features,
        enums = property_data.enum_features.len(),
        "Property data loaded"
    );

    info!("Building spatial grid index (0.01° cells)");
    let grid = grid_index::GridIndex::build(
        &property_data.lat,
        &property_data.lon,
        consts::GRID_CELL_SIZE,
    );

    info!(
        "Precomputing H3 cells at resolution {}",
        consts::H3_PRECOMPUTE_MAX
    );
    let h3_cells = data::precompute_h3(&property_data.lat, &property_data.lon)?;

    let poi_path = cli.pois;

    if !poi_path.exists() {
        bail!("POI parquet file not found: {}", poi_path.display());
    }

    info!("Loading POI data from {}", poi_path.display());
    let poi_data = data::POIData::load(&poi_path)?;
    info!(pois = poi_data.lat.len(), "POI data loaded");

    info!("Building POI spatial grid index");
    let poi_grid =
        grid_index::GridIndex::build(&poi_data.lat, &poi_data.lng, consts::GRID_CELL_SIZE);

    let min_keys: Vec<String> = property_data
        .feature_names
        .iter()
        .map(|name| format!("min_{}", name))
        .collect();
    let max_keys: Vec<String> = property_data
        .feature_names
        .iter()
        .map(|name| format!("max_{}", name))
        .collect();
    let enum_min_keys: Vec<String> = property_data
        .enum_features
        .iter()
        .map(|enum_feature| format!("min_{}", enum_feature.name))
        .collect();
    let enum_max_keys: Vec<String> = property_data
        .enum_features
        .iter()
        .map(|enum_feature| format!("max_{}", enum_feature.name))
        .collect();

    // Precompute POI category groups
    let poi_category_groups = {
        let mut group_cats: std::collections::HashMap<String, std::collections::HashSet<String>> =
            std::collections::HashMap::new();
        let num_pois = poi_data.category.indices.len();
        for row in 0..num_pois {
            let category = poi_data.category.get(row).to_string();
            let group = poi_data.group.get(row).to_string();
            group_cats.entry(group).or_default().insert(category);
        }
        // Validate that data groups match the hardcoded order exactly
        let expected: std::collections::HashSet<&str> =
            consts::POI_GROUP_ORDER.iter().copied().collect();
        let actual: std::collections::HashSet<&str> =
            group_cats.keys().map(|key| key.as_str()).collect();
        let missing_from_data: Vec<&&str> = expected.difference(&actual).collect();
        let missing_from_order: Vec<&&str> = actual.difference(&expected).collect();
        if !missing_from_data.is_empty() || !missing_from_order.is_empty() {
            bail!(
                "POI group mismatch!\n  In POI_GROUP_ORDER but not in data: {:?}\n  In data but not in POI_GROUP_ORDER: {:?}",
                missing_from_data, missing_from_order
            );
        }
        consts::POI_GROUP_ORDER
            .iter()
            .map(|group_name| group_name.to_string())
            .collect::<Vec<_>>()
            .into_iter()
            .map(|name| {
                let mut categories: Vec<String> = group_cats
                    .remove(&name)
                    .context("POI group validated but missing from map")?
                    .into_iter()
                    .collect();
                categories.sort();
                Ok(state::POICategoryGroup { name, categories })
            })
            .collect::<anyhow::Result<Vec<_>>>()?
    };

    // Precompute enum name → index map
    let enum_name_to_idx: rustc_hash::FxHashMap<String, usize> = property_data
        .enum_features
        .iter()
        .enumerate()
        .map(|(index, enum_feature)| (enum_feature.name.clone(), index))
        .collect();

    // Read index.html at startup for crawler OG injection
    let frontend_dist = cli.dist.unwrap_or_else(|| {
        if let Ok(executable) = std::env::current_exe() {
            let executable_dir = executable
                .parent()
                .unwrap_or_else(|| std::path::Path::new("."));
            let dist_next_to_binary = executable_dir.join("dist");
            if dist_next_to_binary.exists() {
                return dist_next_to_binary;
            }
        }
        PathBuf::from("frontend/dist")
    });

    let index_html = if frontend_dist.exists() {
        let index_path = frontend_dist.join("index.html");
        match std::fs::read_to_string(&index_path) {
            Ok(html) => {
                info!("Loaded index.html for OG injection");
                Some(html)
            }
            Err(err) => {
                tracing::warn!("Could not read index.html: {}", err);
                None
            }
        }
    } else {
        None
    };

    let http_client = reqwest::Client::new();

    if cli.og_sidecar_url.is_some() {
        info!(
            "OG sidecar configured: {}",
            cli.og_sidecar_url.as_deref().unwrap()
        );
    }

    let state = Arc::new(AppState {
        data: property_data,
        grid,
        h3_cells,
        poi_data,
        poi_grid,
        min_keys,
        max_keys,
        enum_min_keys,
        enum_max_keys,
        poi_category_groups,
        enum_name_to_idx,
        og_sidecar_url: cli.og_sidecar_url,
        public_url: cli.public_url,
        index_html,
        http_client,
    });

    let cors = CorsLayer::new()
        .allow_origin(Any)
        .allow_methods(Any)
        .allow_headers(Any);

    let state_features = state.clone();
    let state_hexagons = state.clone();
    let state_pois = state.clone();
    let state_poi_categories = state.clone();
    let state_hexagon_properties = state.clone();
    let state_hexagon_stats = state.clone();
    let state_og_image = state.clone();
    let state_crawler = state.clone();

    let api = Router::new()
        .route(
            "/api/features",
            get(move || routes::get_features(state_features.clone())),
        )
        .route(
            "/api/hexagons",
            get(move |query| routes::get_hexagons(state_hexagons.clone(), query)),
        )
        .route(
            "/api/pois",
            get(move |query| routes::get_pois(state_pois.clone(), query)),
        )
        .route(
            "/api/poi-categories",
            get(move || routes::get_poi_categories(state_poi_categories.clone())),
        )
        .route(
            "/api/hexagon-properties",
            get(move |query| {
                routes::get_hexagon_properties(state_hexagon_properties.clone(), query)
            }),
        )
        .route(
            "/api/hexagon-stats",
            get(move |query| routes::get_hexagon_stats(state_hexagon_stats.clone(), query)),
        )
        .route(
            "/api/og-image",
            get(move |query| routes::get_og_image(state_og_image.clone(), query)),
        );

    let app = if frontend_dist.exists() {
        api.fallback_service(ServeDir::new(&frontend_dist))
    } else {
        api
    };

    let app = app
        .layer(middleware::from_fn(
            move |req: axum::extract::Request, next: middleware::Next| {
                let st = state_crawler.clone();
                async move {
                    // Inject state into request extensions for the middleware
                    let (mut parts, body) = req.into_parts();
                    parts.extensions.insert(st);
                    let req = axum::extract::Request::from_parts(parts, body);
                    og_middleware::og_middleware(req, next).await
                }
            },
        ))
        .layer(cors)
        .layer(CompressionLayer::new().zstd(true).gzip(true))
        .layer(TraceLayer::new_for_http());

    let addr = consts::SERVER_ADDRESS;
    let listener = tokio::net::TcpListener::bind(addr)
        .await
        .with_context(|| format!("Failed to bind to {addr}"))?;
    info!("Server listening on {}", addr);
    axum::serve(listener, app).await.context("Server error")?;
    Ok(())
}
