tauq 0.2.0

//! TBF Encoder with schema-aware encoding
//!
//! The encoder supports two modes:
//! 1. **Self-describing mode**: Every value has a type tag (flexible but larger)
//! 2. **Schema mode**: For homogeneous sequences, emit schema once then values without tags

use super::adaptive_encode::CodecAnalyzer;
use super::dictionary::StringDictionary;
use super::schema::{SchemaRegistry, SchemaType};
use super::stats_collector::StatisticsCollector;
use super::varint::*;
use super::{FLAG_CODEC_METADATA, FLAG_DICTIONARY, TBF_MAGIC, TBF_VERSION, TypeTag};
use std::collections::HashMap;

/// Encoding mode flags
/// Mode flag for self-describing values
pub const MODE_SELF_DESCRIBING: u8 = 0x00;
/// Mode flag for schema-based sequence encoding
pub const MODE_SCHEMA: u8 = 0x01;

/// Special marker for schema-encoded sequence
pub const MARKER_SCHEMA_SEQ: u8 = 0xFE;

/// TBF Serializer - encodes values directly to TBF binary format
///
/// Supports both self-describing and schema-based encoding for optimal size.
pub struct TbfSerializer {
    /// Output buffer for data
    pub(crate) buf: Vec<u8>,
    /// String dictionary for deduplication
    pub(crate) dict: StringDictionary,
    /// Schema registry
    pub(crate) schemas: SchemaRegistry,
    /// Schema buffer (written after dictionary)
    pub(crate) schema_buf: Vec<u8>,
    /// Codec metadata buffer (written after schemas, before data)
    pub(crate) codec_metadata_buf: Vec<u8>,
    /// Current encoding context
    pub(crate) context: EncodingContext,
    /// Nesting depth
    pub(crate) depth: usize,
    /// Optional statistics collector (Phase 2)
    pub(crate) stats: Option<StatisticsCollector>,
    /// Optional codec analyzer for adaptive compression (Phase 3)
    pub(crate) codec_analyzer: Option<CodecAnalyzer>,
    /// Selected codecs per field/column (Phase 3)
    pub(crate) selected_codecs: HashMap<String, super::adaptive_encode::CompressionCodec>,
    /// Collected codec metadata (Phase 3)
    pub(crate) codec_metadata: Vec<super::codec_encode::CodecMetadata>,
}

/// Tracks the current encoding context for schema detection
#[derive(Debug, Clone, Default)]
pub struct EncodingContext {
    /// Are we inside a sequence?
    pub in_sequence: bool,
    /// Number of elements seen in current sequence
    pub seq_element_count: usize,
    /// Detected schema for current sequence (if homogeneous)
    pub seq_schema: Option<DetectedSchema>,
    /// Current struct field names being collected
    pub current_fields: Vec<(String, SchemaType)>,
    /// Schema index for current sequence (if schema mode)
    pub seq_schema_idx: Option<u32>,
    /// Are we in schema mode for current sequence?
    pub schema_mode: bool,
}

/// A schema detected during serialization
#[derive(Debug, Clone)]
pub struct DetectedSchema {
    /// Field names in order
    pub fields: Vec<String>,
    /// Field types in order
    pub types: Vec<SchemaType>,
}

impl TbfSerializer {
    /// Create a new serializer
    pub fn new() -> Self {
        Self {
            buf: Vec::with_capacity(1024),
            dict: StringDictionary::new(),
            schemas: SchemaRegistry::new(),
            schema_buf: Vec::new(),
            codec_metadata_buf: Vec::new(),
            context: EncodingContext::default(),
            depth: 0,
            stats: None,
            codec_analyzer: None,
            selected_codecs: HashMap::new(),
            codec_metadata: Vec::new(),
        }
    }

    /// Create a serializer with pre-allocated capacity
    pub fn with_capacity(capacity: usize) -> Self {
        Self {
            buf: Vec::with_capacity(capacity),
            dict: StringDictionary::with_capacity(capacity / 32),
            schemas: SchemaRegistry::new(),
            schema_buf: Vec::new(),
            codec_metadata_buf: Vec::new(),
            context: EncodingContext::default(),
            depth: 0,
            stats: None,
            codec_analyzer: None,
            selected_codecs: HashMap::new(),
            codec_metadata: Vec::new(),
        }
    }

