use std::collections::{HashMap, HashSet};
use std::path::Path;

use anyhow::{bail, Context};
use polars::frame::DataFrame;
use polars::lazy::frame::LazyFrame;
use polars::prelude::*;
use rustc_hash::FxHashSet;
use serde::Serialize;
use tracing::info;

use crate::features::POI_GROUP_ORDER;
use crate::utils::{generate_priorities, InternedColumn};

#[derive(Serialize, Clone)]
pub struct POICategoryGroup {
    pub name: String,
    pub categories: Vec<String>,
}

const GROCERY_DASHBOARD_CATEGORIES: &[&str] = &[
    "Supermarket",
    "Convenience Store",
    "Bakery",
    "Greengrocer",
    "Aldi",
    "Amazon",
    "Asda",
    "Booths",
    "Budgens",
    "Centra",
    "Co-op",
    "COOK",
    "Costco",
    "Dunnes Stores",
    "Farmfoods",
    "Heron Foods",
    "Iceland",
    "Lidl",
    "Makro",
    "M&S",
    "Morrisons",
    "Planet Organic",
    "Sainsbury's",
    "Spar",
    "Tesco",
    "The Food Warehouse",
    "Waitrose",
    "Whole Foods Market",
];

const DASHBOARD_POI_GROUPS: &[(&str, &[&str])] = &[
    (
        "Public Transport",
        &[
            "Rail station",
            "Tube station",
            "DLR station",
            "Tram & Metro stop",
            "Bus station",
            "Bus stop",
            "Airport",
        ],
    ),
    ("Groceries", GROCERY_DASHBOARD_CATEGORIES),
    ("Food & Drink", &["Café", "Restaurant", "Pub", "Fast Food"]),
    ("Green Space", &["Park", "Playground"]),
    (
        "Education",
        &[
            "Nursery school",
            "Primary school",
            "Secondary school",
            "All-through school",
            "Sixth form",
            "Further education college",
            "University",
            "Special school",
            "School",
        ],
    ),
    (
        "Health",
        &["GP Surgery", "Pharmacy", "Dentist", "Hospital", "Clinic"],
    ),
    (
        "Leisure",
        &[
            "Gym & Fitness",
            "Sports Centre",
            "Cinema",
            "Theatre",
            "Library",
        ],
    ),
    (
        "Practical",
        &["Post Office", "Bank", "EV Charging", "Fuel Station"],
    ),
];

fn add_category_filter_index(
    category_values: &[String],
    category: &str,
    selected: &mut FxHashSet<u16>,
) {
    if let Some(pos) = category_values.iter().position(|value| value == category) {
        selected.insert(pos as u16);
    }
}

fn canonical_poi_category(category: &str) -> &str {
    match category {
        "Allendale Co-operative Society"
        | "Central England Co-operative"
        | "Channel Islands Co-operative Society"
        | "Chelmsford Star Co-operative Society"
        | "Clydebank Co-operative"
        | "Coniston Co-operative Society"
        | "Co-op Food"
        | "East of England Co-operative"
        | "Heart of England Co-operative"
        | "Langdale Co-operative Society"
        | "Lincolnshire Co-operative"
        | "Midcounties Co-operative"
        | "Scottish Midland Co-operative"
        | "Tamworth Co-operative Society"
        | "The Co-operative Food"
        | "The Co-operative Food PFS"
        | "The Co-operative Group"
        | "The Radstock Co-operative Society"
        | "The Southern Co-operative" => "Co-op",
        _ => category,
    }
}

/// Categories the pipeline emits for the GIAS-derived school POIs. A bare
/// `poi=School` URL (predating the per-phase split) is expanded to all of these
/// so bookmarked links keep showing schools.
const SCHOOL_CATEGORY_ALIASES: &[&str] = &[
    "Nursery school",
    "Primary school",
    "Secondary school",
    "All-through school",
    "Sixth form",
    "Further education college",
    "University",
    "Special school",
    "School",
];

pub fn resolve_poi_category_filter(category_values: &[String], categories: &str) -> FxHashSet<u16> {
    let mut selected = FxHashSet::default();
    for part in categories.split(',') {
        let category = canonical_poi_category(part.trim());
        if category.is_empty() {
            continue;
        }
        if category == "School" {
            for alias in SCHOOL_CATEGORY_ALIASES {
                add_category_filter_index(category_values, alias, &mut selected);
            }
            continue;
        }
        add_category_filter_index(category_values, category, &mut selected);
    }
    selected
}

