use std::fs;
use std::io;
use std::path::{Path, PathBuf};

use obicompactvec::{
    ColGroup, MatrixGroupOps,
    PersistentBitMatrix, PersistentBitMatrixBuilder,
    PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder,
};
use obilayeredmap::meta::PartitionMeta;
use obilayeredmap::OLMError;
use obiskio::{SKError, SKResult};

use crate::partition::KmerPartition;

const INDEX_SUBDIR: &str = "index";

// ── AggOp ─────────────────────────────────────────────────────────────────────

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AggOp {
    Any,
    All,
    None,
    Sum,
    Min,
    Max,
}

impl AggOp {
    pub fn is_logical(self) -> bool {
        matches!(self, AggOp::Any | AggOp::All | AggOp::None)
    }
}

// ── OutputCol ─────────────────────────────────────────────────────────────────

pub struct OutputCol {
    pub label:   String,
    pub indices: Vec<usize>,
    pub op:      AggOp,
}

// ── Helpers ───────────────────────────────────────────────────────────────────

fn olm_to_sk(e: OLMError) -> SKError {
    match e {
        OLMError::Io(e) => SKError::Io(e),
        other => SKError::InvalidData { context: "select", detail: other.to_string() },
    }
}

/// Copy all plain files (not subdirectories) from `src_dir` to `dst_dir`.
fn copy_layer_files(src_dir: &Path, dst_dir: &Path) -> io::Result<()> {
    for entry in fs::read_dir(src_dir)? {
        let entry = entry?;
        let path = entry.path();
        if path.is_file() {
            fs::copy(&path, dst_dir.join(entry.file_name()))?;
        }
    }
    Ok(())
}

// ── fill_builders ─────────────────────────────────────────────────────────────

fn fill_builders(
    specs: &[OutputCol],
    src_layer_dir: &Path,
    src_is_count: bool,
    threshold: u32,
    output_presence: bool,
    mut dst_bit: Option<&mut PersistentBitMatrixBuilder>,
    mut dst_int: Option<&mut PersistentCompactIntMatrixBuilder>,
) -> SKResult<()> {
    if src_is_count {
        let mat = PersistentCompactIntMatrix::open(src_layer_dir).map_err(SKError::Io)?;
        for spec in specs {
            let g = ColGroup::new(&spec.label, spec.indices.clone());
            if output_presence {
                let b = dst_bit.as_deref_mut().unwrap();
                match spec.op {
                    AggOp::Any  => b.add_col_from    (&mat.partial_group_any (&g, threshold).map_err(SKError::Io)?),
                    AggOp::All  => b.add_col_from    (&mat.partial_group_all (&g, threshold).map_err(SKError::Io)?),
                    AggOp::None => b.add_col_from    (&mat.partial_group_none(&g, threshold).map_err(SKError::Io)?),
                    AggOp::Sum  => b.add_col_from_int(&mat.partial_group_sum (&g).map_err(SKError::Io)?),
                    AggOp::Min  => b.add_col_from_int(&mat.partial_group_min (&g).map_err(SKError::Io)?),
                    AggOp::Max  => b.add_col_from_int(&mat.partial_group_max (&g).map_err(SKError::Io)?),
                }.map_err(SKError::Io)?;
            } else {
                let b = dst_int.as_deref_mut().unwrap();
                match spec.op {
                    AggOp::Sum  => b.add_col_from    (&mat.partial_group_sum (&g).map_err(SKError::Io)?),
                    AggOp::Min  => b.add_col_from    (&mat.partial_group_min (&g).map_err(SKError::Io)?),
                    AggOp::Max  => b.add_col_from    (&mat.partial_group_max (&g).map_err(SKError::Io)?),
                    AggOp::Any  => b.add_col_from_bit(&mat.partial_group_any (&g, threshold).map_err(SKError::Io)?),
                    AggOp::All  => b.add_col_from_bit(&mat.partial_group_all (&g, threshold).map_err(SKError::Io)?),
                    AggOp::None => b.add_col_from_bit(&mat.partial_group_none(&g, threshold).map_err(SKError::Io)?),
                }.map_err(SKError::Io)?;
            }
        }
    } else {
        let mat = PersistentBitMatrix::open(src_layer_dir).map_err(SKError::Io)?;
        for spec in specs {
            let g = ColGroup::new(&spec.label, spec.indices.clone());
            if output_presence {
                let b = dst_bit.as_deref_mut().unwrap();
                match spec.op {
                    AggOp::Any  => b.add_col_from    (&mat.partial_group_any (&g, 1).map_err(SKError::Io)?),
                    AggOp::All  => b.add_col_from    (&mat.partial_group_all (&g, 1).map_err(SKError::Io)?),
                    AggOp::None => b.add_col_from    (&mat.partial_group_none(&g, 1).map_err(SKError::Io)?),
                    AggOp::Sum  => b.add_col_from_int(&mat.partial_group_sum (&g).map_err(SKError::Io)?),
                    AggOp::Min  => b.add_col_from_int(&mat.partial_group_min (&g).map_err(SKError::Io)?),
                    AggOp::Max  => b.add_col_from_int(&mat.partial_group_max (&g).map_err(SKError::Io)?),
                }.map_err(SKError::Io)?;
            } else {
                let b = dst_int.as_deref_mut().unwrap();
                match spec.op {
                    AggOp::Sum  => b.add_col_from    (&mat.partial_group_sum (&g).map_err(SKError::Io)?),
                    AggOp::Min  => b.add_col_from    (&mat.partial_group_min (&g).map_err(SKError::Io)?),
                    AggOp::Max  => b.add_col_from    (&mat.partial_group_max (&g).map_err(SKError::Io)?),
                    AggOp::Any  => b.add_col_from_bit(&mat.partial_group_any (&g, 1).map_err(SKError::Io)?),
                    AggOp::All  => b.add_col_from_bit(&mat.partial_group_all (&g, 1).map_err(SKError::Io)?),
                    AggOp::None => b.add_col_from_bit(&mat.partial_group_none(&g, 1).map_err(SKError::Io)?),
                }.map_err(SKError::Io)?;
            }
        }
    }
    Ok(())
}