    /// Create a serializer with statistics collection enabled
    pub fn with_statistics() -> Self {
        Self {
            buf: Vec::with_capacity(1024),
            dict: StringDictionary::new(),
            schemas: SchemaRegistry::new(),
            schema_buf: Vec::new(),
            codec_metadata_buf: Vec::new(),
            context: EncodingContext::default(),
            depth: 0,
            stats: Some(StatisticsCollector::new()),
            codec_analyzer: None,
            selected_codecs: HashMap::new(),
            codec_metadata: Vec::new(),
        }
    }

    /// Create a serializer with pre-allocated capacity and statistics collection
    pub fn with_capacity_and_statistics(capacity: usize) -> Self {
        Self {
            buf: Vec::with_capacity(capacity),
            dict: StringDictionary::with_capacity(capacity / 32),
            schemas: SchemaRegistry::new(),
            schema_buf: Vec::new(),
            codec_metadata_buf: Vec::new(),
            context: EncodingContext::default(),
            depth: 0,
            stats: Some(StatisticsCollector::new()),
            codec_analyzer: None,
            selected_codecs: HashMap::new(),
            codec_metadata: Vec::new(),
        }
    }

    /// Create a serializer with adaptive codec selection enabled
    pub fn with_codecs() -> Self {
        Self {
            buf: Vec::with_capacity(1024),
            dict: StringDictionary::new(),
            schemas: SchemaRegistry::new(),
            schema_buf: Vec::new(),
            codec_metadata_buf: Vec::new(),
            context: EncodingContext::default(),
            depth: 0,
            stats: None,
            codec_analyzer: Some(CodecAnalyzer::new(100)),
            selected_codecs: HashMap::new(),
            codec_metadata: Vec::new(),
        }
    }

    /// Create a serializer with adaptive codec selection and statistics collection
    pub fn with_codecs_and_statistics() -> Self {
        Self {
            buf: Vec::with_capacity(1024),
            dict: StringDictionary::new(),
            schemas: SchemaRegistry::new(),
            schema_buf: Vec::new(),
            codec_metadata_buf: Vec::new(),
            context: EncodingContext::default(),
            depth: 0,
            stats: Some(StatisticsCollector::new()),
            codec_analyzer: Some(CodecAnalyzer::new(100)),
            selected_codecs: HashMap::new(),
            codec_metadata: Vec::new(),
        }
    }

    /// Finalize and get the encoded bytes
    pub fn into_bytes(mut self) -> Vec<u8> {
        // Encode dictionary
        let mut dict_buf = Vec::new();
        self.dict.encode(&mut dict_buf);

        // Encode schemas
        self.schemas.encode(&mut self.schema_buf, &mut self.dict);
        // Re-encode dictionary since schema encoding may have added strings
        dict_buf.clear();
        self.dict.encode(&mut dict_buf);

        // Encode codec metadata (Phase 3)
        if !self.codec_metadata.is_empty() {
            use super::varint::encode_varint;
            // Write codec count
            encode_varint(
                self.codec_metadata.len() as u64,
                &mut self.codec_metadata_buf,
            );
            // Write each codec metadata
            for metadata in &self.codec_metadata {
                let encoded = metadata.encode();
                encode_varint(encoded.len() as u64, &mut self.codec_metadata_buf);
                self.codec_metadata_buf.extend_from_slice(&encoded);
            }
        }

        // Determine mode
        let mode = if self.schemas.is_empty() {
            MODE_SELF_DESCRIBING
        } else {
            MODE_SCHEMA
        };

        // Determine flags
        let mut flags = FLAG_DICTIONARY;
        if !self.codec_metadata_buf.is_empty() {
            flags |= FLAG_CODEC_METADATA;
        }
        flags |= mode << 4;

        let mut result = Vec::with_capacity(
            8 + dict_buf.len()
                + self.schema_buf.len()
                + self.codec_metadata_buf.len()
                + self.buf.len(),
        );

        // Write header
        result.extend_from_slice(&TBF_MAGIC);
        result.push(TBF_VERSION);
        result.push(flags);
        result.extend_from_slice(&[0u8; 2]); // Reserved

        // Write dictionary
        result.extend_from_slice(&dict_buf);

        // Write schemas (if any)
        if mode == MODE_SCHEMA {
            result.extend_from_slice(&self.schema_buf);
        }

        // Write codec metadata (if any)
        if !self.codec_metadata_buf.is_empty() {
            result.extend_from_slice(&self.codec_metadata_buf);
        }

        // Write data
        result.extend_from_slice(&self.buf);

        // Write statistics footer (Phase 2)
        if let Some(stats) = self.stats
            && let Ok(stats_bytes) = stats.encode_all()
        {
            // Store offset to footer for random access
            let footer_offset = result.len() as u64;
            result.extend_from_slice(&stats_bytes);
            // Append footer offset (8 bytes, little-endian)
            result.extend_from_slice(&footer_offset.to_le_bytes());
        }

        result
    }

