use crate::__private::S3SinkConfig;
use anyhow::Context;
use bytes::Bytes;
use futures::TryStreamExt;
use object_store::aws::AmazonS3Builder;
use object_store::path::Path as ObjectPath;
use object_store::{ObjectStore, ObjectStoreExt};
use parquet::arrow::ArrowWriter;
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
use parquet::basic::Compression;
use parquet::file::properties::WriterProperties;
use std::collections::BTreeMap;
use std::fs::File;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tokio::runtime::Builder as RuntimeBuilder;
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
use tokio::task::JoinSet;

const DEFAULT_MAX_CONCURRENT_DOWNLOADS: usize = 128;

#[derive(Debug, Clone, Default)]
pub struct AnonymizeConfig {
    pub columns_to_pseudonymize: Vec<String>,
    pub columns_to_drop: Vec<String>,
}

#[derive(Debug, Clone)]
pub struct S3ParquetCompactConfig {
    pub s3: S3SinkConfig,
    pub output_path: PathBuf,
    pub max_concurrent_downloads: usize,
    pub anonymize: Option<AnonymizeConfig>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ParquetCompactResult {
    pub source_objects: usize,
    pub rows_written: usize,
    pub output_path: PathBuf,
    pub min_timestamp: u64,
}

pub fn compact_s3_parquet_to_local(
    config: S3ParquetCompactConfig,
) -> anyhow::Result<ParquetCompactResult> {
    let rt = RuntimeBuilder::new_current_thread()
        .enable_all()
        .build()
        .context("failed to build compaction runtime")?;
    rt.block_on(compact_s3_parquet_to_local_async(config))
}

pub async fn compact_s3_parquet_to_local_async(
    config: S3ParquetCompactConfig,
) -> anyhow::Result<ParquetCompactResult> {
    let start = std::time::Instant::now();
    let store = build_s3_store(&config.s3)?;
    let max_concurrent_downloads = match config.max_concurrent_downloads {
        0 => DEFAULT_MAX_CONCURRENT_DOWNLOADS,
        value => value,
    };

    if let Some(ref anon_config) = config.anonymize {
        let temp_path = config.output_path.with_extension("tmp.parquet");
        let result = compact_store_parquet_prefix_to_local(
            store,
            config.s3.prefix.as_str(),
            temp_path.as_path(),
            max_concurrent_downloads,
        )
        .await?;

        let columns_to_pseudonymize_refs: Vec<&str> = anon_config
            .columns_to_pseudonymize
            .iter()
            .map(|s| s.as_str())
            .collect();
        let columns_to_drop_refs: Vec<&str> = anon_config
            .columns_to_drop
            .iter()
            .map(|s| s.as_str())
            .collect();

        let min_ts = crate::anonymize_parquet_file(
            &temp_path,
            &config.output_path,
            &columns_to_pseudonymize_refs,
            &columns_to_drop_refs,
            result.min_timestamp,
        )?;

        std::fs::remove_file(&temp_path)?;

        let duration_ms = start.elapsed().as_millis();
        println!("Duration: {}ms", duration_ms);
        Ok(ParquetCompactResult {
            source_objects: result.source_objects,
            rows_written: result.rows_written,
            output_path: config.output_path,
            min_timestamp: min_ts,
        })
    } else {
        let result = compact_store_parquet_prefix_to_local(
            store,
            config.s3.prefix.as_str(),
            config.output_path.as_path(),
            max_concurrent_downloads,
        )
        .await?;
        let duration_ms = start.elapsed().as_millis();
        println!("Duration: {}ms", duration_ms);
        Ok(result)
    }
}

async fn compact_store_parquet_prefix_to_local(
    store: Arc<dyn ObjectStore>,
    prefix: &str,
    output_path: &Path,
    max_concurrent_downloads: usize,
) -> anyhow::Result<ParquetCompactResult> {
    let prefix = prefix.trim_matches('/');
    let prefix_path = if prefix.is_empty() {
        None
    } else {
        Some(ObjectPath::from(prefix))
    };

    let mut objects = store
        .list(prefix_path.as_ref())
        .try_collect::<Vec<_>>()
        .await
        .context("failed to list parquet objects")?;
    objects.sort_by(|a, b| a.location.cmp(&b.location));
    objects.retain(|meta| meta.location.as_ref().ends_with(".parquet"));

    if objects.is_empty() {
        anyhow::bail!("no parquet objects found under prefix '{prefix}'");
    }

    if let Some(parent) = output_path.parent()
        && !parent.as_os_str().is_empty()
    {
        std::fs::create_dir_all(parent)
            .with_context(|| format!("failed to create output directory {}", parent.display()))?;
    }

    let semaphore = Arc::new(Semaphore::new(max_concurrent_downloads));
    let mut join_set = JoinSet::new();
    let mut ready = BTreeMap::new();
    let mut next_spawn_index = 0usize;
    let mut next_write_index = 0usize;
    let mut in_flight = 0usize;
    let start_time = std::time::Instant::now();

    let mut writer = None;
    let mut expected_schema = None;
    let mut rows_written = 0usize;
    let mut min_timestamp: Option<u64> = None;

    while next_write_index < objects.len() {
        while next_spawn_index < objects.len() && in_flight < max_concurrent_downloads {
            let permit = semaphore
                .clone()
                .acquire_owned()
                .await
                .context("compaction download semaphore closed unexpectedly")?;
            let object = objects[next_spawn_index].clone();
            let store = Arc::clone(&store);
            join_set.spawn(async move {
                fetch_object_payload(store, next_spawn_index, object.location, permit).await
            });
            next_spawn_index += 1;
            in_flight += 1;
        }

        let Some(result) = join_set.join_next().await else {
            anyhow::bail!("download queue terminated before all parquet objects were processed");
        };
        in_flight = in_flight.saturating_sub(1);

        let fetched = match result {
            Ok(Ok(fetched)) => fetched,
            Ok(Err(err)) => {
                join_set.abort_all();
                return Err(err);
            }
            Err(err) => {
                join_set.abort_all();
                return Err(anyhow::anyhow!("compaction download task failed: {err}"));
            }
        };
        ready.insert(fetched.index, fetched);

        while let Some(fetched) = ready.remove(&next_write_index) {
            let file_min_ts = append_parquet_payload(
                fetched.location.as_ref(),
                fetched.payload,
                output_path,
                &mut writer,
                &mut expected_schema,
                &mut rows_written,
            )?;
            if let Some(ts) = file_min_ts {
                min_timestamp = Some(min_timestamp.map_or(ts, |min| min.min(ts)));
            }
            next_write_index += 1;
            if next_write_index.is_multiple_of(10) {
                let elapsed = start_time.elapsed().as_secs_f32();
                eprintln!(
                    "Processed {}/{} files ({:.1}s)",
                    next_write_index,
                    objects.len(),
                    elapsed
                );
            }
        }
    }

    writer
        .context("writer should be initialized when objects are present")?
        .close()
        .with_context(|| format!("failed to finalize {}", output_path.display()))?;

    Ok(ParquetCompactResult {
        source_objects: objects.len(),
        rows_written,
        output_path: output_path.to_path_buf(),
        min_timestamp: min_timestamp.unwrap_or(0),
    })
}

struct FetchedObject {
    index: usize,
    location: ObjectPath,
    payload: Bytes,
}

async fn fetch_object_payload(
    store: Arc<dyn ObjectStore>,
    index: usize,
    location: ObjectPath,
    _permit: OwnedSemaphorePermit,
) -> anyhow::Result<FetchedObject> {
    let payload = store
        .get(&location)
        .await
        .with_context(|| format!("failed to fetch {location}"))?
        .bytes()
        .await
        .with_context(|| format!("failed to read {location}"))?;

    Ok(FetchedObject {
        index,
        location,
        payload,
    })
}

fn append_parquet_payload(
    location: &str,
    payload: Bytes,
    output_path: &Path,
    writer: &mut Option<ArrowWriter<File>>,
    expected_schema: &mut Option<arrow::datatypes::SchemaRef>,
    rows_written: &mut usize,
) -> anyhow::Result<Option<u64>> {
    use arrow::array::UInt64Array;

    let builder = ParquetRecordBatchReaderBuilder::try_new(payload)
        .with_context(|| format!("failed to open parquet {location}"))?;
    let schema = builder.schema().clone();

    if let Some(existing) = expected_schema.as_ref() {
        if existing.as_ref() != schema.as_ref() {
            anyhow::bail!(
                "schema mismatch for {} while compacting prefix '{}'",
                location,
                output_path.display()
            );
        }
    } else {
        let file = File::create(output_path)
            .with_context(|| format!("failed to create {}", output_path.display()))?;
        let props = WriterProperties::builder()
            .set_compression(Compression::SNAPPY)
            .build();
        *writer = Some(
            ArrowWriter::try_new(file, schema.clone(), Some(props))
                .with_context(|| format!("failed to open {}", output_path.display()))?,
        );
        *expected_schema = Some(schema.clone());
    }

    let record_reader = builder
        .build()
        .with_context(|| format!("failed to build batch reader for {location}"))?;

    let mut min_timestamp: Option<u64> = None;
    for batch in record_reader {
        let batch = batch.with_context(|| format!("failed to read batch from {location}"))?;

        // Track min timestamp
        if let Ok(idx) = batch.schema().index_of("timestamp_start_unix_ms")
            && let Some(arr) = batch.column(idx).as_any().downcast_ref::<UInt64Array>()
        {
            for ts in arr.iter().flatten() {
                min_timestamp = Some(min_timestamp.map_or(ts, |min| min.min(ts)));
            }
        }

        *rows_written += batch.num_rows();
        writer
            .as_mut()
            .expect("writer should be initialized from first schema")
            .write(&batch)
            .with_context(|| format!("failed to write batch from {location}"))?;
    }
    Ok(min_timestamp)
}

fn build_s3_store(config: &S3SinkConfig) -> anyhow::Result<Arc<dyn ObjectStore>> {
    let mut builder = AmazonS3Builder::new()
        .with_bucket_name(config.bucket.clone())
        .with_region(config.region.clone());
    if let Some(endpoint) = &config.endpoint {
        builder = builder.with_endpoint(endpoint.clone());
    }
    if let Some(access_key_id) = &config.access_key_id {
        builder = builder.with_access_key_id(resolve_credential(access_key_id));
    }
    if let Some(secret_access_key) = &config.secret_access_key {
        builder = builder.with_secret_access_key(resolve_credential(secret_access_key));
    }
    let store = builder
        .build()
        .map_err(|e| anyhow::anyhow!("failed to build s3 store: {e}"))?;
    Ok(Arc::new(store))
}

fn read_credential_from_path(path: &str) -> anyhow::Result<String> {
    let content = std::fs::read_to_string(path)
        .map_err(|e| anyhow::anyhow!("failed to read credential from {path}: {e}"))?;
    Ok(content.trim().to_string())
}

fn resolve_credential(value: &str) -> String {
    if value.starts_with('/') && Path::new(value).exists() {
        match read_credential_from_path(value) {
            Ok(credential) => credential,
            Err(_) => value.to_string(),
        }
    } else {
        value.to_string()
    }
}

/// Anonymize a parquet file by offsetting timestamps and pseudonymizing sensitive columns.
///
/// # Arguments
/// * `input_path` - Path to input parquet file
/// * `output_path` - Path to output parquet file  
/// * `columns_to_pseudonymize` - List of column names to pseudonymize (e.g., ["provided_session_id", "org_id", "total_hashes"])
/// * `columns_to_drop` - List of column names to drop entirely
/// * `min_timestamp` - Minimum timestamp used for offsetting. If 0, timestamps are not shifted.
///
/// # Behavior
/// - Offsets timestamps by min_timestamp when min_timestamp > 0
/// - Leaves timestamps unchanged when min_timestamp == 0
/// - Assigns unique IDs to specified columns using a HashMap
/// - Drops specified columns (unless they're also in columns_to_pseudonymize)
/// - total_hashes column is treated specially as List<FixedSizeBinary(32)>
pub fn anonymize_parquet_file(
    input_path: &Path,
    output_path: &Path,
    columns_to_pseudonymize: &[&str],
    columns_to_drop: &[&str],
    min_timestamp: u64,
) -> anyhow::Result<u64> {
    use arrow::array::{
        Array, BinaryArray, FixedSizeBinaryArray, ListArray, StringArray, UInt64Array,
    };
    use arrow::record_batch::RecordBatch;
    use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
    use std::collections::{HashMap, HashSet};

    // HashMaps for assigning unique IDs - one per column
    let mut id_maps: HashMap<String, HashMap<Vec<u8>, u64>> = HashMap::new();
    for col in columns_to_pseudonymize {
        id_maps.insert(col.to_string(), HashMap::new());
    }

    // Convert to HashSets for faster lookup
    let columns_to_pseudonymize_set: HashSet<&str> =
        columns_to_pseudonymize.iter().copied().collect();
    let columns_to_drop_set: HashSet<&str> = columns_to_drop.iter().copied().collect();

    // Read input parquet
    let file = std::fs::File::open(input_path)?;
    let reader = ParquetRecordBatchReaderBuilder::try_new(file)?.build()?;

    // Use provided min_timestamp (0 means no offsetting)
    let min_ts = min_timestamp;

    // Process and write in streaming fashion
    let out_file = File::create(output_path)?;
    let props = WriterProperties::builder()
        .set_compression(Compression::SNAPPY)
        .build();

    let mut writer: Option<ArrowWriter<File>> = None;

    for batch in reader {
        let batch = batch?;
        let schema = batch.schema();
        let mut columns: Vec<std::sync::Arc<dyn Array>> = Vec::with_capacity(batch.num_columns());
        let mut fields: Vec<arrow::datatypes::Field> = Vec::new();

        for (idx, field) in schema.fields().iter().enumerate() {
            let column = batch.column(idx);
            let col_name = field.name();

            // Handle timestamp offset
            if col_name == "timestamp_start_unix_ms" && min_ts > 0 {
                let arr = column.as_any().downcast_ref::<UInt64Array>().unwrap();
                let offset_values: Vec<u64> = arr
                    .iter()
                    .map(|v| v.unwrap_or(0).saturating_sub(min_ts))
                    .collect();
                columns.push(std::sync::Arc::new(UInt64Array::from(offset_values)));
                fields.push((**field).clone());
            }
            // Handle columns to pseudonymize
            else if columns_to_pseudonymize_set.contains(col_name.as_str()) {
                if col_name == "total_hashes" {
                    // Handle list of binary hashes - assign unique IDs
                    let list_arr = column.as_any().downcast_ref::<ListArray>().unwrap();
                    let id_map = id_maps.get_mut("total_hashes").unwrap();
                    let mut new_lists: Vec<Option<Vec<u64>>> = Vec::new();

                    for i in 0..list_arr.len() {
                        if list_arr.is_null(i) {
                            new_lists.push(None);
                        } else {
                            let values = list_arr.value(i);
                            let mut ids = Vec::new();

                            // Try FixedSizeBinaryArray first (32-byte fixed), then BinaryArray
                            if let Some(fixed_arr) =
                                values.as_any().downcast_ref::<FixedSizeBinaryArray>()
                            {
                                for j in 0..fixed_arr.len() {
                                    if !fixed_arr.is_null(j) {
                                        let hash_bytes = fixed_arr.value(j).to_vec();
                                        let len = id_map.len() as u64;
                                        let id =
                                            *id_map.entry(hash_bytes).or_insert_with(|| len + 1);
                                        ids.push(id);
                                    }
                                }
                            } else if let Some(binary_arr) =
                                values.as_any().downcast_ref::<BinaryArray>()
                            {
                                for j in 0..binary_arr.len() {
                                    if !binary_arr.is_null(j) {
                                        let hash_bytes = binary_arr.value(j).to_vec();
                                        let len = id_map.len() as u64;
                                        let id =
                                            *id_map.entry(hash_bytes).or_insert_with(|| len + 1);
                                        ids.push(id);
                                    }
                                }
                            }
                            new_lists.push(Some(ids));
                        }
                    }

                    // Build new List<UInt64> for the values
                    let values_builder = arrow::array::builder::UInt64Builder::new();
                    let mut list_builder = arrow::array::builder::ListBuilder::new(values_builder);

                    for list in new_lists {
                        if let Some(ids) = list {
                            for id in ids {
                                list_builder.values().append_value(id);
                            }
                            list_builder.append(true);
                        } else {
                            list_builder.append(false);
                        }
                    }

                    columns.push(std::sync::Arc::new(list_builder.finish()));
                    // Update field type to List<UInt64>
                    let new_field = arrow::datatypes::Field::new(
                        field.name().clone(),
                        arrow::datatypes::DataType::List(std::sync::Arc::new(
                            arrow::datatypes::Field::new(
                                "item",
                                arrow::datatypes::DataType::UInt64,
                                true,
                            ),
                        )),
                        field.is_nullable(),
                    );
                    fields.push(new_field);
                } else {
                    // Handle string columns - assign unique IDs
                    let arr = column
                        .as_any()
                        .downcast_ref::<StringArray>()
                        .ok_or_else(|| anyhow::anyhow!("{} is not StringArray", col_name))?;
                    let id_map = id_maps.get_mut(col_name.as_str()).unwrap();
                    let id_values: Vec<Option<String>> = (0..arr.len())
                        .map(|i| {
                            if arr.is_null(i) {
                                None
                            } else {
                                let value = arr.value(i);
                                let bytes = value.as_bytes().to_vec();
                                let len = id_map.len() as u64;
                                let id = *id_map.entry(bytes).or_insert_with(|| len + 1);
                                Some(id.to_string())
                            }
                        })
                        .collect();
                    columns.push(std::sync::Arc::new(StringArray::from(id_values)));
                    fields.push((**field).clone());
                }
            }
            // Handle columns to drop (only those not pseudonymized)
            else if columns_to_drop_set.contains(col_name.as_str()) {
                // Skip this column
                continue;
            }
            // Pass through other columns unchanged
            else {
                columns.push(column.clone());
                fields.push((**field).clone());
            }
        }

        // Create new schema with only the fields we kept
        let new_schema = arrow::datatypes::Schema::new(fields);
        let new_batch = RecordBatch::try_new(std::sync::Arc::new(new_schema), columns)?;

        if writer.is_none() {
            writer = Some(ArrowWriter::try_new(
                out_file.try_clone()?,
                new_batch.schema(),
                Some(props.clone()),
            )?);
        }

        writer
            .as_mut()
            .expect("writer should be initialized")
            .write(&new_batch)?;
    }

    if let Some(w) = writer {
        w.close()?;
    }

    println!("Anonymization complete.");
    Ok(min_ts)
}

#[cfg(test)]
mod tests {
    use super::*;
    use arrow::array::{Array, Int64Array, StringArray};
    use arrow::datatypes::{DataType, Field, Schema};
    use arrow::record_batch::RecordBatch;
    use bytes::Bytes;
    use object_store::memory::InMemory;
    use object_store::{PutOptions, PutPayload};
    use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
    use std::sync::Arc;
    use std::time::{SystemTime, UNIX_EPOCH};

