//! IndexWriter: buffer documents, analyze text, flush segments, commit.
//!
//! The public write API for Luci. Documents are validated against the
//! schema, text fields are analyzed, and results are accumulated in an
//! in-memory segment buffer. On commit (or memory budget exceeded), the
//! buffer is flushed to a segment and written to storage.
//!
//! See [[architecture-indexing-pipeline]] and [[architecture-overview#Step 9]].

use luci_analysis::{AnalyzerRegistry, Token};
use luci_core::{FieldId, LuciError, Result, SegmentId};
use luci_mapping::{DynamicMode, FieldType,Mapping};
use luci_storage::Storage;

use crate::columnar::writer::ColumnValue;
use crate::spatial::geo::GeoPoint;
use crate::segment::builder::SegmentBuilder;

/// Default memory budget before auto-flush (64 MB).
const DEFAULT_MEMORY_BUDGET: usize = 64 * 1024 * 1024;

/// Writes documents to a Luci index.
///
/// Usage:
/// 1. Create with `IndexWriter::new(storage, schema, analyzers)`
/// 2. Call `put(doc)` for each document
/// 3. Call `commit()` to make documents searchable
pub struct IndexWriter {
    storage: Box<dyn Storage>,
    schema: Mapping,
    analyzers: AnalyzerRegistry,
    builder: SegmentBuilder,
    next_segment_id: u64,
    memory_budget: usize,
    /// Approximate bytes consumed by the current buffer.
    buffer_size: usize,
    /// Merge policy configuration.
    merge_policy: crate::merge_policy::MergePolicy,
    /// Monotonic counter for auto-generated document IDs.
    next_doc_id: u64,
    /// Pending deletions (applied during search, cleaned up on merge).
    pending_deletions: crate::deletion::DeletionMap,
}

impl IndexWriter {
    /// Create a new IndexWriter.
    pub fn new(storage: impl Storage + 'static, schema: Mapping, analyzers: AnalyzerRegistry) -> Self {
        let next_id = storage.generation() + 1;
        let builder = SegmentBuilder::new(SegmentId::new(next_id), &schema);
        Self {
            storage: Box::new(storage),
            schema,
            analyzers,
            builder,
            next_segment_id: next_id,
            memory_budget: DEFAULT_MEMORY_BUDGET,
            buffer_size: 0,
            merge_policy: crate::merge_policy::MergePolicy::default(),
            next_doc_id: std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .map(|d| d.as_millis() as u64 * 1000)
                .unwrap_or(0),
            pending_deletions: crate::deletion::DeletionMap::new(),
        }
    }

    /// Mark a document as deleted. Takes effect on next search (no commit needed).
    pub fn mark_deleted(&mut self, segment_id: SegmentId, doc_id: luci_core::DocId) {
        self.pending_deletions.mark_deleted(segment_id, doc_id);
    }

    /// Get the current deletion map (for search-time filtering).
    pub fn deletions(&self) -> &crate::deletion::DeletionMap {
        &self.pending_deletions
    }

    /// Set the memory budget for auto-flush.
    pub fn set_memory_budget(&mut self, budget: usize) {
        self.memory_budget = budget;
    }