    /// Get reference to output buffer
    pub fn output(&self) -> &[u8] {
        &self.buf
    }

    /// Record codec metadata for a field/column (Phase 3)
    pub fn add_codec_metadata(&mut self, metadata: super::codec_encode::CodecMetadata) {
        self.codec_metadata.push(metadata);
    }

    /// Get the number of codec metadata entries collected
    pub fn codec_metadata_count(&self) -> usize {
        self.codec_metadata.len()
    }

    /// Write a type tag
    #[inline(always)]
    pub(crate) fn write_tag(&mut self, tag: TypeTag) {
        self.buf.push(tag as u8);
    }

    /// Write a varint
    #[inline(always)]
    pub(crate) fn write_varint(&mut self, value: u64) {
        encode_varint(value, &mut self.buf);
    }

    /// Write a signed varint
    #[inline(always)]
    pub(crate) fn write_signed_varint(&mut self, value: i64) {
        encode_signed_varint(value, &mut self.buf);
    }

    /// Intern a string and write its index
    #[inline]
    pub(crate) fn write_string(&mut self, s: &str) {
        let idx = self.dict.intern(s);
        encode_varint(idx as u64, &mut self.buf);
    }

    /// Write a value without type tag (schema mode)
    #[inline]
    pub(crate) fn write_typed_value_bool(&mut self, v: bool) {
        self.buf.push(if v { 1 } else { 0 });
    }

    #[inline]
    pub(crate) fn write_typed_value_int(&mut self, v: i64) {
        encode_signed_varint(v, &mut self.buf);
    }

    #[inline]
    pub(crate) fn write_typed_value_uint(&mut self, v: u64) {
        encode_varint(v, &mut self.buf);
    }

    #[inline]
    pub(crate) fn write_typed_value_f32(&mut self, v: f32) {
        self.buf.extend_from_slice(&v.to_le_bytes());
    }

    #[inline]
    pub(crate) fn write_typed_value_f64(&mut self, v: f64) {
        self.buf.extend_from_slice(&v.to_le_bytes());
    }

    #[inline]
    pub(crate) fn write_typed_value_string(&mut self, s: &str) {
        let idx = self.dict.intern(s);
        encode_varint(idx as u64, &mut self.buf);
    }

    /// Enter a nested structure
    #[inline]
    pub(crate) fn enter(&mut self) {
        self.depth += 1;
    }

    /// Leave a nested structure
    #[inline]
    pub(crate) fn leave(&mut self) {
        self.depth -= 1;
    }

    /// Check if at root level
    #[inline]
    pub(crate) fn is_root(&self) -> bool {
        self.depth == 0
    }

    /// Begin a sequence - returns a helper for schema detection
    pub(crate) fn begin_sequence(&mut self, len: Option<usize>) -> SequenceEncoder<'_> {
        SequenceEncoder::new(self, len)
    }

    /// Check if currently in schema mode for structs
    pub(crate) fn in_schema_mode(&self) -> bool {
        self.context.schema_mode && self.context.seq_schema.is_some()
    }

    /// Get expected field type for current position (schema mode)
    pub(crate) fn get_expected_field_type(&self, field_idx: usize) -> Option<SchemaType> {
        self.context
            .seq_schema
            .as_ref()
            .and_then(|s| s.types.get(field_idx).copied())
    }
}

impl Default for TbfSerializer {
    fn default() -> Self {
        Self::new()
    }
}

/// Helper for encoding sequences with schema detection
pub struct SequenceEncoder<'a> {
    serializer: &'a mut TbfSerializer,
    len: Option<usize>,
    element_count: usize,
    schema_detected: bool,
    first_element_buf: Vec<u8>,
    first_element_fields: Vec<(String, SchemaType)>,
}

impl<'a> SequenceEncoder<'a> {
    fn new(serializer: &'a mut TbfSerializer, len: Option<usize>) -> Self {
        Self {
            serializer,
            len,
            element_count: 0,
            schema_detected: false,
            first_element_buf: Vec::new(),
            first_element_fields: Vec::new(),
        }
    }

