#![cfg(feature = "disabled_tests")]

use anyhow::Result;
use arrow::array::{Array, ArrayRef, Float32Array, Int32Array};
use arrow::buffer::Buffer;
use arrow::datatypes::{DataType, Field, Schema};
use arrow::record_batch::RecordBatch;
use rtx_compress::arrow_integration::{
    ArrowCompressor, ArrowConfig, CompressedArrowArray, ZeroCopyBuffer,
};
use rtx_tensor::{DType, Device, Tensor};
use std::sync::Arc;

#[test]
fn test_arrow_tensor_zero_copy_compression() -> Result<()> {
    let device = Device::try_default()?;

    // Create tensor data
    let tensor_data = Tensor::randn(&[1000, 512], &device)?;

    // Convert to Arrow array (zero-copy)
    let arrow_array = tensor_to_arrow_array(&tensor_data)?;

    let config = ArrowConfig {
        compression_method: "zstd".to_string(),
        compression_level: 6,
        enable_quantization: true,
        quantization_bits: 8,
        preserve_nulls: true,
    };

    let compressor = ArrowCompressor::new(config);

    // Compress Arrow array
    let compressed = compressor.compress_array(&arrow_array)?;

    // Verify compression metrics
    let original_size = arrow_array.get_buffer_memory_size();
    let compressed_size = compressed.compressed_size();
    let compression_ratio = original_size as f64 / compressed_size as f64;

    assert!(
        compression_ratio >= 2.0,
        "Should achieve at least 2x compression, got {:.2}",
        compression_ratio
    );

    // Decompress back to Arrow array (zero-copy)
    let decompressed_array = compressor.decompress_array(&compressed)?;

    // Verify zero-copy: should share memory when possible
    assert!(
        compressed.is_zero_copy(),
        "Should support zero-copy decompression"
    );

    // Convert back to tensor and verify reconstruction
    let reconstructed_tensor = arrow_array_to_tensor(&decompressed_array, &device)?;

    let mse = (&tensor_data - &reconstructed_tensor)?
        .pow(2.0)?
        .mean(&[], false)?
        .to_scalar::<f32>()?;
    assert!(
        mse < 0.01,
        "Reconstruction MSE should be < 0.01, got {}",
        mse
    );

    Ok(())
}

#[test]
fn test_arrow_record_batch_compression() -> Result<()> {
    // Create sample RecordBatch (simulating RAG embeddings + metadata)
    let embeddings: ArrayRef = Arc::new(Float32Array::from(
        (0..5000).map(|_| rand::random::<f32>()).collect::<Vec<_>>(),
    ));
    let doc_ids: ArrayRef = Arc::new(Int32Array::from(
        (0..5000).map(|i| i % 100).collect::<Vec<_>>(),
    ));

    let schema = Arc::new(Schema::new(vec![
        Field::new("embeddings", DataType::Float32, false),
        Field::new("doc_ids", DataType::Int32, false),
    ]));

    let batch = RecordBatch::try_new(schema.clone(), vec![embeddings.clone(), doc_ids.clone()])?;

    let config = ArrowConfig {
        compression_method: "lz4".to_string(),
        compression_level: 1,
        enable_quantization: true,
        quantization_bits: 16, // fp16 for embeddings
        preserve_nulls: false,
    };

    let compressor = ArrowCompressor::new(config);

    // Compress entire record batch
    let compressed_batch = compressor.compress_record_batch(&batch)?;

    // Verify selective compression
    // Embeddings should be compressed/quantized, doc_ids should be uncompressed
    assert!(
        compressed_batch.is_column_compressed(0), // embeddings
        "Embeddings column should be compressed"
    );
    assert!(
        !compressed_batch.is_column_compressed(1), // doc_ids
        "Integer columns should not be compressed"
    );

    // Decompress and verify
    let decompressed_batch = compressor.decompress_record_batch(&compressed_batch)?;

    assert_eq!(decompressed_batch.num_rows(), 5000);
    assert_eq!(decompressed_batch.num_columns(), 2);

    // Verify doc_ids are identical (no quantization)
    let original_doc_ids = batch
        .column(1)
        .as_any()
        .downcast_ref::<Int32Array>()
        .unwrap();
    let decompressed_doc_ids = decompressed_batch
        .column(1)
        .as_any()
        .downcast_ref::<Int32Array>()
        .unwrap();

    for i in 0..5000 {
        assert_eq!(original_doc_ids.value(i), decompressed_doc_ids.value(i));
    }

    Ok(())
}