    /// Add a JSON document to the index.
    ///
    /// The document is validated against the schema, text fields are analyzed,
    /// and the result is buffered. Call `commit()` to make it searchable.
    pub fn add(&mut self, doc: serde_json::Value) -> Result<()> {
        let obj = doc
            .as_object()
            .ok_or_else(|| LuciError::InvalidQuery("document must be a JSON object".into()))?;

        let source = serde_json::to_vec(&doc)
            .map_err(|e| LuciError::InvalidQuery(format!("JSON serialization failed: {e}")))?;

        let mut analyzed_fields: Vec<(FieldId, Vec<Token>)> = Vec::new();
        let mut column_values: Vec<(FieldId, ColumnValue)> = Vec::new();

        // Inject _id: use user-provided or auto-generate
        let doc_id_str = match obj.get("_id").and_then(|v| v.as_str()) {
            Some(id) => id.to_string(),
            None => {
                self.next_doc_id += 1;
                format!("{:x}", self.next_doc_id - 1)
            }
        };
        if let Some(id_field_id) = self.schema.field_id("_id") {
            analyzed_fields.push((id_field_id, vec![Token::new(doc_id_str.clone(), 0, 0, 0)]));
            column_values.push((id_field_id, ColumnValue::Keyword(doc_id_str)));
        }
        let mut vector_fields: Vec<(FieldId, Vec<f32>)> = Vec::new();
        let mut geo_points: Vec<(FieldId, GeoPoint)> = Vec::new();
        let mut geo_shapes: Vec<(FieldId, ::geo::Geometry<f64>)> = Vec::new();
        let mut copy_to_pairs: Vec<(String, String)> = Vec::new(); // (target_field, source_text)

        for (field_name, value) in obj {
            // Look up or dynamically map the field
            let field_id = match self.schema.field_id(field_name) {
                Some(id) => id,
                None => match self.schema.dynamic_mode() {
                    DynamicMode::False => continue, // store in _source but don't index
                    DynamicMode::True => {
                        // TODO: infer field type and add to mapping dynamically.
                        // For now, skip unknown fields (stored in _source only).
                        continue;
                    }
                },
            };

            let mapping = self.schema.field(field_id);

            // Build tokens for the inverted index (skipped if index: false)
            let tokens = match &mapping.field_type {
                FieldType::Text => {
                    let text = value.as_str().unwrap_or_default();
                    let analyzer_name = mapping.analyzer.as_deref().unwrap_or("standard");
                    let analyzer = self.analyzers.get(analyzer_name);
                    analyzer.analyze(text)
                }
                FieldType::Keyword => {
                    let text = match value {
                        serde_json::Value::String(s) => s.clone(),
                        other => other.to_string(),
                    };
                    vec![Token::new(text.clone(), 0, 0, 0)]
                }
                FieldType::Ip => {
                    // Index normalized IP string as keyword token
                    let text = value.as_str().unwrap_or_default();
                    let normalized = crate::ip::normalize_ip(text);
                    if normalized.is_empty() { Vec::new() }
                    else { vec![Token::new(normalized, 0, 0, 0)] }
                }
                _ => {
                    // Numeric/boolean/date fields: only columnar, no inverted index
                    Vec::new()
                }
            };

            if !tokens.is_empty() && mapping.indexed {
                analyzed_fields.push((field_id, tokens));
            }

            // Store geo shape for geo_shape fields
            if matches!(mapping.field_type, FieldType::GeoShape) {
                if let Some(geom) = crate::spatial::shape::parse_geojson(value) {
                    geo_shapes.push((field_id, geom));
                }
            }

            // Store geo point for geo_point fields
            if matches!(mapping.field_type, FieldType::GeoPoint) {
                if let Some(point) = GeoPoint::from_json(value) {
                    geo_points.push((field_id, point));
                }
            }

            // Store vector for dense_vector fields
            if mapping.field_type.is_dense_vector() {
                if let serde_json::Value::Array(arr) = value {
                    let vec: Vec<f32> = arr.iter()
                        .filter_map(|v| v.as_f64().map(|f| f as f32))
                        .collect();
                    if !vec.is_empty() {
                        vector_fields.push((field_id, vec));
                    }
                }
            }

            // Store column value for doc_values fields
            if mapping.doc_values {
                let col_val = match &mapping.field_type {
                    FieldType::Keyword => match value {
                        serde_json::Value::String(s) => ColumnValue::Keyword(s.clone()),
                        serde_json::Value::Null => ColumnValue::Null,
                        other => ColumnValue::Keyword(other.to_string()),
                    },
                    FieldType::Integer | FieldType::Long => match value {
                        serde_json::Value::Number(n) => {
                            ColumnValue::I64(n.as_i64().unwrap_or(0))
                        }
                        _ => ColumnValue::Null,
                    },
                    FieldType::Float | FieldType::Double => match value {
                        serde_json::Value::Number(n) => {
                            ColumnValue::F64(n.as_f64().unwrap_or(0.0))
                        }
                        _ => ColumnValue::Null,
                    },
                    FieldType::Boolean => match value {
                        serde_json::Value::Bool(b) => ColumnValue::Bool(*b),
                        _ => ColumnValue::Null,
                    },
                    FieldType::TokenCount => {
                        let text = value.as_str().unwrap_or_default();
                        let analyzer_name = mapping.analyzer.as_deref().unwrap_or("standard");
                        let analyzer = self.analyzers.get(analyzer_name);
                        ColumnValue::I64(analyzer.analyze(text).len() as i64)
                    },
                    FieldType::Ip => {
                        let text = value.as_str().unwrap_or_default();
                        match crate::ip::ip_to_i64(text) {
                            Some(v) => ColumnValue::I64(v),
                            None => ColumnValue::Null,
                        }
                    },
                    _ => ColumnValue::Null, // Text fields don't get doc_values
                };
                column_values.push((field_id, col_val));
            }

            // Collect copy_to pairs for post-loop processing
            if !mapping.copy_to.is_empty() {
                let source_text = match value {
                    serde_json::Value::String(s) => s.clone(),
                    other => other.to_string(),
                };
                for target in &mapping.copy_to {
                    copy_to_pairs.push((target.clone(), source_text.clone()));
                }
            }
        }

        // Process copy_to: analyze source text with target field's analyzer.
        // See [[feature-mapping-copy-to]].
        for (target_name, source_text) in &copy_to_pairs {
            let target_id = match self.schema.field_id(target_name) {
                Some(id) => id,
                None => continue,
            };
            let target_mapping = self.schema.field(target_id);
            if !target_mapping.indexed { continue; }

            let tokens = match &target_mapping.field_type {
                FieldType::Text => {
                    let analyzer_name = target_mapping.analyzer.as_deref().unwrap_or("standard");
                    let analyzer = self.analyzers.get(analyzer_name);
                    analyzer.analyze(source_text)
                }
                FieldType::Keyword => {
                    vec![Token::new(source_text.clone(), 0, 0, 0)]
                }
                _ => continue,
            };
            if !tokens.is_empty() {
                analyzed_fields.push((target_id, tokens));
            }
        }

        // Process multi-field sub-fields: route parent's source value
        // through the sub-field's analysis chain.
        // See [[feature-mapping-multi-fields]].
        for mapping in self.schema.fields() {
            if let Some(ref parent_name) = mapping.parent_field {
                // Get the parent's value from the document
                let parent_value = match obj.get(parent_name) {
                    Some(v) => v,
                    None => continue,
                };
                let field_id = match self.schema.field_id(&mapping.name) {
                    Some(id) => id,
                    None => continue,
                };

                if mapping.indexed {
                    let tokens = match &mapping.field_type {
                        FieldType::Text => {
                            let text = parent_value.as_str().unwrap_or_default();
                            let analyzer_name = mapping.analyzer.as_deref().unwrap_or("standard");
                            let analyzer = self.analyzers.get(analyzer_name);
                            analyzer.analyze(text)
                        }
                        FieldType::Keyword => {
                            let text = match parent_value {
                                serde_json::Value::String(s) => s.clone(),
                                other => other.to_string(),
                            };
                            vec![Token::new(text, 0, 0, 0)]
                        }
                        _ => continue,
                    };
                    analyzed_fields.push((field_id, tokens));
                }

                if mapping.doc_values {
                    let col_val = match &mapping.field_type {
                        FieldType::Keyword => {
                            let text = match parent_value {
                                serde_json::Value::String(s) => s.clone(),
                                other => other.to_string(),
                            };
                            ColumnValue::Keyword(text)
                        }
                        _ => continue,
                    };
                    column_values.push((field_id, col_val));
                }
            }
        }

        self.buffer_size += source.len();

        // Mark this as a parent doc if we have nested fields
        let has_nested = self.schema.fields().iter()
            .any(|f| matches!(f.field_type, FieldType::Nested));

        self.builder.add_document(&analyzed_fields, &source);

        if has_nested {
            self.builder.mark_parent();
        }

        // Add column values for doc_values fields
        for (field_id, col_val) in column_values {
            self.builder.add_column_value(field_id, col_val);
        }

        // Add vectors for dense_vector fields
        for (field_id, vec) in vector_fields {
            self.builder.add_vector(field_id, vec);
        }

        // Add geo points
        for (field_id, point) in geo_points {
            self.builder.add_geo_point(field_id, point);
        }

        // Add geo shapes
        for (field_id, geom) in &geo_shapes {
            self.builder.add_geo_shape(*field_id, geom);
        }

        // Index nested objects as hidden documents
        for mapping in self.schema.fields() {
            if !matches!(mapping.field_type, FieldType::Nested) {
                continue;
            }
            let field_name = &mapping.name;
            if let Some(serde_json::Value::Array(nested_arr)) = obj.get(field_name) {
                for nested_obj in nested_arr {
                    if let Some(nested_map) = nested_obj.as_object() {
                        // Index each nested object as a hidden document
                        let mut nested_fields: Vec<(FieldId, Vec<Token>)> = Vec::new();
                        for (nested_key, nested_val) in nested_map {
                            // Path-prefixed field name: "offers.seller"
                            let prefixed = format!("{field_name}.{nested_key}");
                            if let Some(fid) = self.schema.field_id(&prefixed) {
                                let m = self.schema.field(fid);
                                let tokens = match &m.field_type {
                                    FieldType::Text => {
                                        let text = nested_val.as_str().unwrap_or_default();
                                        let analyzer = self.analyzers.get(
                                            m.analyzer.as_deref().unwrap_or("standard"),
                                        );
                                        analyzer.analyze(text)
                                    }
                                    FieldType::Keyword => {
                                        let text = match nested_val {
                                            serde_json::Value::String(s) => s.clone(),
                                            other => other.to_string(),
                                        };
                                        vec![Token::new(text, 0, 0, 0)]
                                    }
                                    _ => continue,
                                };
                                if !tokens.is_empty() {
                                    nested_fields.push((fid, tokens));
                                }
                            }
                        }
                        // Add as hidden doc (empty source)
                        self.builder.add_document(&nested_fields, b"{}");
                        self.builder.mark_nested();
                    }
                }
            }
        }

        // Auto-flush if memory budget exceeded
        if self.buffer_size >= self.memory_budget {
            self.flush()?;
        }

        Ok(())
    }

