use std::fmt::{self, Write};
use std::sync::Arc;

use axum::extract::Query;
use axum::http::StatusCode;
use axum::response::IntoResponse;
use rustc_hash::FxHashMap;
use serde::Deserialize;
use tracing::{info, warn};

use crate::consts::{
    BOUNDS_BUFFER_PERCENT, BOUNDS_QUANTIZATION, ENUM_NULL, H3_PRECOMPUTE_MAX, H3_REQUEST_MAX,
    H3_REQUEST_MIN, POSTCODE_MIN_RESOLUTION,
};
use crate::filter::parse_filters;
use crate::state::AppState;

use super::parse::parse_bounds;

struct HumanBytes(usize);

impl fmt::Display for HumanBytes {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        let bytes = self.0;
        if bytes >= 1_000_000 {
            write!(formatter, "{:.1} MB", bytes as f64 / 1_000_000.0)
        } else if bytes >= 1_000 {
            write!(formatter, "{:.1} KB", bytes as f64 / 1_000.0)
        } else {
            write!(formatter, "{} B", bytes)
        }
    }
}

#[derive(Deserialize)]
pub struct HexagonParams {
    resolution: u8,
    bounds: Option<String>,
    /// Comma-separated filters: `name:min:max,...`
    /// Rows must have non-NaN values within [min,max] for each filter.
    filters: Option<String>,
    /// Comma-separated feature names to include in min/max aggregation.
    /// When present (even if empty), only listed features are aggregated and written.
    /// When absent, all features are included (backward compatible).
    fields: Option<String>,
}

/// Per-cell accumulator for aggregating features.
/// Uses Box<[T]> instead of Vec<T> to avoid storing capacity (saves 8 bytes per field per cell).
struct CellAgg {
    count: u32,
    mins: Box<[f32]>,
    maxs: Box<[f32]>,
    /// Min/max ordinal indices for enum features (255 = no data yet)
    enum_mins: Box<[u8]>,
    enum_maxs: Box<[u8]>,
    /// Most common postcode in this cell (only tracked at high resolutions)
    postcode: Option<String>,
    postcode_count: u32,
    lat_sum: f64,
    lon_sum: f64,
}

impl CellAgg {
    fn new(num_features: usize, num_enums: usize) -> Self {
        CellAgg {
            count: 0,
            mins: vec![f32::INFINITY; num_features].into_boxed_slice(),
            maxs: vec![f32::NEG_INFINITY; num_features].into_boxed_slice(),
            enum_mins: vec![ENUM_NULL; num_enums].into_boxed_slice(),
            enum_maxs: vec![0; num_enums].into_boxed_slice(),
            postcode: None,
            postcode_count: 0,
            lat_sum: 0.0,
            lon_sum: 0.0,
        }
    }

    /// Add a row using row-major feature_data layout.
    /// feature_data[row * num_features + feat_idx] — all features for one row
    /// are contiguous, so this reads a single cache line per ~8 features.
    #[inline]
    fn add_row(&mut self, feature_data: &[f32], row: usize, num_features: usize) {
        self.count += 1;
        let base = row * num_features;
        let row_slice = &feature_data[base..base + num_features];
        for (feat_index, &value) in row_slice.iter().enumerate() {
            if value.is_finite() {
                if value < self.mins[feat_index] {
                    self.mins[feat_index] = value;
                }
                if value > self.maxs[feat_index] {
                    self.maxs[feat_index] = value;
                }
            }
        }
    }

    /// Track min/max ordinal index for each enum feature in this cell.
    #[inline]
    fn add_enums(&mut self, enum_data: &[u8], row: usize, num_enums: usize) {
        let base = row * num_enums;
        let row_slice = &enum_data[base..base + num_enums];
        for (enum_index, &value) in row_slice.iter().enumerate() {
            if value != ENUM_NULL {
                if self.enum_mins[enum_index] == ENUM_NULL || value < self.enum_mins[enum_index] {
                    self.enum_mins[enum_index] = value;
                }
                if value > self.enum_maxs[enum_index] {
                    self.enum_maxs[enum_index] = value;
                }
            }
        }
    }