#[test]
fn test_arrow_streaming_compression() -> Result<()> {
    let config = ArrowConfig {
        compression_method: "zstd".to_string(),
        compression_level: 3,
        enable_quantization: true,
        quantization_bits: 8,
        preserve_nulls: true,
    };

    let mut compressor = ArrowCompressor::new(config);

    // Create streaming compressor
    let mut stream_compressor = compressor.create_streaming_compressor()?;

    // Stream multiple batches (simulating large dataset)
    for batch_idx in 0..10 {
        let data: ArrayRef = Arc::new(Float32Array::from(
            (0..1000)
                .map(|i| (batch_idx * 1000 + i) as f32)
                .collect::<Vec<_>>(),
        ));

        let schema = Arc::new(Schema::new(vec![Field::new(
            "values",
            DataType::Float32,
            false,
        )]));

        let batch = RecordBatch::try_new(schema, vec![data])?;

        // Add batch to stream
        stream_compressor.add_batch(&batch)?;
    }

    // Finalize stream compression
    let compressed_stream = stream_compressor.finalize()?;

    // Verify stream metadata
    assert_eq!(compressed_stream.num_batches(), 10);
    assert_eq!(compressed_stream.total_rows(), 10000);

    // Create streaming decompressor
    let mut stream_decompressor = compressor.create_streaming_decompressor(&compressed_stream)?;

    // Read back batches
    let mut total_decompressed_rows = 0;
    while let Some(batch) = stream_decompressor.next_batch()? {
        assert_eq!(batch.num_rows(), 1000);
        total_decompressed_rows += batch.num_rows();
    }

    assert_eq!(total_decompressed_rows, 10000);

    Ok(())
}

#[test]
fn test_arrow_memory_mapped_compression() -> Result<()> {
    let device = Device::try_default()?;

    // Create large tensor data
    let large_tensor = Tensor::randn(&[10000, 256], &device)?;
    let arrow_array = tensor_to_arrow_array(&large_tensor)?;

    let config = ArrowConfig {
        compression_method: "lz4".to_string(),
        compression_level: 1,
        enable_quantization: false, // Test pure compression
        quantization_bits: 32,
        preserve_nulls: false,
    };

    let compressor = ArrowCompressor::new(config);

    // Compress with memory mapping enabled
    let compressed = compressor.compress_array_mmap(&arrow_array)?;

    // Verify memory mapping properties
    assert!(
        compressed.is_memory_mapped(),
        "Should support memory mapped compression"
    );
    assert!(
        compressed.supports_random_access(),
        "Should support random access to compressed data"
    );

    // Test random access decompression
    let chunk_start = 5000;
    let chunk_size = 1000;

    let chunk = compressor.decompress_chunk(&compressed, chunk_start, chunk_size)?;

    // Verify chunk properties
    assert_eq!(chunk.len(), chunk_size);

    // Convert chunk to tensor and verify content
    let chunk_tensor = arrow_array_to_tensor(&chunk, &device)?;
    assert_eq!(chunk_tensor.shape(), &[chunk_size, 256]);

    Ok(())
}

