//! Fast data I/O using Arrow/Parquet.
//!
//! Handles parquet reading and transfer timestamp extraction in Rust.

use hashbrown::{HashMap as HBMap, HashSet as HBSet};
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
use arrow::array::{Array, StringArray, Int64Array, Float64Array};
use rayon::prelude::*;
use std::fs::File;
use std::path::Path;
use std::sync::Mutex;

/// Transfer record for timestamp extraction
#[derive(Clone)]
pub struct TransferRecord {
    pub from_address: String,
    pub to_address: String,
    pub block_timestamp: i64,
    pub amount_usd: f64,
}

/// Timestamp record for burst detection
#[derive(Clone)]
pub struct TimestampRecord {
    pub timestamp: i64,
    pub volume: f64,
    pub counterparty: String,
}

/// Load transfers from parquet file
pub fn load_transfers_from_parquet(
    path: &str,
    start_timestamp_ms: i64,
    end_timestamp_ms: i64,
) -> Result<Vec<TransferRecord>, String> {
    let file = File::open(Path::new(path))
        .map_err(|e| format!("Failed to open file {}: {}", path, e))?;

    let builder = ParquetRecordBatchReaderBuilder::try_new(file)
        .map_err(|e| format!("Failed to create parquet reader: {}", e))?;

    let reader = builder
        .with_batch_size(65536)
        .build()
        .map_err(|e| format!("Failed to build reader: {}", e))?;

    let mut transfers = Vec::new();

    for batch_result in reader {
        let batch = batch_result.map_err(|e| format!("Failed to read batch: {}", e))?;

        let schema = batch.schema();

        // Find column indices
        let from_idx = schema.index_of("from_address")
            .map_err(|_| "Missing from_address column")?;
        let to_idx = schema.index_of("to_address")
            .map_err(|_| "Missing to_address column")?;
        let ts_idx = schema.index_of("block_timestamp")
            .map_err(|_| "Missing block_timestamp column")?;
        let amount_idx = schema.index_of("amount_usd")
            .ok(); // Optional column

        let from_col = batch.column(from_idx)
            .as_any()
            .downcast_ref::<StringArray>()
            .ok_or("from_address is not a string column")?;

        let to_col = batch.column(to_idx)
            .as_any()
            .downcast_ref::<StringArray>()
            .ok_or("to_address is not a string column")?;

        let ts_col = batch.column(ts_idx)
            .as_any()
            .downcast_ref::<Int64Array>()
            .ok_or("block_timestamp is not an int64 column")?;

        let amount_col = amount_idx.and_then(|idx| {
            batch.column(idx)
                .as_any()
                .downcast_ref::<Float64Array>()
        });

        for i in 0..batch.num_rows() {
            let ts = ts_col.value(i);

            // Filter by timestamp range
            if ts < start_timestamp_ms || ts >= end_timestamp_ms {
                continue;
            }

            let from_addr = from_col.value(i).to_string();
            let to_addr = to_col.value(i).to_string();
            let amount = amount_col.map(|col| col.value(i)).unwrap_or(0.0);

            transfers.push(TransferRecord {
                from_address: from_addr,
                to_address: to_addr,
                block_timestamp: ts,
                amount_usd: amount,
            });
        }
    }

    Ok(transfers)
}

/// Build timestamp data for burst detection from transfers.
/// Returns HashMap: address -> Vec<(timestamp, volume, counterparty)>
pub fn build_timestamp_data(
    transfers: &[TransferRecord],
    filter_addresses: &HBSet<String>,
) -> HBMap<String, Vec<TimestampRecord>> {
    // Use parallel processing with thread-safe collection
    let result: Mutex<HBMap<String, Vec<TimestampRecord>>> = Mutex::new(HBMap::new());

    // Process in chunks for better parallelism
    transfers.par_chunks(10000).for_each(|chunk| {
        let mut local_map: HBMap<String, Vec<TimestampRecord>> = HBMap::new();

        for transfer in chunk {
            // Add record for from_address if in filter set
            if filter_addresses.contains(&transfer.from_address) {
                local_map.entry(transfer.from_address.clone())
                    .or_insert_with(Vec::new)
                    .push(TimestampRecord {
                        timestamp: transfer.block_timestamp,
                        volume: transfer.amount_usd,
                        counterparty: transfer.to_address.clone(),
                    });
            }

            // Add record for to_address if in filter set
            if filter_addresses.contains(&transfer.to_address) {
                local_map.entry(transfer.to_address.clone())
                    .or_insert_with(Vec::new)
                    .push(TimestampRecord {
                        timestamp: transfer.block_timestamp,
                        volume: transfer.amount_usd,
                        counterparty: transfer.from_address.clone(),
                    });
            }
        }

        // Merge into global result
        let mut global = result.lock().unwrap();
        for (addr, records) in local_map {
            global.entry(addr)
                .or_insert_with(Vec::new)
                .extend(records);
        }
    });

    result.into_inner().unwrap()
}

