use std::fs::{self, File};
use std::io::{self, BufWriter, Write as _};
use std::path::{Path, PathBuf};

use memmap2::Mmap;
use ndarray::{Array1, Array2};
use rayon::prelude::*;

use crate::bitvec::{PersistentBitVec, PersistentBitVecBuilder};
use crate::colgroup::{ColGroup, MatrixGroupOps};
use crate::layer_meta::LayerMeta;
use crate::meta::MatrixMeta;
use crate::tempbitvec::{TempBitVec, TempBitVecBuilder};
use crate::tempintvec::{TempCompactIntVec, TempCompactIntVecBuilder};
use crate::views::BitSliceView;

fn col_path(dir: &Path, col: usize) -> PathBuf {
    dir.join(format!("col_{col:06}.pbiv"))
}

// ── ColumnarBitMatrix ─────────────────────────────────────────────────────────

/// Per-column file layout (original format).
pub struct ColumnarBitMatrix {
    cols: Vec<PersistentBitVec>,
    n:    usize,
}

impl ColumnarBitMatrix {
    pub(crate) fn open(dir: &Path) -> io::Result<Self> {
        let meta = MatrixMeta::load(dir)?;
        let cols = (0..meta.n_cols)
            .map(|c| PersistentBitVec::open(&col_path(dir, c)))
            .collect::<io::Result<Vec<_>>>()?;
        Ok(Self { cols, n: meta.n })
    }

    pub(crate) fn n(&self)      -> usize { self.n }
    pub(crate) fn n_cols(&self) -> usize { self.cols.len() }
    pub(crate) fn col(&self, c: usize) -> &PersistentBitVec { &self.cols[c] }

    pub(crate) fn row(&self, slot: usize) -> Box<[bool]> {
        self.cols.iter().map(|c| c.get(slot)).collect()
    }

    pub(crate) fn fill_row(&self, slot: usize, buf: &mut [u32]) {
        for (c, col) in self.cols.iter().enumerate() {
            buf[c] = col.get(slot) as u32;
        }
    }

    pub(crate) fn count_ones(&self) -> Array1<u64> {
        let counts: Vec<u64> = (0..self.n_cols())
            .into_par_iter()
            .map(|c| self.col(c).count_ones())
            .collect();
        Array1::from_vec(counts)
    }

    pub(crate) fn partial_jaccard_dist_matrix(&self) -> (Array2<u64>, Array2<u64>) {
        pairwise2_matrix(self.n_cols(), |i, j| self.col(i).partial_jaccard_dist(self.col(j)))
    }

    pub(crate) fn partial_hamming_dist_matrix(&self) -> Array2<u64> {
        pairwise_matrix(self.n_cols(), |i, j| self.col(i).hamming_dist(self.col(j)))
    }

    pub(crate) fn append_column(dir: &Path, value_of: impl Fn(usize) -> bool) -> io::Result<()> {
        let mut meta = MatrixMeta::load(dir)?;
        let mut b = PersistentBitVecBuilder::new(meta.n, &col_path(dir, meta.n_cols))?;
        for slot in 0..meta.n {
            b.set(slot, value_of(slot));
        }
        b.close()?;
        meta.n_cols += 1;
        meta.save(dir)
    }
}

// ── PackedBitMatrix ───────────────────────────────────────────────────────────

const PBMX_MAGIC:  [u8; 4] = *b"PBMX";
const PBMX_HEADER: usize   = 24; // magic(4) + pad(4) + n_rows(8) + n_cols(8)
const PBIV_HEADER: usize   = 16; // magic(4) + pad(4) + n(8)

/// Single-file packed layout: all columns concatenated behind a header.
pub struct PackedBitMatrix {
    mmap:         Mmap,
    n_rows:       usize,
    n_cols:       usize,
    /// Absolute byte offset to the start of each column's bit data
    /// (= file offset of the PBIV blob + PBIV_HEADER).
    data_offsets: Vec<usize>,
}