    /// Add a row, only aggregating the features at the given indices.
    #[inline]
    fn add_row_selective(
        &mut self,
        feature_data: &[f32],
        row: usize,
        num_features: usize,
        indices: &[usize],
    ) {
        self.count += 1;
        let base = row * num_features;
        for &feat_index in indices {
            let value = feature_data[base + feat_index];
            if value.is_finite() {
                if value < self.mins[feat_index] {
                    self.mins[feat_index] = value;
                }
                if value > self.maxs[feat_index] {
                    self.maxs[feat_index] = value;
                }
            }
        }
    }

    /// Track min/max ordinal index for selected enum features only.
    #[inline]
    fn add_enums_selective(
        &mut self,
        enum_data: &[u8],
        row: usize,
        num_enums: usize,
        indices: &[usize],
    ) {
        let base = row * num_enums;
        for &enum_index in indices {
            let value = enum_data[base + enum_index];
            if value != ENUM_NULL {
                if self.enum_mins[enum_index] == ENUM_NULL || value < self.enum_mins[enum_index] {
                    self.enum_mins[enum_index] = value;
                }
                if value > self.enum_maxs[enum_index] {
                    self.enum_maxs[enum_index] = value;
                }
            }
        }
    }

    /// Track postcode and centroid for high-resolution cells.
    /// Uses simple "first seen" approach — at res 11/12, most rows in a cell share a postcode.
    #[inline]
    fn add_postcode(&mut self, postcode: &str, lat: f32, lon: f32) {
        self.lat_sum += lat as f64;
        self.lon_sum += lon as f64;
        if postcode.is_empty() {
            return;
        }
        if self.postcode.is_none() {
            self.postcode = Some(postcode.to_string());
            self.postcode_count = 1;
        } else if self.postcode.as_deref() == Some(postcode) {
            self.postcode_count += 1;
        }
    }
}

/// Escape a string for inclusion in a JSON string literal.
pub(crate) fn write_json_escaped(buf: &mut String, text: &str) {
    for character in text.chars() {
        match character {
            '"' => buf.push_str("\\\""),
            '\\' => buf.push_str("\\\\"),
            '\n' => buf.push_str("\\n"),
            '\r' => buf.push_str("\\r"),
            '\t' => buf.push_str("\\t"),
            ctrl if ctrl < '\x20' => {
                let _ = write!(buf, "\\u{:04x}", ctrl as u32);
            }
            other => buf.push(other),
        }
    }
}

/// Write the hexagons JSON response directly to a String buffer,
/// avoiding serde_json::Value allocations entirely.
/// When `numeric_indices` / `enum_indices` are Some, only those features are written.
#[allow(clippy::too_many_arguments)]
fn write_hexagons_json(
    buf: &mut String,
    groups: &FxHashMap<u64, CellAgg>,
    min_keys: &[String],
    max_keys: &[String],
    num_features: usize,
    enum_min_keys: &[String],
    enum_max_keys: &[String],
    num_enums: usize,
    include_postcode: bool,
    numeric_indices: Option<&[usize]>,
    enum_indices: Option<&[usize]>,
) {
    buf.push_str("{\"features\":[");
    let mut first = true;
    for (&cell_id, aggregation) in groups {
        let Some(cell) = h3o::CellIndex::try_from(cell_id).ok() else {
            continue;
        };

        if !first {
            buf.push(',');
        }
        first = false;

        let _ = write!(buf, "{{\"h3\":\"{}\",\"count\":{}", cell, aggregation.count);

        if let Some(indices) = numeric_indices {
            for &feat_index in indices {
                if aggregation.mins[feat_index].is_finite()
                    && aggregation.maxs[feat_index].is_finite()
                {
                    let _ = write!(
                        buf,
                        ",\"{}\":{},\"{}\":{}",
                        min_keys[feat_index],
                        aggregation.mins[feat_index],
                        max_keys[feat_index],
                        aggregation.maxs[feat_index]
                    );
                }
            }
        } else {
            for feat_index in 0..num_features {
                if aggregation.mins[feat_index].is_finite()
                    && aggregation.maxs[feat_index].is_finite()
                {
                    let _ = write!(
                        buf,
                        ",\"{}\":{},\"{}\":{}",
                        min_keys[feat_index],
                        aggregation.mins[feat_index],
                        max_keys[feat_index],
                        aggregation.maxs[feat_index]
                    );
                }
            }
        }

        if let Some(indices) = enum_indices {
            for &enum_index in indices {
                if aggregation.enum_mins[enum_index] != ENUM_NULL {
                    let _ = write!(
                        buf,
                        ",\"{}\":{},\"{}\":{}",
                        enum_min_keys[enum_index],
                        aggregation.enum_mins[enum_index],
                        enum_max_keys[enum_index],
                        aggregation.enum_maxs[enum_index]
                    );
                }
            }
        } else {
            for enum_index in 0..num_enums {
                if aggregation.enum_mins[enum_index] != ENUM_NULL {
                    let _ = write!(
                        buf,
                        ",\"{}\":{},\"{}\":{}",
                        enum_min_keys[enum_index],
                        aggregation.enum_mins[enum_index],
                        enum_max_keys[enum_index],
                        aggregation.enum_maxs[enum_index]
                    );
                }
            }
        }

        if include_postcode {
            if let Some(ref postcode) = aggregation.postcode {
                let total = aggregation.count as f64;
                let centroid_lat = aggregation.lat_sum / total;
                let centroid_lon = aggregation.lon_sum / total;
                if centroid_lat.is_finite() && centroid_lon.is_finite() {
                    buf.push_str(",\"postcode\":\"");
                    write_json_escaped(buf, postcode);
                    let _ = write!(buf, "\",\"lat\":{},\"lon\":{}", centroid_lat, centroid_lon);
                }
            }
        }

        buf.push('}');
    }
    buf.push_str("]}");
}