/// Metadata for state-funded school POIs (sourced from the DfE GIAS register).
/// Every field is optional because GIAS does not populate every column for every
/// establishment type (e.g. nurseries have no sixth form, FE colleges no FSM).
#[derive(Serialize, Clone, Default)]
pub struct SchoolMetadata {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub phase: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub r#type: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub type_group: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub age_range: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub gender: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub religious_character: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub admissions_policy: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub nursery_provision: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub sixth_form: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub capacity: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub pupils: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub fsm_percent: Option<f32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub trust: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub address: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub postcode: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub local_authority: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub website: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub telephone: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub head_name: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub ofsted_rating: Option<String>,
}

pub struct POIData {
    /// Contiguous buffer holding all POI ID strings end-to-end.
    id_buffer: String,
    /// Byte offset into `id_buffer` where each row's ID starts.
    id_offsets: Vec<u32>,
    /// Length in bytes of each row's ID.
    id_lengths: Vec<u16>,
    pub group: InternedColumn,
    pub category: InternedColumn,
    pub icon_category: InternedColumn,
    pub name: Vec<String>,
    pub lat: Vec<f32>,
    pub lng: Vec<f32>,
    pub emoji: InternedColumn,
    /// Deterministic pseudo-random priority per row, used to select a spatially
    /// uniform subset when the POI count exceeds the per-request limit.
    /// Computed once at load time so the same POIs are always chosen for a given viewport.
    pub priority: Vec<u32>,
    /// Indirection table: row idx → index into `school_meta`, or u32::MAX when
    /// the POI is not a school. Keeps the per-row overhead at 4 bytes regardless
    /// of how many school metadata fields we carry.
    school_meta_idx: Vec<u32>,
    school_meta: Vec<SchoolMetadata>,
}

impl POIData {
    /// Get the ID string for a given row.
    pub fn id(&self, row: usize) -> &str {
        let offset = self.id_offsets[row] as usize;
        let length = self.id_lengths[row] as usize;
        &self.id_buffer[offset..offset + length]
    }

    /// Get the school metadata for a given row, or None if not a school.
    pub fn school(&self, row: usize) -> Option<&SchoolMetadata> {
        let idx = self.school_meta_idx[row];
        if idx == u32::MAX {
            None
        } else {
            Some(&self.school_meta[idx as usize])
        }
    }
}

fn extract_str_col(df: &DataFrame, name: &str) -> anyhow::Result<Vec<String>> {
    let column = df
        .column(name)
        .with_context(|| format!("Missing column '{name}' in POI data"))?;
    let string_column = column
        .str()
        .with_context(|| format!("Column '{name}' is not a string column"))?;
    string_column
        .into_iter()
        .enumerate()
        .map(|(row, value)| {
            value
                .map(ToString::to_string)
                .with_context(|| format!("Column '{name}' has null at row {row}"))
        })
        .collect()
}

fn extract_f32_col(df: &DataFrame, name: &str) -> anyhow::Result<Vec<f32>> {
    let column = df
        .column(name)
        .with_context(|| format!("Missing column '{name}' in POI data"))?;
    let cast = column
        .cast(&DataType::Float32)
        .with_context(|| format!("Failed to cast column '{name}' to Float32"))?;
    let float_column = cast
        .f32()
        .with_context(|| format!("Column '{name}' is not a float32 column"))?;
    float_column
        .into_iter()
        .enumerate()
        .map(|(row, value)| value.with_context(|| format!("Column '{name}' has null at row {row}")))
        .collect()
}

/// Read an optional string column. Returns None when the column itself is missing
/// (older POI parquets without the school_* extension); returns Some(vec) of
/// length row_count where each entry is None for null cells.
fn extract_optional_str_col(
    df: &DataFrame,
    name: &str,
) -> anyhow::Result<Option<Vec<Option<String>>>> {
    let column = match df.column(name) {
        Ok(column) => column,
        Err(_) => return Ok(None),
    };
    let string_column = column
        .str()
        .with_context(|| format!("Column '{name}' is not a string column"))?;
    Ok(Some(
        string_column
            .into_iter()
            .map(|value| value.map(ToString::to_string))
            .collect(),
    ))
}