#[test]
fn test_arrow_schema_preservation() -> Result<()> {
    // Create complex schema with metadata
    let mut metadata = std::collections::HashMap::new();
    metadata.insert("model_version".to_string(), "v1.0".to_string());
    metadata.insert("embedding_dim".to_string(), "768".to_string());

    let schema = Arc::new(Schema::new_with_metadata(
        vec![
            Field::new("embeddings", DataType::Float32, false),
            Field::new("attention_mask", DataType::Int32, true), // nullable
            Field::new("token_ids", DataType::Int32, false),
        ],
        metadata,
    ));

    // Create data with nulls in attention_mask
    let embeddings: ArrayRef = Arc::new(Float32Array::from(vec![1.0, 2.0, 3.0, 4.0]));
    let attention_mask: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), None, Some(1), None]));
    let token_ids: ArrayRef = Arc::new(Int32Array::from(vec![100, 200, 300, 400]));

    let batch = RecordBatch::try_new(schema.clone(), vec![embeddings, attention_mask, token_ids])?;

    let config = ArrowConfig {
        compression_method: "zstd".to_string(),
        compression_level: 4,
        enable_quantization: true,
        quantization_bits: 16,
        preserve_nulls: true, // Important for nullable fields
    };

    let compressor = ArrowCompressor::new(config);
    let compressed = compressor.compress_record_batch(&batch)?;
    let decompressed = compressor.decompress_record_batch(&compressed)?;

    // Verify schema is preserved exactly
    assert_eq!(decompressed.schema().metadata(), schema.metadata());
    assert_eq!(decompressed.schema().fields().len(), 3);

    // Verify nullable field nulls are preserved
    let decompressed_mask = decompressed
        .column(1)
        .as_any()
        .downcast_ref::<Int32Array>()
        .unwrap();

    assert!(decompressed_mask.is_null(1));
    assert!(decompressed_mask.is_null(3));
    assert!(!decompressed_mask.is_null(0));
    assert!(!decompressed_mask.is_null(2));

    Ok(())
}

#[test]
fn test_arrow_quantization_quality_levels() -> Result<()> {
    let device = Device::try_default()?;

    // Create test tensor with known patterns
    let mut data = vec![];
    for i in 0..1000 {
        data.push((i as f32 * 0.001).sin()); // Smooth signal
    }
    let tensor = Tensor::from_slice(&data, &[1000], &device)?;
    let arrow_array = tensor_to_arrow_array(&tensor)?;

    // Test different quantization levels
    let quantization_levels = vec![4, 8, 12, 16];
    let mut results = vec![];

    for bits in quantization_levels {
        let config = ArrowConfig {
            compression_method: "none".to_string(),
            compression_level: 0,
            enable_quantization: true,
            quantization_bits: bits,
            preserve_nulls: false,
        };

        let compressor = ArrowCompressor::new(config);
        let compressed = compressor.compress_array(&arrow_array)?;
        let decompressed = compressor.decompress_array(&compressed)?;

        // Calculate reconstruction quality
        let reconstructed_tensor = arrow_array_to_tensor(&decompressed, &device)?;
        let mse = (&tensor - &reconstructed_tensor)?
            .pow(2.0)?
            .mean(&[], false)?
            .to_scalar::<f32>()?;

        results.push((bits, mse));

        println!("Quantization {}bit: MSE = {:.6}", bits, mse);
    }

    // Verify quality improves with more bits
    for i in 1..results.len() {
        assert!(
            results[i].1 <= results[i - 1].1 * 1.1, // Allow some noise
            "Higher bit quantization should have lower error"
        );
    }

    // 16-bit should be nearly perfect
    assert!(
        results.last().unwrap().1 < 1e-4,
        "16-bit quantization should be very accurate"
    );

    Ok(())
}