    /// Called before each element
    pub fn before_element(&mut self) {
        self.element_count += 1;
    }

    /// Called when a struct starts within the sequence
    pub fn struct_started(&mut self, field_count: usize) {
        if self.element_count == 1 {
            // First element - start collecting field info
            self.serializer.context.current_fields = Vec::with_capacity(field_count);
        }
    }

    /// Called when a struct field is serialized
    pub fn field_serialized(&mut self, name: &str, typ: SchemaType) {
        if self.element_count == 1 {
            self.serializer
                .context
                .current_fields
                .push((name.to_string(), typ));
        }
    }

    /// Called when first struct in sequence ends
    pub fn first_struct_ended(&mut self) {
        if self.element_count == 1 && !self.serializer.context.current_fields.is_empty() {
            // Create schema from collected fields
            let fields: Vec<String> = self
                .serializer
                .context
                .current_fields
                .iter()
                .map(|(n, _)| n.clone())
                .collect();
            let types: Vec<SchemaType> = self
                .serializer
                .context
                .current_fields
                .iter()
                .map(|(_, t)| *t)
                .collect();

            let detected = DetectedSchema { fields, types };
            self.serializer.context.seq_schema = Some(detected);
            self.serializer.context.schema_mode = true;
            self.schema_detected = true;
        }
    }

    /// Finalize the sequence encoder
    pub fn finish(self) {
        // Reset context
        self.serializer.context.in_sequence = false;
        self.serializer.context.seq_element_count = 0;
        self.serializer.context.seq_schema = None;
        self.serializer.context.schema_mode = false;
        self.serializer.context.current_fields.clear();
    }
}

/// Schema-aware struct serializer
pub struct SchemaStructSerializer<'a> {
    serializer: &'a mut TbfSerializer,
    field_idx: usize,
    schema_mode: bool,
    expected_types: Option<Vec<SchemaType>>,
}

impl<'a> SchemaStructSerializer<'a> {
    /// Create a new schema-aware struct serializer
    pub fn new(serializer: &'a mut TbfSerializer, schema_mode: bool) -> Self {
        let expected_types = if schema_mode {
            serializer
                .context
                .seq_schema
                .as_ref()
                .map(|s| s.types.clone())
        } else {
            None
        };

        Self {
            serializer,
            field_idx: 0,
            schema_mode,
            expected_types,
        }
    }

    /// Get expected type for current field
    pub fn expected_type(&self) -> Option<SchemaType> {
        self.expected_types
            .as_ref()
            .and_then(|types| types.get(self.field_idx).copied())
    }

    /// Move to next field
    pub fn next_field(&mut self) {
        self.field_idx += 1;
    }