    fn make_test_parquet(rows: &[(&str, i64)]) -> Vec<u8> {
        let schema = Arc::new(Schema::new(vec![
            Field::new("name", DataType::Utf8, false),
            Field::new("value", DataType::Int64, false),
        ]));
        let names = rows
            .iter()
            .map(|(name, _)| (*name).to_string())
            .collect::<Vec<_>>();
        let values = rows.iter().map(|(_, value)| *value).collect::<Vec<_>>();
        let batch = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(StringArray::from(names)),
                Arc::new(Int64Array::from(values)),
            ],
        )
        .unwrap();

        let mut buf = Vec::new();
        let props = WriterProperties::builder()
            .set_compression(Compression::SNAPPY)
            .build();
        let mut writer = ArrowWriter::try_new(&mut buf, schema, Some(props)).unwrap();
        writer.write(&batch).unwrap();
        writer.close().unwrap();
        buf
    }

    fn temp_output_path(name: &str) -> PathBuf {
        let suffix = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos();
        std::env::temp_dir().join(format!("pm-compact-{name}-{suffix}.parquet"))
    }

    #[tokio::test]
    async fn compact_store_parquet_prefix_to_local_merges_multiple_files() {
        let store: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
        let payload_a = make_test_parquet(&[("a", 1), ("b", 2)]);
        let payload_b = make_test_parquet(&[("c", 3)]);

        store
            .put_opts(
                &ObjectPath::from("prefix/part-0001.parquet"),
                PutPayload::from(Bytes::from(payload_a)),
                PutOptions::default(),
            )
            .await
            .unwrap();
        store
            .put_opts(
                &ObjectPath::from("prefix/part-0002.parquet"),
                PutPayload::from(Bytes::from(payload_b)),
                PutOptions::default(),
            )
            .await
            .unwrap();

        let output_path = temp_output_path("merge");
        let result =
            compact_store_parquet_prefix_to_local(store, "prefix", output_path.as_path(), 2)
                .await
                .unwrap();

        assert_eq!(result.source_objects, 2);
        assert_eq!(result.rows_written, 3);
        assert_eq!(result.output_path, output_path);

        let output_bytes = std::fs::read(output_path.as_path()).unwrap();
        let record_reader = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(output_bytes))
            .unwrap()
            .build()
            .unwrap();
        let mut names = Vec::new();
        let mut values = Vec::new();
        for batch in record_reader {
            let batch = batch.unwrap();
            let name_col = batch
                .column(0)
                .as_any()
                .downcast_ref::<StringArray>()
                .unwrap();
            let value_col = batch
                .column(1)
                .as_any()
                .downcast_ref::<Int64Array>()
                .unwrap();
            for i in 0..batch.num_rows() {
                names.push(name_col.value(i).to_string());
                values.push(value_col.value(i));
            }
        }
        assert_eq!(names, vec!["a", "b", "c"]);
        assert_eq!(values, vec![1, 2, 3]);

        std::fs::remove_file(output_path).unwrap();
    }

    #[tokio::test]
    async fn compact_store_parquet_prefix_to_local_ignores_non_parquet_objects() {
        let store: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
        let payload = make_test_parquet(&[("a", 1)]);

        store
            .put_opts(
                &ObjectPath::from("prefix/part-0001.parquet"),
                PutPayload::from(Bytes::from(payload)),
                PutOptions::default(),
            )
            .await
            .unwrap();
        store
            .put_opts(
                &ObjectPath::from("prefix/notes.txt"),
                PutPayload::from(Bytes::from_static(b"ignored")),
                PutOptions::default(),
            )
            .await
            .unwrap();

        let output_path = temp_output_path("filter");
        let result =
            compact_store_parquet_prefix_to_local(store, "prefix", output_path.as_path(), 2)
                .await
                .unwrap();
        assert_eq!(result.source_objects, 1);
        assert_eq!(result.rows_written, 1);

        std::fs::remove_file(output_path).unwrap();
    }

    #[test]
    fn anonymize_parquet_file_offsets_timestamps_and_hashes_columns() {
        use arrow::array::{Int64Array, StringArray, UInt64Array};
        use arrow::datatypes::{DataType, Field, Schema};
        use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;

        // Create test data with timestamps, session IDs, and total_hashes
        let schema = Arc::new(Schema::new(vec![
            Field::new("timestamp_start_unix_ms", DataType::UInt64, false),
            Field::new("provided_session_id", DataType::Utf8, true),
            Field::new("org_id", DataType::Utf8, true),
            Field::new("value", DataType::Int64, false),
        ]));

        // Create test data: timestamps [1000, 2000, 500, 3000], session IDs, org IDs
        let timestamps = vec![1000u64, 2000, 500, 3000];
        let sessions: Vec<Option<String>> = vec![
            Some("session-1".to_string()),
            Some("session-2".to_string()),
            Some("session-1".to_string()),
            None,
        ];
        let orgs: Vec<Option<String>> = vec![
            Some("org-1".to_string()),
            Some("org-2".to_string()),
            Some("org-1".to_string()),
            Some("org-3".to_string()),
        ];
        let values = vec![1i64, 2, 3, 4];

        let batch = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(UInt64Array::from(timestamps)),
                Arc::new(StringArray::from(sessions)),
                Arc::new(StringArray::from(orgs)),
                Arc::new(Int64Array::from(values)),
            ],
        )
        .unwrap();

        // Write test input file
        let input_path = temp_output_path("anonymize_input");
        let output_path = temp_output_path("anonymize_output");

        let mut buf = Vec::new();
        let props = WriterProperties::builder()
            .set_compression(Compression::SNAPPY)
            .build();
        let mut writer = ArrowWriter::try_new(&mut buf, schema, Some(props)).unwrap();
        writer.write(&batch).unwrap();
        writer.close().unwrap();

        std::fs::write(&input_path, &buf).unwrap();

        // Run anonymization with min_timestamp=500 (the actual minimum in the data)
        anonymize_parquet_file(
            &input_path,
            &output_path,
            &["provided_session_id", "org_id"],
            &["org_id"], // org_id is in both, so it will be hashed not dropped
            500,         // min timestamp to offset by
        )
        .unwrap();

        // Verify output
        let output_bytes = std::fs::read(&output_path).unwrap();
        let record_reader = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(output_bytes))
            .unwrap()
            .build()
            .unwrap();

        let mut timestamps = Vec::new();
        let mut sessions = Vec::new();
        let mut orgs = Vec::new();
        let mut values = Vec::new();

        for batch in record_reader {
            let batch = batch.unwrap();
            let schema = batch.schema();

            // Get column indices by name
            let ts_idx = schema.index_of("timestamp_start_unix_ms").unwrap();
            let session_idx = schema.index_of("provided_session_id").unwrap();
            let org_idx = schema.index_of("org_id").unwrap();
            let value_idx = schema.index_of("value").unwrap();

            let ts_col = batch
                .column(ts_idx)
                .as_any()
                .downcast_ref::<UInt64Array>()
                .unwrap();
            let session_col = batch
                .column(session_idx)
                .as_any()
                .downcast_ref::<StringArray>()
                .unwrap();
            let org_col = batch
                .column(org_idx)
                .as_any()
                .downcast_ref::<StringArray>()
                .unwrap();
            let value_col = batch
                .column(value_idx)
                .as_any()
                .downcast_ref::<Int64Array>()
                .unwrap();

            for i in 0..batch.num_rows() {
                timestamps.push(ts_col.value(i));
                sessions.push(session_col.value(i).to_string());
                orgs.push(org_col.value(i).to_string());
                values.push(value_col.value(i));
            }
        }

        // Verify timestamps are offset (minimum should be 0)
        assert!(timestamps.contains(&0), "Minimum timestamp should be 0");
        assert!(timestamps.contains(&500), "Should have timestamp 500");
        assert!(timestamps.contains(&1500), "Should have timestamp 1500");
        assert!(timestamps.contains(&2500), "Should have timestamp 2500");

        // Verify session IDs are replaced with unique IDs
        // Note: sessions[3] is null, so it will be an empty string
        assert_eq!(sessions.len(), 4);
        for (i, session) in sessions.iter().enumerate() {
            if i == 3 {
                // Null value - should be empty string
                assert_eq!(session.len(), 0, "Null session should be empty string");
            } else {
                // Non-null values - should be unique IDs (numeric strings)
                assert_ne!(session, "session-1", "Session should be anonymized");
                assert_ne!(session, "session-2", "Session should be anonymized");
                assert!(
                    session.parse::<u64>().is_ok(),
                    "Session should be a numeric ID"
                );
            }
        }
        // Verify same session values get same IDs
        assert_eq!(
            sessions[0], sessions[2],
            "Same session values should get same ID"
        );

        // Verify org IDs are replaced with unique IDs (not dropped, since they're in both lists)
        assert_eq!(orgs.len(), 4);
        for org in &orgs {
            assert_ne!(org, "org-1", "Org should be anonymized");
            assert_ne!(org, "org-2", "Org should be anonymized");
            assert_ne!(org, "org-3", "Org should be anonymized");
            assert!(org.parse::<u64>().is_ok(), "Org should be a numeric ID");
        }
        // Verify same org values get same IDs
        assert_eq!(orgs[0], orgs[2], "Same org values should get same ID");

        // Verify values are unchanged
        assert_eq!(values, vec![1, 2, 3, 4]);

        // Cleanup
        std::fs::remove_file(input_path).unwrap();
        std::fs::remove_file(output_path).unwrap();
    }

    #[test]
    fn anonymize_parquet_file_zero_min_timestamp_does_not_shift_timestamps() {
        use arrow::array::{Int64Array, UInt64Array};
        use arrow::datatypes::{DataType, Field, Schema};
        use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;

        let schema = Arc::new(Schema::new(vec![
            Field::new("timestamp_start_unix_ms", DataType::UInt64, false),
            Field::new("value", DataType::Int64, false),
        ]));

        let input_timestamps = vec![1000u64, 2000u64, 3500u64];
        let values = vec![1i64, 2, 3];

        let batch = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(UInt64Array::from(input_timestamps.clone())),
                Arc::new(Int64Array::from(values)),
            ],
        )
        .unwrap();

        let input_path = temp_output_path("anonymize_zero_shift_input");
        let output_path = temp_output_path("anonymize_zero_shift_output");

        let mut buf = Vec::new();
        let props = WriterProperties::builder()
            .set_compression(Compression::SNAPPY)
            .build();
        let mut writer = ArrowWriter::try_new(&mut buf, schema, Some(props)).unwrap();
        writer.write(&batch).unwrap();
        writer.close().unwrap();
        std::fs::write(&input_path, &buf).unwrap();

        let returned_min_ts =
            anonymize_parquet_file(&input_path, &output_path, &[], &[], 0).unwrap();
        assert_eq!(returned_min_ts, 0);

        let output_bytes = std::fs::read(&output_path).unwrap();
        let record_reader = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(output_bytes))
            .unwrap()
            .build()
            .unwrap();

        let mut output_timestamps = Vec::new();
        for batch in record_reader {
            let batch = batch.unwrap();
            let ts_col = batch
                .column(batch.schema().index_of("timestamp_start_unix_ms").unwrap())
                .as_any()
                .downcast_ref::<UInt64Array>()
                .unwrap();
            for i in 0..batch.num_rows() {
                output_timestamps.push(ts_col.value(i));
            }
        }

        assert_eq!(output_timestamps, input_timestamps);

        std::fs::remove_file(input_path).unwrap();
        std::fs::remove_file(output_path).unwrap();
    }
}
