//! Columnar search result rows built from Tantivy documents or fast fields.
//!
//! **Invariant:** For each value type, the order of entries in the
//! corresponding internal list ([`SearchRows::i64_columns`],
//! [`SearchRows::u64_columns`], [`SearchRows::f64_columns`],
//! [`SearchRows::bytes_columns`]) matches the order of fields with that type in
//! the `return_fields` slice passed to constructors and push methods (first I64
//! in `return_fields` → `i64_columns[0]`, and so on).

use anyhow::{Context, Result, bail};
use tantivy::{DocAddress, Searcher, SegmentOrdinal, TantivyDocument, schema::Type};

use crate::{reader::query_info::DocField, utils::tantivy_doc::DocFieldAccess};

enum CachedColumn {
    I64(tantivy::fastfield::Column<i64>),
    U64(tantivy::fastfield::Column<u64>),
    F64(tantivy::fastfield::Column<f64>),
}

#[derive(Default)]
pub(crate) struct FastFieldColumnCache {
    current_segment_ord: Option<SegmentOrdinal>,
    return_columns: Vec<CachedColumn>,
}

impl FastFieldColumnCache {
    pub(crate) fn ensure_segment(
        &mut self,
        searcher: &Searcher,
        segment_ord: SegmentOrdinal,
        return_fields: &[DocField],
    ) -> Result<()> {
        if self.current_segment_ord == Some(segment_ord) {
            return Ok(());
        }

        self.current_segment_ord = Some(segment_ord);
        self.return_columns.clear();

        let seg_reader = searcher.segment_reader(segment_ord);
        let ff_readers = seg_reader.fast_fields();

        for field in return_fields {
            if field.is_bm25_score() {
                self.return_columns.push(CachedColumn::F64(
                    tantivy::fastfield::Column::build_empty_column(seg_reader.max_doc()),
                ));
                continue;
            }
            match field.value_type {
                Type::I64 => {
                    let column = ff_readers.i64(&field.name).with_context(|| {
                        format!("I64 fast field column '{:?}' not available", field)
                    })?;
                    self.return_columns.push(CachedColumn::I64(column));
                }
                Type::U64 => {
                    let column = ff_readers.u64(&field.name).with_context(|| {
                        format!("U64 fast field column '{:?}' not available", field)
                    })?;
                    self.return_columns.push(CachedColumn::U64(column));
                }
                Type::F64 => {
                    let column = ff_readers.f64(&field.name).with_context(|| {
                        format!("F64 fast field column '{:?}' not available", field)
                    })?;
                    self.return_columns.push(CachedColumn::F64(column));
                }
                _ => {
                    bail!("Cannot use FAST fields for this field type: {:?}", field);
                }
            }
        }

        Ok(())
    }
}

macro_rules! define_numeric_column_data {
    ($name:ident, $ty:ty, $null_value:expr, $type_name:literal) => {
        #[derive(Debug, Default)]
        pub struct $name {
            pub col_name: String,
            pub values: Vec<$ty>,
            pub null_map: Vec<u8>,
        }

        impl $name {
            fn with_capacity(col_name: String, capacity: usize) -> Self {
                Self {
                    col_name,
                    values: Vec::with_capacity(capacity),
                    null_map: Vec::new(),
                }
            }

            fn push_value(&mut self, value: Option<$ty>) {
                match value {
                    Some(value) => {
                        self.values.push(value);
                        if !self.null_map.is_empty() {
                            self.null_map.push(0);
                        }
                    }
                    None => {
                        if self.null_map.is_empty() {
                            self.null_map.resize(self.values.len(), 0);
                        }
                        self.values.push($null_value);
                        self.null_map.push(1);
                    }
                }
            }

            fn append(&mut self, other: $name) -> Result<()> {
                if self.col_name != other.col_name {
                    bail!(
                        "{} column mismatch: expect {}, got {}",
                        $type_name,
                        self.col_name,
                        other.col_name
                    );
                }

                let self_len = self.values.len();
                let other_len = other.values.len();
                self.values.extend(other.values);
                append_null_map(&mut self.null_map, other.null_map, self_len, other_len);
                Ok(())
            }

            pub fn value_at(&self, row_idx: usize) -> Option<$ty> {
                if self.is_null(row_idx) {
                    None
                } else {
                    self.values.get(row_idx).copied()
                }
            }

            pub fn is_null(&self, row_idx: usize) -> bool {
                is_null_at(&self.null_map, row_idx)
            }
        }
    };
}