    /// Check if in schema mode
    pub fn is_schema_mode(&self) -> bool {
        self.schema_mode
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::tbf::adaptive_encode::CompressionCodec;
    use serde_json::json;

    #[test]
    fn test_serializer_with_codecs_creation() {
        let serializer = TbfSerializer::with_codecs();
        assert!(serializer.codec_analyzer.is_some());
        assert!(serializer.stats.is_none());
        assert!(serializer.selected_codecs.is_empty());
    }

    #[test]
    fn test_serializer_with_codecs_and_statistics() {
        let serializer = TbfSerializer::with_codecs_and_statistics();
        assert!(serializer.codec_analyzer.is_some());
        assert!(serializer.stats.is_some());
        assert!(serializer.selected_codecs.is_empty());
    }

    #[test]
    fn test_serializer_without_codecs_unchanged() {
        let serializer = TbfSerializer::new();
        assert!(serializer.codec_analyzer.is_none());
        assert!(serializer.stats.is_none());
        assert!(serializer.selected_codecs.is_empty());
    }

    #[test]
    fn test_codec_analyzer_accessibility() {
        let mut serializer = TbfSerializer::with_codecs();

        // Verify analyzer exists and can be used
        if let Some(analyzer) = &mut serializer.codec_analyzer {
            // Add 20 sorted samples (enough for Delta detection)
            for i in 0..20 {
                analyzer.add_sample(Some(json!(i)));
            }

            let codec = analyzer.choose_codec();
            // Sorted integers should select Delta
            assert_eq!(codec, CompressionCodec::Delta);
        } else {
            panic!("codec_analyzer should be Some");
        }
    }

    #[test]
    fn test_codec_roundtrip_basic() {
        let _value = json!({
            "id": 1,
            "name": "test"
        });

        let serializer = TbfSerializer::new();
        let bytes = serializer.into_bytes();

        // Should produce some output
        assert!(!bytes.is_empty());
    }

    #[test]
    fn test_selected_codecs_storage() {
        let mut serializer = TbfSerializer::with_codecs();

        // Simulate selecting codecs for fields
        serializer
            .selected_codecs
            .insert("field1".to_string(), CompressionCodec::Dictionary);
        serializer
            .selected_codecs
            .insert("field2".to_string(), CompressionCodec::Delta);

        assert_eq!(serializer.selected_codecs.len(), 2);
        assert_eq!(
            serializer.selected_codecs.get("field1"),
            Some(&CompressionCodec::Dictionary)
        );
        assert_eq!(
            serializer.selected_codecs.get("field2"),
            Some(&CompressionCodec::Delta)
        );
    }

    #[test]
    fn test_codec_and_stats_together() {
        let serializer = TbfSerializer::with_codecs_and_statistics();

        // Both features should be enabled independently
        assert!(serializer.codec_analyzer.is_some());
        assert!(serializer.stats.is_some());

        // They should not interfere with each other
        assert_eq!(serializer.depth, 0);
        assert!(serializer.buf.is_empty());
    }

    #[test]
    fn test_codec_metadata_collection() {
        use crate::tbf::codec_encode::CodecMetadata;

        let mut serializer = TbfSerializer::new();

        // Add codec metadata
        serializer.add_codec_metadata(CodecMetadata::Delta { initial_value: 100 });
        serializer.add_codec_metadata(CodecMetadata::Dictionary {
            dictionary_size: 50,
        });

        assert_eq!(serializer.codec_metadata_count(), 2);
    }

    #[test]
    fn test_codec_metadata_binary_encoding() {
        use crate::tbf::codec_encode::CodecMetadata;

        let mut serializer = TbfSerializer::new();

        // Add metadata for a Delta codec
        serializer.add_codec_metadata(CodecMetadata::Delta { initial_value: 42 });

        let bytes = serializer.into_bytes();

        // Check header has codec metadata flag
        assert!(bytes.len() > 8);
        let flags = bytes[5];
        // FLAG_CODEC_METADATA = 0x04
        assert!(flags & 0x04 != 0, "Codec metadata flag should be set");
    }

    #[test]
    fn test_codec_metadata_format_section() {
        use crate::tbf::codec_encode::CodecMetadata;

        let mut serializer = TbfSerializer::with_codecs();

        // Add codec metadata
        serializer.add_codec_metadata(CodecMetadata::RLE);
        serializer.add_codec_metadata(CodecMetadata::Dictionary {
            dictionary_size: 100,
        });

        let bytes = serializer.into_bytes();

        // Verify structure: [header][dict][codecs][data][footer?]
        assert!(bytes.len() > 8);
        // Verify codec metadata flag is set
        let flags = bytes[5];
        assert!(flags & 0x04 != 0);
    }

    #[test]
    fn test_codec_metadata_with_statistics() {
        use crate::tbf::codec_encode::CodecMetadata;

        let mut serializer = TbfSerializer::with_codecs_and_statistics();

        serializer.add_codec_metadata(CodecMetadata::Delta { initial_value: 0 });

        let bytes = serializer.into_bytes();

        // Should have both codec metadata and stats footer
        assert!(bytes.len() > 8);
        let flags = bytes[5];
        // Should have codec metadata flag set (0x04)
        assert!(flags & 0x04 != 0);
    }

    #[test]
    fn test_no_codec_metadata_no_flag() {
        let serializer = TbfSerializer::new();
        let bytes = serializer.into_bytes();

        // Check header does NOT have codec metadata flag
        let flags = bytes[5];
        // FLAG_CODEC_METADATA = 0x04 should not be set
        assert!(
            flags & 0x04 == 0,
            "Codec metadata flag should not be set when no codecs"
        );
    }

    #[test]
    fn test_multiple_codec_metadata_entries() {
        use crate::tbf::codec_encode::CodecMetadata;

        let mut serializer = TbfSerializer::new();

        // Add multiple codec metadata entries
        for i in 0..5 {
            serializer.add_codec_metadata(CodecMetadata::Delta {
                initial_value: i as i64,
            });
        }

        assert_eq!(serializer.codec_metadata_count(), 5);

        let bytes = serializer.into_bytes();

        // Verify codec metadata flag is set
        let flags = bytes[5];
        assert!(flags & 0x04 != 0);
    }
}