//! Individual police.uk crime records (last 7 years) backing the right pane's
//! "individual crimes" list and the `/api/crime-records` endpoint.
//!
//! This table is enormous — ~500M rows, because each incident is replicated to
//! every postcode whose buffer covers it (see [`gather`](CrimeRecords::gather)),
//! so it is NOT held as a `Vec<struct>`: each field is a flat columnar
//! [`SpillVec`] (mmap-backed and kernel-reclaimable when `--spill-dir` is set),
//! small string fields are dictionary-encoded, and the parquet is pre-sorted by
//! postcode so each postcode's records are a contiguous `[start, start+count)`
//! slice located via a CSR-style offset index. Resident RSS is ~0 until records
//! are actually read.
//!
//! At ~500M rows the parquet's string columns (postcode/type/location/outcome)
//! decode to tens of GB if read whole, so the loader never materialises the
//! whole `DataFrame`: it streams the file in bounded row-count chunks (only the
//! row groups overlapping each slice are decoded) and writes each column
//! straight into its (optionally spilled) backing store via [`SpillVecBuilder`],
//! keeping the transient footprint to one chunk plus the index maps.

use std::fs::File;
use std::path::Path;

use anyhow::{bail, Context};
use lasso::{Rodeo, RodeoReader, Spur};
use polars::prelude::*;
use rustc_hash::FxHashMap;
use tracing::info;

use super::run_polars_io;
use super::spill::{SpillVec, SpillVecBuilder};

/// Rows decoded per streaming slice. `with_slice` decodes only the row groups
/// overlapping the slice, so the transient decode is roughly one chunk's worth
/// (~tens of MB at the writer's ~123k-row groups) instead of the tens-of-GB
/// whole-file `DataFrame`. The bound's only dependency on file layout is a
/// reasonable input row-group size, which our pipeline writer produces.
const CHUNK_ROWS: usize = 2_000_000;

/// A resolved view of one record (strings dereferenced from the dictionaries).
pub struct CrimeRecordView<'a> {
    /// `year * 12 + (month - 1)`.
    pub month_index: u32,
    pub crime_type: &'a str,
    pub outcome: Option<&'a str>,
    pub location: Option<&'a str>,
    pub lat: f32,
    pub lon: f32,
}

pub struct CrimeRecords {
    month: SpillVec<u32>,
    ctype: SpillVec<u8>,
    outcome: SpillVec<u8>,
    location: SpillVec<Spur>,
    lat: SpillVec<f32>,
    lon: SpillVec<f32>,
    /// Dictionary for `ctype` (bare crime type names, e.g. "Burglary").
    crime_type_dict: Vec<String>,
    /// Dictionary for `outcome`; index 0 is the empty/unknown sentinel.
    outcome_dict: Vec<String>,
    /// Resolver for the interned `location` strings (`""` means withheld).
    location_resolver: RodeoReader,
    /// Postcode → `(start, count)` into the columnar arrays (records for a
    /// postcode are contiguous because the parquet is sorted by postcode).
    by_postcode: FxHashMap<String, (u32, u32)>,
}

impl CrimeRecords {
    #[cfg(test)]
    pub fn empty() -> Self {
        Self {
            month: SpillVec::owned(Vec::new()),
            ctype: SpillVec::owned(Vec::new()),
            outcome: SpillVec::owned(Vec::new()),
            location: SpillVec::owned(Vec::new()),
            lat: SpillVec::owned(Vec::new()),
            lon: SpillVec::owned(Vec::new()),
            crime_type_dict: Vec::new(),
            outcome_dict: vec![String::new()],
            location_resolver: Rodeo::default().into_reader(),
            by_postcode: FxHashMap::default(),
        }
    }

    /// Number of records stored for a postcode (0 if none).
    pub fn total_for(&self, postcode: &str) -> u32 {
        self.by_postcode.get(postcode).map_or(0, |&(_, c)| c)
    }