define_numeric_column_data!(I64ColumnData, i64, Default::default(), "i64");
define_numeric_column_data!(U64ColumnData, u64, Default::default(), "u64");
define_numeric_column_data!(F64ColumnData, f64, Default::default(), "f64");

#[derive(Debug, Default)]
pub struct BytesColumnData {
    pub col_name: String,
    pub chars: Vec<u8>,
    pub offsets: Vec<u64>,
    pub null_map: Vec<u8>,
}

impl BytesColumnData {
    fn with_capacity(col_name: String, capacity: usize) -> Self {
        Self {
            col_name,
            chars: Vec::new(),
            offsets: Vec::with_capacity(capacity),
            null_map: Vec::new(),
        }
    }

    fn push_value(&mut self, value: Option<&[u8]>) {
        let is_null = value.is_none();
        if is_null {
            if self.null_map.is_empty() {
                self.null_map.resize(self.offsets.len(), 0);
            }
            self.null_map.push(1);
        } else if !self.null_map.is_empty() {
            self.null_map.push(0);
        }

        self.chars.extend_from_slice(value.unwrap_or(&[]));
        self.chars.push(0);
        self.offsets.push(self.chars.len() as u64);
    }

    fn append(&mut self, other: BytesColumnData) -> Result<()> {
        if self.col_name != other.col_name {
            bail!(
                "bytes column mismatch: expect {}, got {}",
                self.col_name,
                other.col_name
            );
        }

        let self_len = self.offsets.len();
        let other_len = other.offsets.len();
        let char_base = self.chars.len() as u64;
        self.chars.extend(other.chars);
        self.offsets
            .extend(other.offsets.into_iter().map(|offset| offset + char_base));
        append_null_map(&mut self.null_map, other.null_map, self_len, other_len);
        Ok(())
    }

    pub fn value_at(&self, row_idx: usize) -> Option<&[u8]> {
        if self.is_null(row_idx) {
            return None;
        }

        let end = *self.offsets.get(row_idx)? as usize;
        let start = if row_idx == 0 {
            0
        } else {
            *self.offsets.get(row_idx - 1)? as usize
        };
        self.chars.get(start..end - 1)
    }

    pub fn is_null(&self, row_idx: usize) -> bool {
        is_null_at(&self.null_map, row_idx)
    }
}

#[derive(Debug, Default)]
pub struct SearchRows {
    pub row_count: u64,
    pub i64_columns: Vec<I64ColumnData>,
    pub u64_columns: Vec<U64ColumnData>,
    pub f64_columns: Vec<F64ColumnData>,
    pub bytes_columns: Vec<BytesColumnData>,
}

impl SearchRows {
    pub fn empty(return_fields: &[DocField]) -> Result<Self> {
        Self::with_capacity(return_fields, 0)
    }

    pub fn with_capacity(return_fields: &[DocField], capacity: usize) -> Result<Self> {
        let mut rows = Self {
            row_count: 0,
            ..Default::default()
        };

        for field in return_fields {
            match field.value_type {
                Type::I64 => rows
                    .i64_columns
                    .push(I64ColumnData::with_capacity(field.name.clone(), capacity)),
                Type::U64 => rows
                    .u64_columns
                    .push(U64ColumnData::with_capacity(field.name.clone(), capacity)),
                Type::F64 => rows
                    .f64_columns
                    .push(F64ColumnData::with_capacity(field.name.clone(), capacity)),
                Type::Str | Type::Bytes => rows
                    .bytes_columns
                    .push(BytesColumnData::with_capacity(field.name.clone(), capacity)),
                _ => {
                    bail!(
                        "unsupported return field type {:?} for column {}",
                        field.value_type,
                        field.name
                    );
                }
            }
        }

        Ok(rows)
    }