fn extract_optional_u32_col(
    df: &DataFrame,
    name: &str,
) -> anyhow::Result<Option<Vec<Option<u32>>>> {
    let column = match df.column(name) {
        Ok(column) => column,
        Err(_) => return Ok(None),
    };
    let cast = column
        .cast(&DataType::Int64)
        .with_context(|| format!("Failed to cast column '{name}' to Int64"))?;
    let int_column = cast
        .i64()
        .with_context(|| format!("Column '{name}' is not an integer column"))?;
    Ok(Some(
        int_column
            .into_iter()
            .map(|value| value.and_then(|v| if v < 0 { None } else { Some(v as u32) }))
            .collect(),
    ))
}

fn extract_optional_f32_col(
    df: &DataFrame,
    name: &str,
) -> anyhow::Result<Option<Vec<Option<f32>>>> {
    let column = match df.column(name) {
        Ok(column) => column,
        Err(_) => return Ok(None),
    };
    let cast = column
        .cast(&DataType::Float32)
        .with_context(|| format!("Failed to cast column '{name}' to Float32"))?;
    let float_column = cast
        .f32()
        .with_context(|| format!("Column '{name}' is not a float32 column"))?;
    Ok(Some(float_column.into_iter().collect()))
}

fn build_school_meta(
    row_count: usize,
    df: &DataFrame,
) -> anyhow::Result<(Vec<u32>, Vec<SchoolMetadata>)> {
    let phase = extract_optional_str_col(df, "school_phase")?;
    if phase.is_none() {
        // POI parquet predates the school metadata extension — record an empty
        // table and a sentinel-filled index, so callers transparently see None.
        return Ok((vec![u32::MAX; row_count], Vec::new()));
    }

    let phase = phase.unwrap();
    let r#type = extract_optional_str_col(df, "school_type")?.unwrap_or_default();
    let type_group = extract_optional_str_col(df, "school_type_group")?.unwrap_or_default();
    let age_range = extract_optional_str_col(df, "school_age_range")?.unwrap_or_default();
    let gender = extract_optional_str_col(df, "school_gender")?.unwrap_or_default();
    let religious_character =
        extract_optional_str_col(df, "school_religious_character")?.unwrap_or_default();
    let admissions_policy =
        extract_optional_str_col(df, "school_admissions_policy")?.unwrap_or_default();
    let nursery_provision =
        extract_optional_str_col(df, "school_nursery_provision")?.unwrap_or_default();
    let sixth_form = extract_optional_str_col(df, "school_sixth_form")?.unwrap_or_default();
    let capacity = extract_optional_u32_col(df, "school_capacity")?.unwrap_or_default();
    let pupils = extract_optional_u32_col(df, "school_pupils")?.unwrap_or_default();
    let fsm_percent = extract_optional_f32_col(df, "school_fsm_percent")?.unwrap_or_default();
    let trust = extract_optional_str_col(df, "school_trust")?.unwrap_or_default();
    let address = extract_optional_str_col(df, "school_address")?.unwrap_or_default();
    let postcode = extract_optional_str_col(df, "school_postcode")?.unwrap_or_default();
    let local_authority =
        extract_optional_str_col(df, "school_local_authority")?.unwrap_or_default();
    let website = extract_optional_str_col(df, "school_website")?.unwrap_or_default();
    let telephone = extract_optional_str_col(df, "school_telephone")?.unwrap_or_default();
    let head_name = extract_optional_str_col(df, "school_head_name")?.unwrap_or_default();
    let ofsted_rating = extract_optional_str_col(df, "school_ofsted_rating")?.unwrap_or_default();

    let fetch_str = |col: &Vec<Option<String>>, row: usize| -> Option<String> {
        col.get(row).cloned().flatten()
    };
    let fetch_u32 =
        |col: &Vec<Option<u32>>, row: usize| -> Option<u32> { col.get(row).copied().flatten() };
    let fetch_f32 =
        |col: &Vec<Option<f32>>, row: usize| -> Option<f32> { col.get(row).copied().flatten() };

    let mut idx = vec![u32::MAX; row_count];
    let mut meta = Vec::new();
    for (row, meta_idx) in idx.iter_mut().enumerate().take(row_count) {
        let type_group_val = fetch_str(&type_group, row);
        let type_val = fetch_str(&r#type, row);
        // type_group is present for every GIAS row, so use it as the sentinel
        // for "this POI is a school" — matches the pipeline guarantee.
        if type_group_val.is_none() && type_val.is_none() {
            continue;
        }
        *meta_idx = meta.len() as u32;
        meta.push(SchoolMetadata {
            phase: fetch_str(&phase, row),
            r#type: type_val,
            type_group: type_group_val,
            age_range: fetch_str(&age_range, row),
            gender: fetch_str(&gender, row),
            religious_character: fetch_str(&religious_character, row),
            admissions_policy: fetch_str(&admissions_policy, row),
            nursery_provision: fetch_str(&nursery_provision, row),
            sixth_form: fetch_str(&sixth_form, row),
            capacity: fetch_u32(&capacity, row),
            pupils: fetch_u32(&pupils, row),
            fsm_percent: fetch_f32(&fsm_percent, row),
            trust: fetch_str(&trust, row),
            address: fetch_str(&address, row),
            postcode: fetch_str(&postcode, row),
            local_authority: fetch_str(&local_authority, row),
            website: fetch_str(&website, row),
            telephone: fetch_str(&telephone, row),
            head_name: fetch_str(&head_name, row),
            ofsted_rating: fetch_str(&ofsted_rating, row),
        });
    }
    Ok((idx, meta))
}

impl POIData {
    pub fn load(parquet_path: &Path) -> anyhow::Result<Self> {
        super::run_polars_io(|| Self::load_inner(parquet_path))
    }

    fn load_inner(parquet_path: &Path) -> anyhow::Result<Self> {
        info!("Loading POI data from {:?}...", parquet_path);

        let parquet_path = PlRefPath::try_from_path(parquet_path)
            .context("Failed to normalize POI parquet path")?;
        let df = LazyFrame::scan_parquet(parquet_path, Default::default())
            .context("Failed to scan POI parquet")?
            .collect()
            .context("Failed to read POI parquet")?;

        let row_count = df.height();
        info!("Loaded {} POIs", row_count);

        let id_raw: Vec<String> = extract_str_col(&df, "id")?;
        let name = extract_str_col(&df, "name")?;
        let category_raw: Vec<String> = extract_str_col(&df, "category")?
            .into_iter()
            .map(|category| canonical_poi_category(&category).to_string())
            .collect();
        let group_raw = extract_str_col(&df, "group")?;
        let lat = extract_f32_col(&df, "lat")?;
        let lng = extract_f32_col(&df, "lng")?;
        let emoji_raw = extract_str_col(&df, "emoji")?;
        let icon_category_raw: Vec<String> = extract_str_col(&df, "icon_category")?
            .into_iter()
            .map(|category| canonical_poi_category(&category).to_string())
            .collect();

        // Pack POI IDs into a contiguous buffer
        let total_id_bytes: usize = id_raw.iter().map(|s| s.len()).sum();
        let mut id_buffer = String::with_capacity(total_id_bytes);
        let mut id_offsets = Vec::with_capacity(row_count);
        let mut id_lengths = Vec::with_capacity(row_count);
        for s in &id_raw {
            let offset = id_buffer.len() as u32;
            let length = s.len().min(u16::MAX as usize) as u16;
            id_offsets.push(offset);
            id_lengths.push(length);
            id_buffer.push_str(&s[..length as usize]);
        }

        let category = InternedColumn::build(&category_raw);
        let icon_category = InternedColumn::build(&icon_category_raw);
        let group = InternedColumn::build(&group_raw);
        let emoji = InternedColumn::build(&emoji_raw);

        info!(
            category_unique = category.values.len(),
            icon_category_unique = icon_category.values.len(),
            group_unique = group.values.len(),
            emoji_unique = emoji.values.len(),
            "POI string columns interned"
        );

        // Assign a deterministic pseudo-random priority to each row.
        // This ensures the same POIs are selected across requests,
        // preventing visual "shuffling" when panning the map.
        let priority = generate_priorities(row_count);

        let (school_meta_idx, school_meta) = build_school_meta(row_count, &df)?;
        info!(schools = school_meta.len(), "Loaded GIAS school metadata");

        info!("POI data loading complete.");

        Ok(POIData {
            id_buffer,
            id_offsets,
            id_lengths,
            name,
            category,
            icon_category,
            group,
            lat,
            lng,
            emoji,
            priority,
            school_meta_idx,
            school_meta,
        })
    }

    /// Build dashboard category groups from every category present in the loaded POI data.
    pub fn category_groups(&self) -> anyhow::Result<Vec<POICategoryGroup>> {
        let mut group_cats: HashMap<String, HashSet<String>> = HashMap::new();
        let num_pois = self.category.indices.len();
        for row in 0..num_pois {
            let category = self.category.get(row).to_string();
            let group = self.group.get(row).to_string();
            group_cats.entry(group).or_default().insert(category);
        }

        // Validate that data groups match the hardcoded order exactly
        let expected: HashSet<&str> = POI_GROUP_ORDER.iter().copied().collect();
        let actual: 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
            );
        }

        let preferred_order: HashMap<&str, HashMap<&str, usize>> = DASHBOARD_POI_GROUPS
            .iter()
            .map(|(group, categories)| {
                (
                    *group,
                    categories
                        .iter()
                        .enumerate()
                        .map(|(idx, category)| (*category, idx))
                        .collect(),
                )
            })
            .collect();

        let groups: Vec<POICategoryGroup> = POI_GROUP_ORDER
            .iter()
            .filter_map(|group_name| {
                let mut categories: Vec<String> = group_cats
                    .get(*group_name)
                    .map(|categories| categories.iter().cloned().collect())
                    .unwrap_or_default();
                if categories.is_empty() {
                    return None;
                }
                let group_order = preferred_order.get(*group_name);
                categories.sort_by(|a, b| {
                    let a_order = group_order.and_then(|order| order.get(a.as_str())).copied();
                    let b_order = group_order.and_then(|order| order.get(b.as_str())).copied();
                    match (a_order, b_order) {
                        (Some(left), Some(right)) => left.cmp(&right),
                        (Some(_), None) => std::cmp::Ordering::Less,
                        (None, Some(_)) => std::cmp::Ordering::Greater,
                        (None, None) => a.cmp(b),
                    }
                });
                Some(POICategoryGroup {
                    name: (*group_name).to_string(),
                    categories,
                })
            })
            .collect();

        Ok(groups)
    }
}