    /// Flush the current buffer to a segment (without committing).
    fn flush(&mut self) -> Result<()> {
        if self.builder.is_empty() {
            return Ok(());
        }

        let segment_id = SegmentId::new(self.next_segment_id);
        self.next_segment_id += 1;

        // Take the current builder and replace with a new one
        let builder = std::mem::replace(
            &mut self.builder,
            SegmentBuilder::new(SegmentId::new(self.next_segment_id), &self.schema),
        );

        let segment_data = builder.build();
        self.storage.write_segment(segment_id, &segment_data)?;
        self.buffer_size = 0;

        Ok(())
    }

    /// Flush the buffer and commit all pending segments to storage.
    ///
    /// Persists the current field mappings alongside segment metadata.
    /// After commit returns, all documents added via `add()` are searchable.
    /// If the merge policy triggers, a synchronous merge is executed before
    /// returning.
    pub fn commit(&mut self) -> Result<()> {
        self.flush()?;
        // Persist mapping as user metadata
        let mapping_json = self.schema.to_json();
        let mapping_bytes = serde_json::to_vec(&mapping_json)
            .map_err(|e| LuciError::Io(std::io::Error::new(
                std::io::ErrorKind::Other, format!("failed to serialize mapping: {e}")
            )))?;
        self.storage.set_user_metadata(mapping_bytes);
        self.storage.commit()?;

        // Check if merge is needed after commit
        self.maybe_merge()?;

        Ok(())
    }