    pub fn from_docs(docs: &[TantivyDocument], return_fields: &[DocField]) -> Result<Self> {
        Self::from_doc_refs(docs.iter(), docs.len(), return_fields)
    }

    pub fn from_doc_refs<'a>(
        docs: impl IntoIterator<Item = &'a TantivyDocument>,
        capacity: usize,
        return_fields: &[DocField],
    ) -> Result<Self> {
        let mut rows = Self::with_capacity(return_fields, capacity)?;
        for doc in docs {
            rows.push_document(doc, return_fields)?;
        }
        Ok(rows)
    }

    pub fn from_fast_fields(
        searcher: &Searcher,
        doc_addresses: &[DocAddress],
        return_fields: &[DocField],
    ) -> Result<Self> {
        if doc_addresses.is_empty() {
            return Self::with_capacity(return_fields, 0);
        }

        let mut rows = Self::with_capacity(return_fields, doc_addresses.len())?;
        let mut cached_columns = FastFieldColumnCache::default();

        for doc_address in doc_addresses {
            cached_columns.ensure_segment(searcher, doc_address.segment_ord, return_fields)?;
            rows.push_fast_fields(*doc_address, return_fields, &cached_columns)?;
        }

        Ok(rows)
    }

    pub(crate) fn push_fast_fields(
        &mut self,
        doc_address: DocAddress,
        return_fields: &[DocField],
        cached_columns: &FastFieldColumnCache,
    ) -> Result<()> {
        if cached_columns.return_columns.len() != return_fields.len() {
            bail!("Fast field column cache is not initialized");
        }

        let mut i64_idx = 0_usize;
        let mut u64_idx = 0_usize;
        let mut f64_idx = 0_usize;
        for (doc_field, column) in return_fields
            .iter()
            .zip(cached_columns.return_columns.iter())
        {
            match (doc_field.value_type, column) {
                (Type::I64, CachedColumn::I64(column)) => {
                    self.i64_columns
                        .get_mut(i64_idx)
                        .with_context(|| format!("missing i64 column '{}'", doc_field.name))?
                        .push_value(column.first(doc_address.doc_id));
                    i64_idx += 1;
                }
                (Type::U64, CachedColumn::U64(column)) => {
                    self.u64_columns
                        .get_mut(u64_idx)
                        .with_context(|| format!("missing u64 column '{}'", doc_field.name))?
                        .push_value(column.first(doc_address.doc_id));
                    u64_idx += 1;
                }
                (Type::F64, CachedColumn::F64(column)) => {
                    let value = if doc_field.is_bm25_score() {
                        Some(0.0)
                    } else {
                        column.first(doc_address.doc_id)
                    };
                    self.f64_columns
                        .get_mut(f64_idx)
                        .with_context(|| format!("missing f64 column '{}'", doc_field.name))?
                        .push_value(value);
                    f64_idx += 1;
                }
                _ => bail!("Unsupported field type: {:?}", doc_field.value_type),
            }
        }

        self.row_count += 1;
        Ok(())
    }

    pub fn push_document(
        &mut self,
        doc: &TantivyDocument,
        return_fields: &[DocField],
    ) -> Result<()> {
        let mut i64_idx = 0_usize;
        let mut u64_idx = 0_usize;
        let mut f64_idx = 0_usize;
        let mut bytes_idx = 0_usize;

        for field in return_fields {
            match field.value_type {
                Type::I64 => {
                    self.i64_columns
                        .get_mut(i64_idx)
                        .with_context(|| format!("missing i64 column '{}'", field.name))?
                        .push_value(doc.get_i64(field)?);
                    i64_idx += 1;
                }
                Type::U64 => {
                    self.u64_columns
                        .get_mut(u64_idx)
                        .with_context(|| format!("missing u64 column '{}'", field.name))?
                        .push_value(doc.get_u64(field)?);
                    u64_idx += 1;
                }
                Type::Str => {
                    self.bytes_columns
                        .get_mut(bytes_idx)
                        .with_context(|| format!("missing bytes column '{}'", field.name))?
                        .push_value(doc.get_str(field)?.map(str::as_bytes));
                    bytes_idx += 1;
                }
                Type::F64 => {
                    self.f64_columns
                        .get_mut(f64_idx)
                        .with_context(|| format!("missing f64 column '{}'", field.name))?
                        .push_value(if field.is_bm25_score() {
                            Some(0.0)
                        } else {
                            doc.get_f64(field)?
                        });
                    f64_idx += 1;
                }
                Type::Bytes => {
                    self.bytes_columns
                        .get_mut(bytes_idx)
                        .with_context(|| format!("missing bytes column '{}'", field.name))?
                        .push_value(doc.get_bytes(field)?);
                    bytes_idx += 1;
                }
                _ => {
                    bail!("Unsupported field type for primary key: {:?}", field);
                }
            }
        }

        self.row_count += 1;
        Ok(())
    }

    /// Returns `Err` if `self` and `other` do not have the same column layout:
    /// for each of i64 / u64 / f64 / bytes, the per-type count, column
    /// order, and names must match. (Column type is fixed by which Vec
    /// holds it.)
    fn ensure_append_layout_compatible(&self, other: &SearchRows) -> Result<()> {
        macro_rules! check {
            ($field:ident, $label:literal) => {{
                if self.$field.len() != other.$field.len() {
                    bail!(
                        "cannot merge SearchRows: {} column count mismatch: left={}, right={} (left names: {:?}, right names: {:?})",
                        $label,
                        self.$field.len(),
                        other.$field.len(),
                        self.$field.iter().map(|c| c.col_name.as_str()).collect::<Vec<_>>(),
                        other.$field.iter().map(|c| c.col_name.as_str()).collect::<Vec<_>>(),
                    );
                }
                for (i, (a, b)) in self.$field.iter().zip(other.$field.iter()).enumerate() {
                    if a.col_name != b.col_name {
                        bail!(
                            "cannot merge SearchRows: {} column at index {}: expected name {:?}, got {:?}",
                            $label, i, a.col_name, b.col_name
                        );
                    }
                }
            }};
        }
        check!(i64_columns, "i64");
        check!(u64_columns, "u64");
        check!(f64_columns, "f64");
        check!(bytes_columns, "bytes");
        Ok(())
    }

    pub fn append(&mut self, other: SearchRows) -> Result<()> {
        self.ensure_append_layout_compatible(&other)?;

        for (column, other_column) in self.i64_columns.iter_mut().zip(other.i64_columns) {
            column.append(other_column)?;
        }
        for (column, other_column) in self.u64_columns.iter_mut().zip(other.u64_columns) {
            column.append(other_column)?;
        }
        for (column, other_column) in self.f64_columns.iter_mut().zip(other.f64_columns) {
            column.append(other_column)?;
        }
        for (column, other_column) in self.bytes_columns.iter_mut().zip(other.bytes_columns) {
            column.append(other_column)?;
        }

        self.row_count += other.row_count;
        Ok(())
    }

    pub fn len(&self) -> usize {
        self.row_count as usize
    }

    pub fn is_empty(&self) -> bool {
        self.row_count == 0
    }

    pub fn column_count(&self) -> usize {
        self.i64_columns.len()
            + self.u64_columns.len()
            + self.f64_columns.len()
            + self.bytes_columns.len()
    }

    /// Column names in storage order: all i64 fields, then u64, f64, and bytes.
    pub fn column_names(&self) -> Vec<&str> {
        self.i64_columns
            .iter()
            .map(|c| c.col_name.as_str())
            .chain(self.u64_columns.iter().map(|c| c.col_name.as_str()))
            .chain(self.f64_columns.iter().map(|c| c.col_name.as_str()))
            .chain(self.bytes_columns.iter().map(|c| c.col_name.as_str()))
            .collect()
    }

    pub fn i64_column(&self, field_name: &str) -> Option<&I64ColumnData> {
        self.i64_columns
            .iter()
            .find(|column| column.col_name == field_name)
    }

    pub fn u64_column(&self, field_name: &str) -> Option<&U64ColumnData> {
        self.u64_columns
            .iter()
            .find(|column| column.col_name == field_name)
    }

    pub fn f64_column(&self, field_name: &str) -> Option<&F64ColumnData> {
        self.f64_columns
            .iter()
            .find(|column| column.col_name == field_name)
    }

    pub fn bytes_column(&self, field_name: &str) -> Option<&BytesColumnData> {
        self.bytes_columns
            .iter()
            .find(|column| column.col_name == field_name)
    }
}