pub async fn get_hexagons(
    state: Arc<AppState>,
    Query(params): Query<HexagonParams>,
) -> Result<impl IntoResponse, (StatusCode, String)> {
    let resolution = params.resolution;
    if !(H3_REQUEST_MIN..=H3_REQUEST_MAX).contains(&resolution) {
        warn!(
            resolution,
            "Resolution out of range [{}, {}]", H3_REQUEST_MIN, H3_REQUEST_MAX
        );
        return Err((
            StatusCode::BAD_REQUEST,
            format!(
                "resolution must be between {} and {}",
                H3_REQUEST_MIN, H3_REQUEST_MAX
            ),
        ));
    }

    let bounds_str = params.bounds.ok_or((
        StatusCode::BAD_REQUEST,
        "bounds parameter is required".into(),
    ))?;

    let (mut south, mut west, mut north, mut east) = parse_bounds(&bounds_str)?;

    let lat_range = north - south;
    let lng_range = east - west;
    south -= lat_range * BOUNDS_BUFFER_PERCENT;
    north += lat_range * BOUNDS_BUFFER_PERCENT;
    west -= lng_range * BOUNDS_BUFFER_PERCENT;
    east += lng_range * BOUNDS_BUFFER_PERCENT;

    south = (south / BOUNDS_QUANTIZATION).floor() * BOUNDS_QUANTIZATION;
    west = (west / BOUNDS_QUANTIZATION).floor() * BOUNDS_QUANTIZATION;
    north = (north / BOUNDS_QUANTIZATION).ceil() * BOUNDS_QUANTIZATION;
    east = (east / BOUNDS_QUANTIZATION).ceil() * BOUNDS_QUANTIZATION;

    let filters_str = params.filters.clone();
    let (parsed_filters, parsed_enum_filters) = parse_filters(
        params.filters.as_deref(),
        &state.data.feature_names,
        &state.data.enum_features,
    );
    let num_filters = parsed_filters.len() + parsed_enum_filters.len();

    // Parse optional `fields` param into numeric and enum index sets.
    // If `fields` is absent (None), all features are included.
    // If `fields` is present (even empty string), only listed features are included.
    let field_indices: Option<(Vec<usize>, Vec<usize>)> =
        params.fields.as_ref().map(|fields_str| {
            let mut numeric_indices = Vec::new();
            let mut enum_indices = Vec::new();
            if !fields_str.is_empty() {
                for name in fields_str.split(',') {
                    let name = name.trim();
                    if name.is_empty() {
                        continue;
                    }
                    if let Some(idx) = state
                        .data
                        .feature_names
                        .iter()
                        .position(|feat| feat == name)
                    {
                        numeric_indices.push(idx);
                    } else if let Some(&idx) = state.enum_name_to_idx.get(name) {
                        enum_indices.push(idx);
                    }
                }
            }
            (numeric_indices, enum_indices)
        });

    let json_body = tokio::task::spawn_blocking(move || -> Result<String, String> {
        let t0 = std::time::Instant::now();

        let num_features = state.data.num_features;
        let num_enums = state.data.num_enums;
        let feature_data = &state.data.feature_data;
        let enum_data = &state.data.enum_data;

        let min_keys = &state.min_keys;
        let max_keys = &state.max_keys;
        let enum_min_keys = &state.enum_min_keys;
        let enum_max_keys = &state.enum_max_keys;

        let h3_res = h3o::Resolution::try_from(resolution)
            .map_err(|error| format!("Invalid H3 resolution {}: {}", resolution, error))?;
        let precomputed = &state.h3_cells;
        let need_parent = resolution < H3_PRECOMPUTE_MAX;

        let mut groups: FxHashMap<u64, CellAgg> = FxHashMap::default();

        let include_postcode = resolution >= POSTCODE_MIN_RESOLUTION;

        // Row-level filter check: numeric must be non-NaN and within [min, max],
        // enum must have value index in the allowed set
        let row_passes = |row: usize| -> bool {
            parsed_filters.iter().all(|filter| {
                let value = feature_data[row * num_features + filter.feat_idx];
                value.is_finite() && value >= filter.min && value <= filter.max
            }) && parsed_enum_filters.iter().all(|enum_filter| {
                let value = enum_data[row * num_enums + enum_filter.enum_idx];
                value != ENUM_NULL && enum_filter.allowed.contains(&value)
            })
        };

        // Choose aggregation strategy based on whether fields are specified
        let has_selective = field_indices.is_some();
        let (sel_numeric, sel_enum) = field_indices
            .as_ref()
            .map_or((&[][..], &[][..]), |(ni, ei)| {
                (ni.as_slice(), ei.as_slice())
            });

        let aggregate_row = |groups: &mut FxHashMap<u64, CellAgg>, cell_id: u64, row: usize| {
            let aggregation = groups
                .entry(cell_id)
                .or_insert_with(|| CellAgg::new(num_features, num_enums));
            if has_selective {
                aggregation.add_row_selective(feature_data, row, num_features, sel_numeric);
                aggregation.add_enums_selective(enum_data, row, num_enums, sel_enum);
            } else {
                aggregation.add_row(feature_data, row, num_features);
                aggregation.add_enums(enum_data, row, num_enums);
            }
            if include_postcode {
                aggregation.add_postcode(
                    state.data.postcode(row),
                    state.data.lat[row],
                    state.data.lon[row],
                );
            }
        };

        // Resolve cell at requested resolution from precomputed max-resolution cell.
        // For max resolution, use directly; for lower resolutions, derive parent.
        let cell_for_row = |row: usize| -> u64 {
            let max_cell = precomputed[row];
            if !need_parent || max_cell == 0 {
                return max_cell;
            }
            h3o::CellIndex::try_from(max_cell)
                .ok()
                .and_then(|ci| ci.parent(h3_res))
                .map(u64::from)
                .unwrap_or(0)
        };

        state
            .grid
            .for_each_in_bounds(south, west, north, east, |row_idx| {
                let row = row_idx as usize;
                if !row_passes(row) {
                    return;
                }
                aggregate_row(&mut groups, cell_for_row(row), row);
            });

        let t_agg = t0.elapsed();

        let mut json_buf = String::with_capacity(groups.len() * 128);
        write_hexagons_json(
            &mut json_buf,
            &groups,
            min_keys,
            max_keys,
            num_features,
            enum_min_keys,
            enum_max_keys,
            num_enums,
            include_postcode,
            field_indices.as_ref().map(|(ni, _)| ni.as_slice()),
            field_indices.as_ref().map(|(_, ei)| ei.as_slice()),
        );

        let t_total = t0.elapsed();
        info!(
            resolution,
            cells = groups.len(),
            filters = num_filters,
            filters_raw = filters_str.as_deref().unwrap_or("-"),
            agg_ms = format_args!("{:.1}", t_agg.as_secs_f64() * 1000.0),
            total_ms = format_args!("{:.1}", t_total.as_secs_f64() * 1000.0),
            size = format_args!("{}", HumanBytes(json_buf.len())),
            "GET /api/hexagons"
        );

        Ok(json_buf)
    })
    .await
    .map_err(|error| (StatusCode::INTERNAL_SERVER_ERROR, error.to_string()))?
    .map_err(|error| (StatusCode::INTERNAL_SERVER_ERROR, error))?;

    Ok(([("content-type", "application/json")], json_body))
}