#[cfg(test)]
impl POIData {
    /// Minimal empty instance for integration tests that need an `AppState`
    /// but never touch POI data.
    pub(crate) fn empty_for_tests() -> Self {
        POIData {
            id_buffer: String::new(),
            id_offsets: Vec::new(),
            id_lengths: Vec::new(),
            group: InternedColumn::build(&[]),
            category: InternedColumn::build(&[]),
            icon_category: InternedColumn::build(&[]),
            name: Vec::new(),
            lat: Vec::new(),
            lng: Vec::new(),
            emoji: InternedColumn::build(&[]),
            priority: Vec::new(),
            school_meta_idx: Vec::new(),
            school_meta: Vec::new(),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn category_filter_matches_exact_present_categories() {
        let values = vec![
            "Supermarket".to_string(),
            "Tesco".to_string(),
            "Aldi".to_string(),
            "Rail station".to_string(),
        ];

        let selected = resolve_poi_category_filter(&values, "Supermarket,Rail station");

        assert!(selected.contains(&0));
        assert!(selected.contains(&3));
        assert_eq!(selected.len(), 2);
    }

    #[test]
    fn unknown_category_filter_matches_nothing() {
        let values = vec!["Supermarket".to_string()];

        let selected = resolve_poi_category_filter(&values, "Unknown");

        assert!(selected.is_empty());
    }

    #[test]
    fn legacy_school_filter_expands_to_all_school_categories() {
        // Bookmarked URLs from before the per-phase split sent `poi=School`;
        // they should still match every school category that's loaded.
        let values = vec![
            "Primary school".to_string(),
            "Secondary school".to_string(),
            "University".to_string(),
            "Tesco".to_string(),
        ];

        let selected = resolve_poi_category_filter(&values, "School");

        assert!(selected.contains(&0));
        assert!(selected.contains(&1));
        assert!(selected.contains(&2));
        assert!(!selected.contains(&3));
        assert_eq!(selected.len(), 3);
    }

    #[test]
    fn coop_category_aliases_resolve_to_single_category() {
        let values = vec!["Co-op".to_string(), "Tesco".to_string()];

        let selected = resolve_poi_category_filter(
            &values,
            "Central England Co-operative,The Southern Co-operative",
        );

        assert!(selected.contains(&0));
        assert_eq!(selected.len(), 1);
        assert_eq!(canonical_poi_category("Lincolnshire Co-operative"), "Co-op");
        assert_eq!(canonical_poi_category("Tesco"), "Tesco");
    }
}