fn is_null_at(null_map: &[u8], row_idx: usize) -> bool {
    !null_map.is_empty() && null_map[row_idx] != 0
}

fn append_null_map(dest: &mut Vec<u8>, src: Vec<u8>, dest_rows_before: usize, src_rows: usize) {
    if dest.is_empty() && src.is_empty() {
        return;
    }

    if dest.is_empty() {
        dest.resize(dest_rows_before, 0);
    }

    if src.is_empty() {
        dest.resize(dest_rows_before + src_rows, 0);
    } else {
        dest.extend(src);
    }
}

#[cfg(test)]
mod tests {
    use tantivy::schema::Field;

    use super::*;

    fn test_doc_field(name: &str, field_id: u32, value_type: Type) -> DocField {
        DocField {
            name: name.to_string(),
            field: Field::from_field_id(field_id),
            value_type,
        }
    }

    #[test]
    fn test_search_rows_empty_initializes_columns() -> Result<()> {
        let return_fields = vec![
            test_doc_field("i64_col", 0, Type::I64),
            test_doc_field("u64_col", 1, Type::U64),
            test_doc_field("f64_col", 2, Type::F64),
            test_doc_field("bytes_col", 3, Type::Bytes),
        ];
        let rows = SearchRows::empty(&return_fields)?;

        assert_eq!(rows.row_count, 0);
        assert_eq!(rows.i64_columns.len(), 1);
        assert_eq!(rows.u64_columns.len(), 1);
        assert_eq!(rows.f64_columns.len(), 1);
        assert_eq!(rows.bytes_columns.len(), 1);
        assert_eq!(rows.i64_columns[0].col_name, "i64_col");
        assert_eq!(rows.u64_columns[0].col_name, "u64_col");
        assert_eq!(rows.f64_columns[0].col_name, "f64_col");
        assert_eq!(rows.bytes_columns[0].col_name, "bytes_col");
        Ok(())
    }