    /// Resolve a record index to a borrowing view.
    pub fn view(&self, idx: u32) -> CrimeRecordView<'_> {
        let i = idx as usize;
        let outcome_idx = self.outcome[i] as usize;
        let outcome = self
            .outcome_dict
            .get(outcome_idx)
            .filter(|s| !s.is_empty())
            .map(String::as_str);
        let location = self.location_resolver.resolve(&self.location[i]);
        CrimeRecordView {
            month_index: self.month[i],
            crime_type: self
                .crime_type_dict
                .get(self.ctype[i] as usize)
                .map_or("", String::as_str),
            outcome,
            location: (!location.is_empty()).then_some(location),
            lat: self.lat[i],
            lon: self.lon[i],
        }
    }

    /// Record indices across `postcodes`, newest first, optionally restricted to
    /// months `>= since_month`. These are exactly the incidents counted for the
    /// selected postcodes — for a single postcode that is its precise incident
    /// list; for a multi-postcode selection a boundary incident counted for
    /// several postcodes appears once per postcode, matching the count. We do not
    /// de-duplicate because police.uk snaps many genuinely distinct incidents
    /// (especially anti-social behaviour) to the same point/month and provides no
    /// per-incident id to tell a true duplicate from two real incidents apart.
    pub fn gather(&self, postcodes: &[&str], since_month: Option<u32>) -> Vec<u32> {
        let month = self.month.as_slice();
        let mut out: Vec<u32> = Vec::new();
        for pc in postcodes {
            let Some(&(start, count)) = self.by_postcode.get(*pc) else {
                continue;
            };
            for i in start..start + count {
                if since_month.map_or(true, |s| month[i as usize] >= s) {
                    out.push(i);
                }
            }
        }
        out.sort_unstable_by(|&a, &b| month[b as usize].cmp(&month[a as usize]));
        out
    }

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

    fn load_inner(
        path: &Path,
        spill_dir: Option<&Path>,
        chunk_rows: usize,
    ) -> anyhow::Result<Self> {
        // A zero chunk size would loop forever; the public entry point passes the
        // const, but tests parameterise this to exercise chunk boundaries.
        let chunk_rows = chunk_rows.max(1);
        info!("Loading crime records from {}", path.display());

        // Read the footer once for the row count, and keep it to hand to every
        // per-chunk reader so the 2.9MB metadata is never re-parsed.
        let metadata = {
            let file = File::open(path).with_context(|| {
                format!("Failed to open crime-records parquet at {}", path.display())
            })?;
            ParquetReader::new(file)
                .get_metadata()
                .with_context(|| {
                    format!("Failed to read crime-records parquet metadata at {}", path.display())
                })?
                .clone()
        };
        let n = metadata.num_rows;
        // Record indices are stored as `u32` (and `by_postcode` holds `(start,
        // count)` as `u32`), so the table must fit in that index space.
        if n > u32::MAX as usize {
            bail!("crime-records parquet has {n} rows, exceeding the u32 record-index limit");
        }

        // Columns that, when spilling, are written straight into mmap-backed files
        // as we stream — so the ~9GB of columnar data never lands on the heap.
        let mut month = SpillVecBuilder::<u32>::with_len(n, spill_dir, "crime_month")?;
        let mut ctype = SpillVecBuilder::<u8>::with_len(n, spill_dir, "crime_ctype")?;
        let mut outcome = SpillVecBuilder::<u8>::with_len(n, spill_dir, "crime_outcome")?;
        let mut location = SpillVecBuilder::<Spur>::with_len(n, spill_dir, "crime_location")?;
        let mut lat = SpillVecBuilder::<f32>::with_len(n, spill_dir, "crime_lat")?;
        let mut lon = SpillVecBuilder::<f32>::with_len(n, spill_dir, "crime_lon")?;

        let mut crime_type_dict: Vec<String> = Vec::new();
        let mut type_index: FxHashMap<String, u8> = FxHashMap::default();
        // Outcome index 0 is the empty/unknown sentinel.
        let mut outcome_dict: Vec<String> = vec![String::new()];
        let mut outcome_index: FxHashMap<String, u8> = FxHashMap::default();
        let mut rodeo = Rodeo::default();
        let empty_spur = rodeo.get_or_intern("");

        let mut by_postcode: FxHashMap<String, (u32, u32)> = FxHashMap::default();
        let mut cur_pc: Option<String> = None;
        let mut cur_start: u32 = 0;
        // Absolute record index across all chunks; drives both the CSR index and
        // the column builders' write order.
        let mut global_row: u32 = 0;

        let columns: Vec<String> = ["postcode", "month_index", "crime_type", "location", "outcome", "lat", "lon"]
            .iter()
            .map(|s| s.to_string())
            .collect();

        let mut offset = 0usize;
        while offset < n {
            let len = chunk_rows.min(n - offset);
            let file = File::open(path).with_context(|| {
                format!("Failed to open crime-records parquet at {}", path.display())
            })?;
            let mut reader = ParquetReader::new(file);
            reader.set_metadata(metadata.clone());
            // `with_slice` only decodes the row groups overlapping `[offset,
            // offset+len)` (the file is memory-mapped, so untouched groups are
            // never faulted in), capping the transient decode to one chunk.
            let df = reader
                .with_columns(Some(columns.clone()))
                .with_slice(Some((offset, len)))
                .finish()
                .with_context(|| {
                    format!(
                        "Failed to read crime-records rows [{offset}, {}) from {}",
                        offset + len,
                        path.display()
                    )
                })?;

            let postcode_col = df
                .column("postcode")
                .context("crime-records parquet missing 'postcode'")?
                .str()
                .context("'postcode' is not a string")?;
            let month_col = df
                .column("month_index")
                .context("crime-records parquet missing 'month_index'")?
                .cast(&DataType::Int32)
                .context("'month_index' not castable to i32")?;
            let month_ca = month_col.i32().context("'month_index' is not i32")?;
            // Months are `year*12 + month0` (~24_000), always positive. A null or
            // non-positive value means a corrupt parquet; fail loudly rather than
            // silently clamping it to 0 and later rendering it as "0000-01".
            if month_ca.null_count() > 0 {
                bail!("crime-records 'month_index' has null values (corrupt parquet)");
            }
            match month_ca.min() {
                Some(m) if m > 0 => {}
                _ => bail!("crime-records 'month_index' must be a positive year*12+month index"),
            }
            let type_col = df
                .column("crime_type")
                .context("crime-records parquet missing 'crime_type'")?
                .str()
                .context("'crime_type' is not a string")?;
            let location_col = df
                .column("location")
                .context("crime-records parquet missing 'location'")?
                .str()
                .context("'location' is not a string")?;
            let outcome_col = df
                .column("outcome")
                .context("crime-records parquet missing 'outcome'")?
                .str()
                .context("'outcome' is not a string")?;
            let lat_col = df
                .column("lat")
                .context("crime-records parquet missing 'lat'")?
                .cast(&DataType::Float32)?;
            let lat_ca = lat_col.f32().context("'lat' is not f32")?;
            let lon_col = df
                .column("lon")
                .context("crime-records parquet missing 'lon'")?
                .cast(&DataType::Float32)?;
            let lon_ca = lon_col.f32().context("'lon' is not f32")?;

            let height = df.height();
            for row in 0..height {
                // CSR index: the parquet is sorted by postcode, so a change in the
                // postcode value (across chunk boundaries too) closes the previous
                // run and opens a new one.
                let pc = postcode_col
                    .get(row)
                    .with_context(|| {
                        format!("crime-records row {} has null postcode", offset + row)
                    })?
                    .trim();
                if cur_pc.as_deref() != Some(pc) {
                    if let Some(prev) = cur_pc.take() {
                        by_postcode.insert(prev, (cur_start, global_row - cur_start));
                    }
                    cur_pc = Some(pc.to_string());
                    cur_start = global_row;
                }

                month.push(month_ca.get(row).unwrap() as u32);

                let ty = type_col.get(row).unwrap_or("");
                let ty_id = match type_index.get(ty) {
                    Some(&id) => id,
                    None => {
                        let id = u8::try_from(crime_type_dict.len())
                            .context("more than 256 distinct crime types")?;
                        crime_type_dict.push(ty.to_string());
                        type_index.insert(ty.to_string(), id);
                        id
                    }
                };
                ctype.push(ty_id);

                let oc = outcome_col.get(row).unwrap_or("");
                let oc_id = if oc.is_empty() {
                    0
                } else {
                    match outcome_index.get(oc) {
                        Some(&id) => id,
                        None => {
                            let id = u8::try_from(outcome_dict.len())
                                .context("more than 256 distinct outcomes")?;
                            outcome_dict.push(oc.to_string());
                            outcome_index.insert(oc.to_string(), id);
                            id
                        }
                    }
                };
                outcome.push(oc_id);

                let loc = location_col.get(row).unwrap_or("");
                location.push(if loc.is_empty() {
                    empty_spur
                } else {
                    rodeo.get_or_intern(loc)
                });

                lat.push(lat_ca.get(row).unwrap_or(f32::NAN));
                lon.push(lon_ca.get(row).unwrap_or(f32::NAN));

                global_row += 1;
            }

            offset += len;
        }
        if let Some(prev) = cur_pc.take() {
            by_postcode.insert(prev, (cur_start, global_row - cur_start));
        }
        debug_assert_eq!(global_row as usize, n, "streamed fewer rows than the parquet declares");

        let records = Self {
            month: month.finish()?,
            ctype: ctype.finish()?,
            outcome: outcome.finish()?,
            location: location.finish()?,
            lat: lat.finish()?,
            lon: lon.finish()?,
            crime_type_dict,
            outcome_dict,
            location_resolver: rodeo.into_reader(),
            by_postcode,
        };

        info!(
            records = n,
            postcodes = records.by_postcode.len(),
            crime_types = records.crime_type_dict.len(),
            outcomes = records.outcome_dict.len(),
            "Crime records loaded"
        );
        Ok(records)
    }
}

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

    fn write_fixture(path: &Path) {
        // Two postcodes, postcode-sorted. AA1 1AA has 3 records across two
        // months (a null outcome and a null location), BB2 2BB has 1.
        let mut df = df!(
            "postcode" => ["AA1 1AA", "AA1 1AA", "AA1 1AA", "BB2 2BB"],
            "month_index" => [24300i32, 24300, 24290, 24305],
            "crime_type" => ["Burglary", "Burglary", "Robbery", "Drugs"],
            "location" => [Some("On or near A St"), Some("On or near A St"), None, Some("On or near B Rd")],
            "outcome" => [Some("Under investigation"), None, Some("Court result"), None],
            "lat" => [51.5f32, 51.5, 51.6, 52.0],
            "lon" => [-0.1f32, -0.1, -0.2, -1.0],
        )
        .unwrap();
        let mut file = std::fs::File::create(path).unwrap();
        ParquetWriter::new(&mut file).finish(&mut df).unwrap();
    }

    #[test]
    fn loads_indexes_and_gathers() {
        let dir = std::env::temp_dir().join(format!("crimerec-{}", std::process::id()));
        std::fs::create_dir_all(&dir).unwrap();
        let path = dir.join("records.parquet");
        write_fixture(&path);

        let recs = CrimeRecords::load(&path, None).unwrap();
        assert_eq!(recs.total_for("AA1 1AA"), 3);
        assert_eq!(recs.total_for("BB2 2BB"), 1);
        assert_eq!(recs.total_for("ZZ9 9ZZ"), 0);

        // Newest-first across the two postcodes.
        let all = recs.gather(&["AA1 1AA", "BB2 2BB"], None);
        assert_eq!(all.len(), 4);
        let months: Vec<u32> = all.iter().map(|&i| recs.view(i).month_index).collect();
        assert_eq!(months, vec![24305, 24300, 24300, 24290]);

        // `since` window filter (keep months >= 24300).
        assert_eq!(recs.gather(&["AA1 1AA"], Some(24300)).len(), 2);

        // String resolution + null handling.
        let robbery = all
            .iter()
            .map(|&i| recs.view(i))
            .find(|v| v.crime_type == "Robbery")
            .unwrap();
        assert_eq!(robbery.outcome, Some("Court result"));
        assert_eq!(robbery.location, None); // null location → None
        // Two records have a null outcome (an AA1 Burglary and the BB2 Drugs).
        let null_outcomes = all
            .iter()
            .map(|&i| recs.view(i))
            .filter(|v| v.outcome.is_none())
            .count();
        assert_eq!(null_outcomes, 2);

        std::fs::remove_dir_all(&dir).ok();
    }

    /// The CSR per-postcode index and the column builders must compose correctly
    /// across streaming chunk boundaries — including a postcode run split between
    /// two chunks. Forces `chunk_rows = 2` over the 4-row fixture so AA1 1AA's
    /// three records straddle the boundary (rows 0,1 in chunk 0; row 2 in chunk 1)
    /// and is exercised both heap-backed (no spill) and mmap-backed (spill).
    #[test]
    fn streams_across_chunk_boundaries() {
        let base = std::env::temp_dir().join(format!("crimerec-chunk-{}", std::process::id()));
        std::fs::create_dir_all(&base).unwrap();
        let path = base.join("records.parquet");
        write_fixture(&path);
        let spill = base.join("spill");
        std::fs::create_dir_all(&spill).unwrap();

        for spill_dir in [None, Some(spill.as_path())] {
            let recs = CrimeRecords::load_inner(&path, spill_dir, 2).unwrap();
            // Counts match regardless of how the runs were split across chunks.
            assert_eq!(recs.total_for("AA1 1AA"), 3);
            assert_eq!(recs.total_for("BB2 2BB"), 1);
            assert_eq!(recs.total_for("ZZ9 9ZZ"), 0);

            // Full gather, newest-first, identical to the single-chunk load.
            let all = recs.gather(&["AA1 1AA", "BB2 2BB"], None);
            assert_eq!(all.len(), 4);
            let months: Vec<u32> = all.iter().map(|&i| recs.view(i).month_index).collect();
            assert_eq!(months, vec![24305, 24300, 24300, 24290]);

            // The run that straddled the boundary still resolves its strings.
            let robbery = all
                .iter()
                .map(|&i| recs.view(i))
                .find(|v| v.crime_type == "Robbery")
                .unwrap();
            assert_eq!(robbery.outcome, Some("Court result"));
            assert_eq!(robbery.location, None);
        }

        std::fs::remove_dir_all(&base).ok();
    }

    /// Peak/resident RSS in MiB from `/proc/self/status` (Linux only).
    fn rss_mib() -> (f64, f64) {
        let status = std::fs::read_to_string("/proc/self/status").unwrap_or_default();
        let field = |key: &str| -> f64 {
            status
                .lines()
                .find(|l| l.starts_with(key))
                .and_then(|l| l.split_whitespace().nth(1))
                .and_then(|kb| kb.parse::<f64>().ok())
                .map_or(0.0, |kb| kb / 1024.0)
        };
        (field("VmHWM:"), field("VmRSS:"))
    }

    /// Manual, real-data smoke test: load the actual ~500M-row parquet and report
    /// peak RSS, proving the streaming + spill load completes without the
    /// tens-of-GB `DataFrame` materialisation that OOMed the old `.collect()`.
    ///
    /// Run with:
    ///   PPC_REAL_CRIME_RECORDS=/path/to/crime_records.parquet \
    ///     cargo test --bins -- --ignored --nocapture real_crime_records_load_is_bounded
    #[test]
    #[ignore = "needs the full crime_records.parquet; run manually"]
    fn real_crime_records_load_is_bounded() {
        let path = std::env::var("PPC_REAL_CRIME_RECORDS")
            .unwrap_or_else(|_| "../property-data/crime_records.parquet".to_string());
        let path = Path::new(&path);
        if !path.exists() {
            eprintln!("skipping: {} not found", path.display());
            return;
        }
        let spill = std::env::var("PPC_REAL_SPILL")
            .unwrap_or_else(|_| "../.tmp/crime-spill-realtest".to_string());
        let spill = Path::new(&spill);
        std::fs::create_dir_all(spill).unwrap();

        let (hwm_before, _rss_before) = rss_mib();
        let start = std::time::Instant::now();
        let recs = CrimeRecords::load(path, Some(spill)).unwrap();
        let elapsed = start.elapsed();
        let (hwm_after, rss_after) = rss_mib();

        let total: u64 = recs.by_postcode.values().map(|&(_, c)| c as u64).sum();
        eprintln!(
            "loaded {} records across {} postcodes in {:.1}s | RSS peak {:.0}->{:.0} MiB (Δ{:.0}) resident now {:.0} MiB",
            total,
            recs.by_postcode.len(),
            elapsed.as_secs_f64(),
            hwm_before,
            hwm_after,
            hwm_after - hwm_before,
            rss_after,
        );

        assert!(recs.by_postcode.len() > 0, "expected at least one postcode");
        assert!(total > 0, "expected at least one record");
        // The old `.collect()` decoded all rows' string columns at once (tens of
        // GB). Streaming must keep the peak growth far below that; a generous 20GiB
        // ceiling still proves we never materialise the whole file.
        assert!(
            hwm_after - hwm_before < 20_480.0,
            "peak RSS grew by {:.0} MiB during load — streaming/spill not bounding memory",
            hwm_after - hwm_before
        );

        std::fs::remove_dir_all(spill).ok();
    }
}