#[test]
fn test_arrow_rag_integration_workflow() -> Result<()> {
    let device = Device::try_default()?;

    // Simulate RAG pipeline: Documents -> Embeddings -> Compressed Storage -> Retrieval

    // Step 1: Create document embeddings
    let num_docs = 10000;
    let embedding_dim = 384;

    let embeddings = Tensor::randn(&[num_docs, embedding_dim], &device)?;
    let doc_scores = Tensor::rand(&[num_docs], &device)?; // Relevance scores

    // Step 2: Convert to Arrow format
    let embedding_array = tensor_to_arrow_array(&embeddings)?;
    let scores_array = tensor_to_arrow_array(&doc_scores)?;

    let schema = Arc::new(Schema::new(vec![
        Field::new("embeddings", DataType::Float32, false),
        Field::new("scores", DataType::Float32, false),
    ]));

    let batch = RecordBatch::try_new(schema, vec![embedding_array, scores_array])?;

    // Step 3: Compress for storage (aggressive compression for embeddings)
    let config = ArrowConfig {
        compression_method: "zstd".to_string(),
        compression_level: 9, // High compression for storage
        enable_quantization: true,
        quantization_bits: 8, // 8-bit embeddings
        preserve_nulls: false,
    };

    let compressor = ArrowCompressor::new(config);
    let compressed_storage = compressor.compress_record_batch(&batch)?;

    // Step 4: Simulate query-time retrieval with partial decompression
    let query_embedding = Tensor::randn(&[1, embedding_dim], &device)?;

    // Retrieve top-k documents without full decompression
    let top_k = 100;
    let relevant_doc_indices = compressor.search_compressed(
        &compressed_storage,
        &tensor_to_arrow_array(&query_embedding)?,
        top_k,
        "cosine_similarity",
    )?;

    assert_eq!(relevant_doc_indices.len(), top_k);

    // Step 5: Decompress only relevant documents
    let relevant_batch =
        compressor.decompress_selective(&compressed_storage, &relevant_doc_indices)?;

    assert_eq!(relevant_batch.num_rows(), top_k);
    assert_eq!(relevant_batch.num_columns(), 2);

    // Step 6: Verify retrieval quality
    let retrieved_embeddings = arrow_array_to_tensor(relevant_batch.column(0), &device)?;
    assert_eq!(retrieved_embeddings.shape(), &[top_k, embedding_dim]);

    // Compute actual similarities to verify retrieval correctness
    let similarities = compute_cosine_similarity(&query_embedding, &retrieved_embeddings)?;

    // Similarities should be in descending order (top-k retrieval)
    for i in 1..similarities.len() {
        let curr_sim = similarities[i - 1];
        let next_sim = similarities[i];
        assert!(
            curr_sim >= next_sim - 0.01, // Allow small numerical errors
            "Retrieved documents should be ordered by similarity"
        );
    }

    Ok(())
}

// Helper functions
fn tensor_to_arrow_array(tensor: &Tensor) -> Result<ArrayRef> {
    // This would need proper implementation based on tensor data layout
    // For now, create a placeholder Float32Array
    let data: Vec<f32> = (0..tensor.numel()).map(|_| rand::random::<f32>()).collect();
    Ok(Arc::new(Float32Array::from(data)))
}

fn arrow_array_to_tensor(array: &ArrayRef, device: &Device) -> Result<Tensor> {
    // This would need proper implementation based on Arrow array type
    let float_array = array.as_any().downcast_ref::<Float32Array>().unwrap();

    let data: Vec<f32> = (0..float_array.len())
        .map(|i| float_array.value(i))
        .collect();

    Tensor::from_slice(&data, &[float_array.len()], device)
}

fn compute_cosine_similarity(query: &Tensor, docs: &Tensor) -> Result<Vec<f32>> {
    // Simplified cosine similarity computation
    let query_norm = query.pow(2.0)?.sum(None)?.sqrt()?;
    let doc_norms = docs.pow(2.0)?.sum(&[-1])?.sqrt()?;

    let dots = (query.unsqueeze(0)? * docs)?.sum(&[-1])?;
    let similarities = dots / (query_norm * doc_norms)?;

    let mut result = vec![];
    for i in 0..docs.shape()[0] {
        result.push(similarities.get(&[i])?.to_scalar::<f32>()?);
    }

    Ok(result)
}