    #[test]
    fn test_fast_field_bm25_score_returns_zero_from_f64_column() -> Result<()> {
        let return_fields = vec![test_doc_field("column_-2050", 0, Type::F64)];
        let cached_columns = FastFieldColumnCache {
            current_segment_ord: Some(0),
            return_columns: vec![CachedColumn::F64(
                tantivy::fastfield::Column::build_empty_column(2),
            )],
        };
        let mut rows = SearchRows::with_capacity(&return_fields, 1)?;

        rows.push_fast_fields(
            DocAddress {
                segment_ord: 0,
                doc_id: 1,
            },
            &return_fields,
            &cached_columns,
        )?;

        let score_col = rows.f64_column("column_-2050").unwrap();
        assert_eq!(score_col.values, vec![0.0]);
        assert!(score_col.null_map.is_empty());
        Ok(())
    }

    #[test]
    fn test_search_rows_returns_zero_bm25_score() -> Result<()> {
        let return_fields = vec![test_doc_field("column_-2050", 0, Type::F64)];
        let rows = SearchRows::from_docs(&[TantivyDocument::default()], &return_fields)?;

        assert_eq!(rows.row_count, 1);
        let score_col = rows.f64_column("column_-2050").unwrap();
        assert_eq!(score_col.values, vec![0.0]);
        assert!(score_col.null_map.is_empty());
        Ok(())
    }