impl PackedBitMatrix {
    pub(crate) fn open(path: &Path) -> io::Result<Self> {
        let mmap = unsafe { Mmap::map(&File::open(path)?)? };
        if mmap.len() < PBMX_HEADER {
            return Err(io::Error::new(io::ErrorKind::InvalidData, "PBMX file too short"));
        }
        if &mmap[0..4] != &PBMX_MAGIC {
            return Err(io::Error::new(io::ErrorKind::InvalidData, "bad PBMX magic"));
        }
        let n_rows = u64::from_le_bytes(mmap[8..16].try_into().unwrap())  as usize;
        let n_cols = u64::from_le_bytes(mmap[16..24].try_into().unwrap()) as usize;

        let mut data_offsets = Vec::with_capacity(n_cols);
        for c in 0..n_cols {
            let off_pos = PBMX_HEADER + c * 8;
            let col_file_off = u64::from_le_bytes(mmap[off_pos..off_pos+8].try_into().unwrap()) as usize;
            data_offsets.push(col_file_off + PBIV_HEADER);
        }

        Ok(Self { mmap, n_rows, n_cols, data_offsets })
    }

    #[inline]
    pub(crate) fn fill_row(&self, slot: usize, buf: &mut [u32]) {
        for (c, &data_off) in self.data_offsets.iter().enumerate() {
            buf[c] = ((self.mmap[data_off + (slot >> 3)] >> (slot & 7)) & 1) as u32;
        }
    }

    pub(crate) fn row(&self, slot: usize) -> Box<[bool]> {
        (0..self.n_cols).map(|c| {
            (self.mmap[self.data_offsets[c] + (slot >> 3)] >> (slot & 7)) & 1 != 0
        }).collect()
    }

    fn col_bytes(&self, c: usize) -> &[u8] {
        let start = self.data_offsets[c];
        &self.mmap[start..start + self.n_rows.div_ceil(8)]
    }

    fn col_words(&self, c: usize) -> &[u64] {
        let nw = self.n_rows.div_ceil(64);
        // SAFETY: data_offsets[c] is always 8-byte aligned.
        // PBMX header = 24 + n_cols×8 (multiple of 8); each PBIV blob =
        // 16 + nwords×8 (multiple of 8); mmap base is page-aligned.
        let ptr = self.mmap[self.data_offsets[c]..].as_ptr() as *const u64;
        unsafe { std::slice::from_raw_parts(ptr, nw) }
    }

    pub(crate) fn col_slice(&self, c: usize) -> BitSliceView<'_> {
        BitSliceView::new(self.col_words(c), self.n_rows)
    }

    pub(crate) fn col_persist(&self, c: usize, path: &Path) -> io::Result<PersistentBitVecBuilder> {
        PersistentBitVecBuilder::from_raw_bytes(self.col_bytes(c), self.n_rows, path)
    }

    pub(crate) fn count_ones(&self) -> Array1<u64> {
        Array1::from_vec(
            (0..self.n_cols).into_par_iter()
                .map(|c| self.col_slice(c).count_ones())
                .collect()
        )
    }

    pub(crate) fn partial_jaccard_dist_matrix(&self) -> (Array2<u64>, Array2<u64>) {
        pairwise2_matrix(self.n_cols, |i, j| {
            self.col_slice(i).partial_jaccard_dist(self.col_slice(j))
        })
    }

    pub(crate) fn partial_hamming_dist_matrix(&self) -> Array2<u64> {
        pairwise_matrix(self.n_cols, |i, j| {
            self.col_slice(i).hamming_dist(self.col_slice(j))
        })
    }
}

/// Build `presence/matrix.pbmx` from existing `col_*.pbiv` files.
pub fn pack_bit_matrix(dir: &Path) -> io::Result<()> {
    let packed_path = dir.join("matrix.pbmx");
    if packed_path.exists() {
        // Matrix complete; remove any leftover column files from a killed cleanup.
        if let Ok(meta) = MatrixMeta::load(dir) {
            for c in 0..meta.n_cols { let _ = fs::remove_file(col_path(dir, c)); }
            let _ = fs::remove_file(dir.join("meta.json"));
        }
        return Ok(());
    }

    let meta = MatrixMeta::load(dir)?;
    let n_cols = meta.n_cols;

    // Compute offsets from file sizes — no column data loaded into RAM.
    let col_sizes: Vec<u64> = (0..n_cols)
        .map(|c| fs::metadata(col_path(dir, c)).map(|m| m.len()))
        .collect::<io::Result<_>>()?;

    let header_size = (PBMX_HEADER + n_cols * 8) as u64;
    let mut col_offset = header_size;
    let mut offsets = Vec::with_capacity(n_cols);
    for &size in &col_sizes {
        offsets.push(col_offset);
        col_offset += size;
    }

    // Write to a temp file; rename atomically so a killed process never leaves
    // a truncated matrix.pbmx that would be mistaken for a complete file.
    let tmp_path = dir.join("matrix.pbmx.tmp");
    let mut out = BufWriter::new(File::create(&tmp_path)?);
    out.write_all(&PBMX_MAGIC)?;
    out.write_all(&[0u8; 4])?;
    out.write_all(&(meta.n as u64).to_le_bytes())?;
    out.write_all(&(n_cols as u64).to_le_bytes())?;
    for &off in &offsets { out.write_all(&off.to_le_bytes())?; }
    for c in 0..n_cols {
        io::copy(&mut File::open(col_path(dir, c))?, &mut out)?;
    }
    out.flush()?;
    drop(out);
    fs::rename(&tmp_path, &packed_path)?;

    for c in 0..n_cols { fs::remove_file(col_path(dir, c))?; }
    fs::remove_file(dir.join("meta.json"))?;
    Ok(())
}