    /// Check the merge policy and execute a merge if needed.
    fn maybe_merge(&mut self) -> Result<()> {
        use crate::merge_policy::{find_merge, SegmentInfo};

        let infos: Vec<SegmentInfo> = self.storage.segments().iter().map(|e| {
            SegmentInfo {
                segment_id: e.segment_id,
                size_bytes: e.data_len,
                doc_count: 0, // TODO: store doc_count in SegmentEntry
                deletion_count: 0,
            }
        }).collect();

        let candidate = match find_merge(&self.merge_policy, &infos) {
            Some(c) => c,
            None => return Ok(()),
        };

        self.execute_merge(&candidate.segment_ids)
    }

    /// Execute a merge: read source segments, merge, write result, replace.
    fn execute_merge(&mut self, source_ids: &[SegmentId]) -> Result<()> {
        use crate::deletion::DeletionMap;
        use crate::segment::reader::SegmentReader;

        // Open source segment readers
        let mut readers = Vec::new();
        let mut segment_data = Vec::new();
        for &seg_id in source_ids {
            let data = self.storage.read_segment(seg_id)?;
            segment_data.push(data);
        }
        for data in &segment_data {
            readers.push(SegmentReader::open(data.clone())?);
        }
        let reader_refs: Vec<&SegmentReader> = readers.iter().collect();

        // Merge into a new segment
        let new_id = SegmentId::new(self.next_segment_id);
        self.next_segment_id += 1;

        let deletions = DeletionMap::new();
        let merged_data = crate::merge::merge_segments(
            new_id,
            &reader_refs,
            &deletions,
            &self.schema,
            &self.analyzers,
        );

        // Write the merged segment
        self.storage.write_segment(new_id, &merged_data)?;

        // Remove source segments
        self.storage.remove_segments(source_ids);

        // Commit the replacement
        self.storage.commit()?;

        Ok(())
    }