    #[test]
    fn test_search_rows_append_preserves_values() -> Result<()> {
        let mut left = SearchRows {
            row_count: 1,
            i64_columns: vec![I64ColumnData {
                col_name: "i64_col".to_string(),
                values: vec![11],
                null_map: vec![],
            }],
            u64_columns: vec![U64ColumnData {
                col_name: "u64_col".to_string(),
                values: vec![0],
                null_map: vec![1],
            }],
            f64_columns: vec![F64ColumnData {
                col_name: "f64_col".to_string(),
                values: vec![3.5],
                null_map: vec![],
            }],
            bytes_columns: vec![BytesColumnData {
                col_name: "bytes_col".to_string(),
                chars: b"foo\0".to_vec(),
                offsets: vec![4],
                null_map: vec![],
            }],
        };
        let right = SearchRows {
            row_count: 1,
            i64_columns: vec![I64ColumnData {
                col_name: "i64_col".to_string(),
                values: vec![-9],
                null_map: vec![],
            }],
            u64_columns: vec![U64ColumnData {
                col_name: "u64_col".to_string(),
                values: vec![42],
                null_map: vec![0],
            }],
            f64_columns: vec![F64ColumnData {
                col_name: "f64_col".to_string(),
                values: vec![-1.25],
                null_map: vec![],
            }],
            bytes_columns: vec![BytesColumnData {
                col_name: "bytes_col".to_string(),
                chars: vec![0],
                offsets: vec![1],
                null_map: vec![1],
            }],
        };

        left.append(right)?;
        assert_eq!(left.row_count, 2);
        assert_eq!(left.i64_column("i64_col").unwrap().values, vec![11, -9]);
        assert_eq!(left.u64_column("u64_col").unwrap().values, vec![0, 42]);
        assert_eq!(left.f64_column("f64_col").unwrap().values, vec![3.5, -1.25]);
        assert_eq!(
            left.bytes_column("bytes_col").unwrap().value_at(0),
            Some(b"foo".as_slice())
        );
        assert_eq!(left.bytes_column("bytes_col").unwrap().value_at(1), None);
        assert_eq!(left.u64_column("u64_col").unwrap().null_map, vec![1, 0]);
        assert_eq!(left.bytes_column("bytes_col").unwrap().null_map, vec![0, 1]);
        Ok(())
    }

    #[test]
    fn test_search_rows_append_rejects_column_mismatch() {
        let mut left = SearchRows {
            i64_columns: vec![I64ColumnData {
                col_name: "left".to_string(),
                values: vec![],
                null_map: vec![],
            }],
            ..Default::default()
        };
        let right = SearchRows {
            i64_columns: vec![I64ColumnData {
                col_name: "right".to_string(),
                values: vec![],
                null_map: vec![],
            }],
            ..Default::default()
        };

        let err = left.append(right).unwrap_err();
        let msg = err.to_string();
        assert!(
            msg.contains("i64") && msg.contains("left") && msg.contains("right"),
            "unexpected error: {msg}"
        );
    }

    #[test]
    fn test_numeric_columns_track_nulls_consistently() {
        let mut i64_column = I64ColumnData::with_capacity("i64_col".to_string(), 3);
        i64_column.push_value(Some(11));
        i64_column.push_value(None);
        i64_column.push_value(Some(-7));
        assert_eq!(i64_column.values, vec![11, 0, -7]);
        assert_eq!(i64_column.null_map, vec![0, 1, 0]);
        assert_eq!(i64_column.value_at(1), None);

        let mut u64_column = U64ColumnData::with_capacity("u64_col".to_string(), 3);
        u64_column.push_value(Some(9));
        u64_column.push_value(None);
        u64_column.push_value(Some(5));
        assert_eq!(u64_column.values, vec![9, 0, 5]);
        assert_eq!(u64_column.null_map, vec![0, 1, 0]);
        assert_eq!(u64_column.value_at(1), None);

        let mut f64_column = F64ColumnData::with_capacity("f64_col".to_string(), 3);
        f64_column.push_value(Some(1.25));
        f64_column.push_value(None);
        f64_column.push_value(Some(-3.5));
        assert_eq!(f64_column.values, vec![1.25, 0.0, -3.5]);
        assert_eq!(f64_column.null_map, vec![0, 1, 0]);
        assert_eq!(f64_column.value_at(1), None);
    }
}