// ── KmerPartition::select_partition ──────────────────────────────────────────

impl KmerPartition {
    /// Rewrite the data matrices of partition `i` in `src` into `self`.
    ///
    /// `specs` defines the output columns (projection/aggregation).
    /// `output_presence` — if true, all output builders use bit (0/1) format.
    /// `in_place` — `self` and `src` share the same root; write to temp dirs then swap.
    pub fn select_partition(
        &self,
        src: &KmerPartition,
        i: usize,
        specs: &[OutputCol],
        _n_src_genomes: usize,
        threshold: u32,
        output_presence: bool,
        in_place: bool,
    ) -> SKResult<()> {
        let src_index_dir = src.part_dir(i).join(INDEX_SUBDIR);
        if !src_index_dir.exists() {
            return Ok(());
        }

        let src_meta = PartitionMeta::load(&src_index_dir).map_err(olm_to_sk)?;
        if src_meta.n_layers == 0 {
            return Ok(());
        }

        let dst_index_dir = self.part_dir(i).join(INDEX_SUBDIR);
        if !in_place {
            fs::create_dir_all(&dst_index_dir)?;
        }

        let data_subdir = if output_presence { "presence" } else { "counts" };

        for l in 0..src_meta.n_layers {
            let src_layer_dir = src_index_dir.join(format!("layer_{l}"));
            if !src_layer_dir.exists() { continue; }

            let dst_layer_dir = dst_index_dir.join(format!("layer_{l}"));

            let counts_dir   = src_layer_dir.join("counts");
            let presence_dir = src_layer_dir.join("presence");
            let src_is_count = counts_dir.exists() && !presence_dir.exists();

            // Determine number of slots and detect implicit layers.
            let n = if counts_dir.exists() {
                PersistentCompactIntMatrix::open(&src_layer_dir).map_err(SKError::Io)?.n()
            } else if presence_dir.exists() {
                PersistentBitMatrix::open(&src_layer_dir).map_err(SKError::Io)?.n()
            } else {
                // Implicit single-genome layer: no data matrix needed in output either.
                if !in_place {
                    fs::create_dir_all(&dst_layer_dir)?;
                    copy_layer_files(&src_layer_dir, &dst_layer_dir)?;
                }
                continue;
            };

            // Choose the output data directory (temp name for in-place).
            let (dst_data_dir, final_data_dir): (PathBuf, PathBuf) = if in_place {
                let tmp  = dst_layer_dir.join(format!("{data_subdir}_new"));
                let perm = dst_layer_dir.join(data_subdir);
                (tmp, perm)
            } else {
                let perm = dst_layer_dir.join(data_subdir);
                (perm.clone(), perm)
            };

            if !in_place {
                fs::create_dir_all(&dst_layer_dir)?;
                copy_layer_files(&src_layer_dir, &dst_layer_dir)?;
            }
            fs::create_dir_all(&dst_data_dir)?;

            let (mut dst_bit, mut dst_int) = if output_presence {
                (Some(PersistentBitMatrixBuilder::new(n, &dst_data_dir).map_err(SKError::Io)?), None)
            } else {
                (None, Some(PersistentCompactIntMatrixBuilder::new(n, &dst_data_dir).map_err(SKError::Io)?))
            };

            fill_builders(
                specs, &src_layer_dir, src_is_count, threshold, output_presence,
                dst_bit.as_mut(), dst_int.as_mut(),
            )?;

            if output_presence {
                dst_bit.unwrap().close().map_err(SKError::Io)?;
            } else {
                dst_int.unwrap().close().map_err(SKError::Io)?;
            }

            // In-place: swap old data dir for new.
            if in_place {
                let old_data_dir = if src_is_count {
                    dst_layer_dir.join("counts")
                } else {
                    dst_layer_dir.join("presence")
                };
                if old_data_dir.exists() {
                    fs::remove_dir_all(&old_data_dir)?;
                }
                fs::rename(&dst_data_dir, &final_data_dir)?;
            }
        }

        if !in_place {
            PartitionMeta::load(&src_index_dir).map_err(olm_to_sk)?
                .save(&dst_index_dir).map_err(olm_to_sk)?;
        }

        Ok(())
    }
}