/// Load money flows from parquet and build edge list for graph.
/// Returns (sources, targets, weights, edge_attrs)
pub fn load_money_flows_from_parquet(
    path: &str,
    start_timestamp_ms: i64,
    end_timestamp_ms: i64,
) -> Result<(Vec<String>, Vec<String>, Vec<f64>), String> {
    let file = File::open(Path::new(path))
        .map_err(|e| format!("Failed to open file {}: {}", path, e))?;

    let builder = ParquetRecordBatchReaderBuilder::try_new(file)
        .map_err(|e| format!("Failed to create parquet reader: {}", e))?;

    let reader = builder
        .with_batch_size(65536)
        .build()
        .map_err(|e| format!("Failed to build reader: {}", e))?;

    let mut sources = Vec::new();
    let mut targets = Vec::new();
    let mut weights = Vec::new();

    for batch_result in reader {
        let batch = batch_result.map_err(|e| format!("Failed to read batch: {}", e))?;

        let schema = batch.schema();

        let from_idx = schema.index_of("from_address")
            .map_err(|_| "Missing from_address column")?;
        let to_idx = schema.index_of("to_address")
            .map_err(|_| "Missing to_address column")?;
        let amount_idx = schema.index_of("amount_usd_sum")
            .map_err(|_| "Missing amount_usd_sum column")?;
        let ts_idx = schema.index_of("last_seen_timestamp")
            .ok();

        let from_col = batch.column(from_idx)
            .as_any()
            .downcast_ref::<StringArray>()
            .ok_or("from_address is not a string column")?;

        let to_col = batch.column(to_idx)
            .as_any()
            .downcast_ref::<StringArray>()
            .ok_or("to_address is not a string column")?;

        let amount_col = batch.column(amount_idx)
            .as_any()
            .downcast_ref::<Float64Array>()
            .ok_or("amount_usd_sum is not a float64 column")?;

        let ts_col = ts_idx.and_then(|idx| {
            batch.column(idx)
                .as_any()
                .downcast_ref::<Int64Array>()
        });

        for i in 0..batch.num_rows() {
            // Filter by last_seen_timestamp if available
            if let Some(ts) = ts_col {
                let timestamp = ts.value(i);
                if timestamp < start_timestamp_ms || timestamp >= end_timestamp_ms {
                    continue;
                }
            }

            sources.push(from_col.value(i).to_string());
            targets.push(to_col.value(i).to_string());
            weights.push(amount_col.value(i));
        }
    }

    Ok((sources, targets, weights))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_build_timestamp_data() {
        let transfers = vec![
            TransferRecord {
                from_address: "A".to_string(),
                to_address: "B".to_string(),
                block_timestamp: 1000,
                amount_usd: 100.0,
            },
            TransferRecord {
                from_address: "B".to_string(),
                to_address: "C".to_string(),
                block_timestamp: 2000,
                amount_usd: 200.0,
            },
        ];

        let mut filter: HBSet<String> = HBSet::new();
        filter.insert("A".to_string());
        filter.insert("B".to_string());

        let result = build_timestamp_data(&transfers, &filter);

        assert!(result.contains_key("A"));
        assert!(result.contains_key("B"));
        assert_eq!(result["A"].len(), 1);
        assert_eq!(result["B"].len(), 2); // B appears in both transfers
    }
}
