//! Per-unit-postcode usual-resident headcounts (ONS Census 2021, table P001),
//! loaded from a side parquet and shown in the right pane. This is display-only
//! area data: it is never a filterable attribute and never enters the feature
//! matrix, mirroring the crime-by-year side table.

use std::path::Path;

use anyhow::{bail, Context};
use polars::prelude::PlRefPath;
use polars::prelude::*;
use rustc_hash::FxHashMap;
use tracing::info;

use crate::utils::normalize_postcode;

use super::run_polars_io;

pub struct PostcodePopulation {
    /// Canonical spaced postcode (e.g. "AL1 1AG") → usual residents (Census 2021).
    by_postcode: FxHashMap<String, u32>,
}

impl PostcodePopulation {
    /// Empty table — used in tests and when no --population-path is supplied.
    pub fn empty() -> Self {
        Self {
            by_postcode: FxHashMap::default(),
        }
    }

    pub fn load(path: &Path) -> anyhow::Result<Self> {
        run_polars_io(|| Self::load_inner(path))
    }

    fn load_inner(path: &Path) -> anyhow::Result<Self> {
        info!("Loading postcode population from {}", path.display());
        let pl_path = PlRefPath::try_from_path(path).with_context(|| {
            format!(
                "Failed to normalize population parquet path {}",
                path.display()
            )
        })?;
        let df = LazyFrame::scan_parquet(pl_path, Default::default())
            .with_context(|| format!("Failed to scan population parquet at {}", path.display()))?
            .collect()
            .with_context(|| format!("Failed to read population parquet at {}", path.display()))?;

        let postcode_col = df
            .column("postcode")
            .context("population parquet missing 'postcode' column")?
            .str()
            .context("'postcode' column is not a string")?;
        // Accept whatever integer width the parquet writer used.
        let population_cast = df
            .column("population")
            .context("population parquet missing 'population' column")?
            .cast(&DataType::Int64)
            .context("'population' column is not an integer")?;
        let population_col = population_cast
            .i64()
            .context("failed to read 'population' as i64")?;

        let mut by_postcode: FxHashMap<String, u32> = FxHashMap::default();
        by_postcode.reserve(df.height());
        for (postcode, population) in postcode_col.into_iter().zip(population_col) {
            let (Some(postcode), Some(population)) = (postcode, population) else {
                continue;
            };
            let trimmed = postcode.trim();
            if trimmed.is_empty() || population <= 0 {
                continue;
            }
            // Normalize to the exact canonical form the routes look up with, so
            // a stray double-space or lowercase in the source can't miss.
            by_postcode.insert(
                normalize_postcode(trimmed),
                population.min(u32::MAX as i64) as u32,
            );
        }

        if by_postcode.is_empty() {
            bail!("population parquet at {} produced no rows", path.display());
        }

        info!(
            postcodes = by_postcode.len(),
            "Postcode population loaded"
        );
        Ok(Self { by_postcode })
    }

    /// Usual-resident count for a single canonical (spaced, upper-case) postcode.
    pub fn for_postcode(&self, postcode: &str) -> Option<u32> {
        self.by_postcode.get(postcode).copied()
    }
}

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

    /// Integration smoke test against the real Census 2021 parquet. Skips when
    /// the data file is absent (CI without a data build) so it never blocks.
    #[test]
    fn loads_real_census_parquet_if_present() {
        let path = std::path::Path::new("../property-data/population_by_postcode.parquet");
        if !path.exists() {
            eprintln!("skipping: {} not present", path.display());
            return;
        }
        let pop = PostcodePopulation::load(path).expect("population parquet should load");

        // Covers the whole of England & Wales (~1.37M unit postcodes).
        assert!(pop.by_postcode.len() > 1_000_000);
        // A residential postcode has a positive headcount...
        assert!(pop.for_postcode("AL1 1AG").is_some_and(|n| n > 0));
        // ...reachable from non-canonical input via normalize-on-load.
        assert_eq!(
            pop.for_postcode(&normalize_postcode("al1  1ag")),
            pop.for_postcode("AL1 1AG"),
        );
        // Postcodes with zero usual residents are absent from P001.
        assert_eq!(pop.for_postcode("EC1A 1BB"), None);
    }
}