    /// Force-merge all segments down to at most `max_segments`.
    ///
    /// Repeatedly merges until the segment count is at or below the target.
    /// This is expensive and should only be called after bulk indexing is
    /// complete, not during normal operation.
    pub fn force_merge(&mut self, max_segments: usize) -> Result<()> {
        loop {
            let segments = self.storage.segments();
            if segments.len() <= max_segments {
                break;
            }

            // Merge all segments into one batch (up to max_merge_at_once)
            let ids: Vec<SegmentId> = segments.iter()
                .take(self.merge_policy.max_merge_at_once)
                .map(|e| e.segment_id)
                .collect();

            if ids.len() < 2 {
                break;
            }

            self.execute_merge(&ids)?;
        }
        Ok(())
    }

    /// Number of documents in the current (unflushed) buffer.
    pub fn buffered_doc_count(&self) -> u32 {
        self.builder.doc_count()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use luci_mapping::FieldType;
    use luci_storage::SingleFileDirectory;
    use std::path::PathBuf;

    fn test_dir(name: &str) -> PathBuf {
        let dir = std::env::temp_dir().join(format!(
            "luci_writer_test_{}_{name}",
            std::process::id()
        ));
        let _ = std::fs::remove_dir_all(&dir);
        dir
    }

    fn cleanup(path: &std::path::Path) {
        let _ = std::fs::remove_dir_all(path);
    }

    fn basic_setup(name: &str) -> (PathBuf, IndexWriter) {
        let path = test_dir(name);
        let storage = SingleFileDirectory::create(&path).unwrap();
        let schema = Mapping::builder()
            .field("title", FieldType::Text)
            .field("status", FieldType::Keyword)
            .build();
        let analyzers = AnalyzerRegistry::new();
        let writer = IndexWriter::new(storage, schema, analyzers);
        (path, writer)
    }

    #[test]
    fn put_single_doc_and_commit() {
        let (path, mut writer) = basic_setup("single");
        writer
            .add(serde_json::json!({
                "title": "hello world",
                "status": "active"
            }))
            .unwrap();
        writer.commit().unwrap();

        // Verify segment was written
        let storage = SingleFileDirectory::open(&path).unwrap();
        assert_eq!(storage.segments().len(), 1);

        cleanup(&path);
    }

    #[test]
    fn put_multiple_docs_and_commit() {
        let (path, mut writer) = basic_setup("multi");
        for i in 0..10 {
            writer
                .add(serde_json::json!({
                    "title": format!("document {i}"),
                    "status": "published"
                }))
                .unwrap();
        }
        writer.commit().unwrap();

        let storage = SingleFileDirectory::open(&path).unwrap();
        assert_eq!(storage.segments().len(), 1);

        cleanup(&path);
    }

    #[test]
    fn text_fields_analyzed() {
        let (path, mut writer) = basic_setup("analyzed");
        writer
            .add(serde_json::json!({
                "title": "The Quick Brown Fox",
                "status": "active"
            }))
            .unwrap();
        writer.commit().unwrap();

        // Read the segment and verify lowercase terms are in the index
        let storage = SingleFileDirectory::open(&path).unwrap();
        let seg_id = storage.segments()[0].segment_id;
        let data = storage.read_segment(seg_id).unwrap();

        use crate::segment::reader::SegmentReader;
        let reader = SegmentReader::open(data).unwrap();

        // Standard analyzer should lowercase
        assert!(reader.postings(FieldId::new(0), "the").is_some());
        assert!(reader.postings(FieldId::new(0), "quick").is_some());
        assert!(reader.postings(FieldId::new(0), "brown").is_some());
        assert!(reader.postings(FieldId::new(0), "fox").is_some());

        // Original casing should not be found
        assert!(reader.postings(FieldId::new(0), "The").is_none());
        assert!(reader.postings(FieldId::new(0), "Quick").is_none());

        cleanup(&path);
    }

    #[test]
    fn keyword_fields_exact() {
        let (path, mut writer) = basic_setup("keyword");
        writer
            .add(serde_json::json!({
                "title": "test",
                "status": "Active"
            }))
            .unwrap();
        writer.commit().unwrap();

        let storage = SingleFileDirectory::open(&path).unwrap();
        let data = storage.read_segment(storage.segments()[0].segment_id).unwrap();

        use crate::segment::reader::SegmentReader;
        let reader = SegmentReader::open(data).unwrap();

        // Keyword field should preserve case
        assert!(reader.postings(FieldId::new(1), "Active").is_some());
        assert!(reader.postings(FieldId::new(1), "active").is_none());

        cleanup(&path);
    }

    #[test]
    fn commit_with_no_docs_is_noop() {
        let (path, mut writer) = basic_setup("empty_commit");
        writer.commit().unwrap();

        let storage = SingleFileDirectory::open(&path).unwrap();
        assert!(storage.segments().is_empty());

        cleanup(&path);
    }

    #[test]
    fn auto_flush_on_memory_budget() {
        let (path, mut writer) = basic_setup("autoflush");
        writer.set_memory_budget(100); // Very small budget

        for i in 0..5 {
            writer
                .add(serde_json::json!({
                    "title": format!("document number {i} with some extra text to exceed the budget"),
                    "status": "active"
                }))
                .unwrap();
        }
        writer.commit().unwrap();

        // Multiple segments should have been created due to auto-flush
        let storage = SingleFileDirectory::open(&path).unwrap();
        assert!(
            storage.segments().len() > 1,
            "expected multiple segments from auto-flush, got {}",
            storage.segments().len()
        );

        cleanup(&path);
    }

    #[test]
    fn dynamic_false_ignores_unknown() {
        let path = test_dir("dynamic_false");
        let storage = SingleFileDirectory::create(&path).unwrap();
        let schema = Mapping::builder()
            .field("title", FieldType::Text)
            .dynamic(DynamicMode::False)
            .build();
        let analyzers = AnalyzerRegistry::new();
        let mut writer = IndexWriter::new(storage, schema, analyzers);

        // Should succeed — unknown field is silently ignored
        writer
            .add(serde_json::json!({
                "title": "hello",
                "unknown_field": "value"
            }))
            .unwrap();
        writer.commit().unwrap();

        cleanup(&path);
    }

    #[test]
    fn multiple_commits() {
        let (path, mut writer) = basic_setup("multi_commit");
        writer
            .add(serde_json::json!({"title": "first", "status": "a"}))
            .unwrap();
        writer.commit().unwrap();

        writer
            .add(serde_json::json!({"title": "second", "status": "b"}))
            .unwrap();
        writer.commit().unwrap();

        let storage = SingleFileDirectory::open(&path).unwrap();
        assert_eq!(storage.segments().len(), 2);

        cleanup(&path);
    }

    #[test]
    fn source_stored_correctly() {
        let (path, mut writer) = basic_setup("source");
        let doc = serde_json::json!({"title": "hello world", "status": "active"});
        writer.add(doc.clone()).unwrap();
        writer.commit().unwrap();

        let storage = SingleFileDirectory::open(&path).unwrap();
        let data = storage.read_segment(storage.segments()[0].segment_id).unwrap();

        use crate::segment::reader::SegmentReader;
        let reader = SegmentReader::open(data).unwrap();
        let source = reader.doc_store().get(0).unwrap();
        let stored: serde_json::Value = serde_json::from_slice(&source).unwrap();
        assert_eq!(stored, doc);

        cleanup(&path);
    }
}