// ── PersistentBitMatrix — public enum ────────────────────────────────────────

/// Bit matrix that transparently handles columnar, packed, and implicit formats.
///
/// - `Columnar`: per-column `.pbiv` files (original format, used during build)
/// - `Packed`:   single `matrix.pbmx` file (optimised for query — one `mmap`)
/// - `Implicit`: no file — all values are 1 (mono-genome presence/absence)
pub enum PersistentBitMatrix {
    Columnar(ColumnarBitMatrix),
    Packed(PackedBitMatrix),
    Implicit { n_rows: usize, n_cols: usize },
}

impl PersistentBitMatrix {
    /// Open from `layer_dir`, auto-detecting the format.
    ///
    /// Checks (in order):
    /// 1. `layer_dir/presence/matrix.pbmx`  → Packed
    /// 2. `layer_dir/presence/meta.json`    → Columnar
    /// 3. `layer_dir/layer_meta.json`        → Implicit (new index)
    /// 4. `layer_dir/unitigs.bin`            → Implicit with warning (old index)
    pub fn open(layer_dir: &Path) -> io::Result<Self> {
        let presence_dir = layer_dir.join("presence");

        if presence_dir.join("matrix.pbmx").exists() {
            return Ok(Self::Packed(PackedBitMatrix::open(&presence_dir.join("matrix.pbmx"))?));
        }

        if MatrixMeta::load(&presence_dir).is_ok() {
            return Ok(Self::Columnar(ColumnarBitMatrix::open(&presence_dir)?));
        }

        // No presence matrix → Implicit; requires layer_meta.json
        let meta = LayerMeta::load(layer_dir).map_err(|_| io::Error::new(
            io::ErrorKind::NotFound,
            format!(
                "no presence matrix and no layer_meta.json in {} — run 'obikmer upgrade'",
                layer_dir.display()
            ),
        ))?;
        Ok(Self::Implicit { n_rows: meta.n, n_cols: 1 })
    }

    pub fn n(&self) -> usize {
        match self {
            Self::Columnar(m)            => m.n(),
            Self::Packed(m)              => m.n_rows,
            Self::Implicit { n_rows, .. } => *n_rows,
        }
    }

    pub fn n_cols(&self) -> usize {
        match self {
            Self::Columnar(m)            => m.n_cols(),
            Self::Packed(m)              => m.n_cols,
            Self::Implicit { n_cols, .. } => *n_cols,
        }
    }

    pub fn col(&self, c: usize) -> &PersistentBitVec {
        match self {
            Self::Columnar(m) => m.col(c),
            _ => panic!("col() only available on Columnar PersistentBitMatrix"),
        }
    }

    pub fn col_view(&self, c: usize) -> BitSliceView<'_> {
        match self {
            Self::Columnar(m) => m.col(c).view(),
            Self::Packed(m)   => m.col_slice(c),
            Self::Implicit { .. } => panic!("col_view() not available on Implicit PersistentBitMatrix"),
        }
    }

    pub fn col_persist(&self, c: usize, path: &Path) -> io::Result<PersistentBitVecBuilder> {
        match self {
            Self::Columnar(m) => PersistentBitVecBuilder::build_from(m.col(c), path),
            Self::Packed(m)   => m.col_persist(c, path),
            Self::Implicit { n_rows, .. } => {
                PersistentBitVecBuilder::new_ones(*n_rows, path)
            }
        }
    }

    pub fn row(&self, slot: usize) -> Box<[bool]> {
        match self {
            Self::Columnar(m)             => m.row(slot),
            Self::Packed(m)               => m.row(slot),
            Self::Implicit { n_cols, .. } => vec![true; *n_cols].into_boxed_slice(),
        }
    }

    /// Fill `buf[i]` with `col_i[slot]` as 0/1 u32, without allocating.
    pub fn fill_row(&self, slot: usize, buf: &mut [u32]) {
        match self {
            Self::Columnar(m)             => m.fill_row(slot, buf),
            Self::Packed(m)               => m.fill_row(slot, buf),
            Self::Implicit { n_cols, .. } => buf[..*n_cols].fill(1),
        }
    }

    pub fn count_ones(&self) -> Array1<u64> {
        match self {
            Self::Columnar(m)              => m.count_ones(),
            Self::Packed(m)                => m.count_ones(),
            Self::Implicit { n_rows, n_cols } => Array1::from_elem(*n_cols, *n_rows as u64),
        }
    }

    pub fn partial_jaccard_dist_matrix(&self) -> (Array2<u64>, Array2<u64>) {
        match self {
            Self::Columnar(m)                 => m.partial_jaccard_dist_matrix(),
            Self::Packed(m)                   => m.partial_jaccard_dist_matrix(),
            Self::Implicit { n_rows, n_cols } => {
                let v = *n_rows as u64;
                let n = *n_cols;
                let mut inter = Array2::zeros((n, n));
                let mut union = Array2::zeros((n, n));
                for i in 0..n { for j in 0..n {
                    inter[[i, j]] = v; union[[i, j]] = v;
                }}
                (inter, union)
            }
        }
    }

    pub fn partial_hamming_dist_matrix(&self) -> Array2<u64> {
        match self {
            Self::Columnar(m)              => m.partial_hamming_dist_matrix(),
            Self::Packed(m)                => m.partial_hamming_dist_matrix(),
            Self::Implicit { n_cols, .. }  => Array2::zeros((*n_cols, *n_cols)),
        }
    }

    /// Append a new column to an on-disk Columnar matrix.
    pub fn append_column(dir: &Path, value_of: impl Fn(usize) -> bool) -> io::Result<()> {
        ColumnarBitMatrix::append_column(dir, value_of)
    }
}

// ── Trait impls ───────────────────────────────────────────────────────────────

use crate::traits::{BitPartials, ColumnWeights};

impl ColumnWeights for PersistentBitMatrix {
    fn col_weights(&self) -> Array1<u64> { self.count_ones() }
}

impl BitPartials for PersistentBitMatrix {
    fn partial_jaccard(&self) -> (Array2<u64>, Array2<u64>) {
        self.partial_jaccard_dist_matrix()
    }
    fn partial_hamming(&self) -> Array2<u64> {
        self.partial_hamming_dist_matrix()
    }
}

// ── Builder (unchanged — always builds Columnar) ──────────────────────────────

pub struct PersistentBitMatrixBuilder {
    dir:    PathBuf,
    n:      usize,
    n_cols: usize,
}

impl PersistentBitMatrixBuilder {
    pub fn new(n: usize, dir: &Path) -> io::Result<Self> {
        fs::create_dir_all(dir)?;
        Ok(Self { dir: dir.to_path_buf(), n, n_cols: 0 })
    }

    pub fn n(&self)      -> usize { self.n }
    pub fn n_cols(&self) -> usize { self.n_cols }

    pub fn add_col(&mut self) -> io::Result<PersistentBitVecBuilder> {
        let path = col_path(&self.dir, self.n_cols);
        self.n_cols += 1;
        PersistentBitVecBuilder::new(self.n, &path)
    }

    pub fn add_col_ones(&mut self) -> io::Result<PersistentBitVecBuilder> {
        let path = col_path(&self.dir, self.n_cols);
        self.n_cols += 1;
        PersistentBitVecBuilder::new_ones(self.n, &path)
    }

    pub fn add_col_from(&mut self, src: &TempBitVec) -> io::Result<()> {
        src.make_persistent(&col_path(&self.dir, self.n_cols))?;
        self.n_cols += 1;
        Ok(())
    }

    pub fn add_col_from_int(&mut self, src: &TempCompactIntVec) -> io::Result<()> {
        let path = col_path(&self.dir, self.n_cols);
        self.n_cols += 1;
        let mut b = PersistentBitVecBuilder::new(self.n, &path)?;
        b.or_where(src.view(), |v| v > 0);
        b.close()
    }

    pub fn close(self) -> io::Result<()> {
        MatrixMeta { n: self.n, n_cols: self.n_cols }.save(&self.dir)
    }
}

// ── MatrixGroupOps ────────────────────────────────────────────────────────────

impl MatrixGroupOps for PersistentBitMatrix {
    fn partial_group_presence_count(&self, g: &ColGroup, _threshold: u32) -> io::Result<TempCompactIntVec> {
        // Bit matrices store 0/1 — threshold is structurally always 1.
        let n = self.n();
        if g.indices.len() < 255 {
            let mut builder = TempCompactIntVecBuilder::new(n)?;
            for &c in &g.indices {
                builder.inc_present_fast(self.col_view(c));
            }
            builder.freeze()
        } else {
            let mut result = TempCompactIntVecBuilder::new(n)?;
            for chunk in g.indices.chunks(254) {
                let mut chunk_b = TempCompactIntVecBuilder::new(n)?;
                for &c in chunk {
                    chunk_b.inc_present_fast(self.col_view(c));
                }
                let frozen = chunk_b.freeze()?;
                result.add(frozen.view());
            }
            result.freeze()
        }
    }

    fn partial_group_sum(&self, g: &ColGroup) -> io::Result<TempCompactIntVec> {
        // For bit matrices, sum = count of 1-bits — identical to presence_count.
        self.partial_group_presence_count(g, 1)
    }

    fn partial_group_any(&self, g: &ColGroup, _threshold: u32) -> io::Result<TempBitVec> {
        let n = self.n();
        let mut result = TempBitVecBuilder::new(n)?;
        for &c in &g.indices {
            result.or(self.col_view(c));
        }
        result.freeze()
    }

    fn partial_group_min(&self, g: &ColGroup) -> io::Result<TempCompactIntVec> {
        // min of 0/1 values = AND: 1 only if ALL columns are 1
        let n = self.n();
        let mut result = TempCompactIntVecBuilder::new(n)?;
        if let Some((&first, rest)) = g.indices.split_first() {
            result.inc_present_fast(self.col_view(first));
            for &c in rest { result.mask_with(self.col_view(c)); }
        }
        result.freeze()
    }

    fn partial_group_max(&self, g: &ColGroup) -> io::Result<TempCompactIntVec> {
        // max of 0/1 values = OR: 1 if any column is 1
        let any = self.partial_group_any(g, 1)?;
        let n = any.len();
        let mut result = TempCompactIntVecBuilder::new(n)?;
        result.inc_present(any.view());
        result.freeze()
    }
}

// ── Shared matrix helpers (also used by intmatrix.rs) ─────────────────────────

fn upper_pairs(n: usize) -> Vec<(usize, usize)> {
    (0..n).flat_map(|i| (i + 1..n).map(move |j| (i, j))).collect()
}

fn fill_symmetric<T>(n: usize, vals: impl Iterator<Item = (usize, usize, T, T)>) -> Array2<T>
where T: Clone + Default {
    let mut m = Array2::from_elem((n, n), T::default());
    for (i, j, vij, vji) in vals { m[[i, j]] = vij; m[[j, i]] = vji; }
    m
}

/// Compute a symmetric `n×n` matrix in parallel by evaluating `f(i,j)` for
/// all upper-triangle pairs. `T: Copy` avoids the `.clone()` needed for the
/// lower-triangle mirror.
pub(crate) fn pairwise_matrix<T>(n: usize, f: impl Fn(usize, usize) -> T + Sync) -> Array2<T>
where T: Copy + Default + Send {
    let results: Vec<(usize, usize, T)> = upper_pairs(n)
        .into_par_iter().map(|(i, j)| (i, j, f(i, j))).collect();
    fill_symmetric(n, results.into_iter().map(|(i, j, v)| (i, j, v, v)))
}

/// Same as `pairwise_matrix` but `f` returns two values that fill two
/// symmetric matrices simultaneously (e.g. intersection + union for Jaccard).
pub(crate) fn pairwise2_matrix<T>(n: usize, f: impl Fn(usize, usize) -> (T, T) + Sync) -> (Array2<T>, Array2<T>)
where T: Copy + Default + Send {
    let results: Vec<(usize, usize, T, T)> = upper_pairs(n)
        .into_par_iter()
        .map(|(i, j)| { let (a, b) = f(i, j); (i, j, a, b) })
        .collect();
    let mut m0 = Array2::from_elem((n, n), T::default());
    let mut m1 = Array2::from_elem((n, n), T::default());
    for (i, j, a, b) in results {
        m0[[i, j]] = a; m0[[j, i]] = a;
        m1[[i, j]] = b; m1[[j, i]] = b;
    }
    (m0, m1)
}
