use arrow::array::*;
use arrow::datatypes::*;
use arrow::record_batch::RecordBatch;
use arrow::ipc::writer::{StreamWriter, IpcWriteOptions};
use std::io::{Write, Seek, SeekFrom};
use std::collections::HashMap;
use std::sync::Arc;
use serde::{Serialize, Deserialize};
use crate::dna::hel::error::HlxError;

type ArrayRef = Arc<dyn Array>;
pub const HLX_MAGIC: &[u8; 4] = b"HLX\x01";
pub const HLX_FOOTER_MAGIC: &[u8; 4] = b"\xFF\xFF\xFF\xFF";
pub const HLX_VERSION: u8 = 1;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HlxHeader {
    pub schema: serde_json::Value,
    pub metadata: HashMap<String, serde_json::Value>,
    pub flags: u8,
    pub row_count: u64,
    pub preview_rows: Option<Vec<serde_json::Value>>,
}
impl HlxHeader {
    pub fn new(schema: &Schema, metadata: HashMap<String, serde_json::Value>) -> Self {
        Self {
            schema: serde_json::Value::String(format!("{:?}", schema)),
            metadata,
            flags: 0,
            row_count: 0,
            preview_rows: None,
        }
    }
    pub fn with_compression(mut self, compressed: bool) -> Self {
        if compressed {
            self.flags |= 0x01;
        } else {
            self.flags &= !0x01;
        }
        self
    }
    pub fn with_row_count(mut self, count: u64) -> Self {
        self.row_count = count;
        self
    }
    pub fn with_preview(mut self, preview: Vec<serde_json::Value>) -> Self {
        self.preview_rows = Some(preview);
        self
    }
    pub fn is_compressed(&self) -> bool {
        (self.flags & 0x01) != 0
    }
    pub fn to_json_bytes(&self) -> Result<Vec<u8>, HlxError> {
        serde_json::to_vec(self).map_err(|e| HlxError::json_error(e.to_string(), ""))
    }
    pub fn from_json_bytes(bytes: &[u8]) -> Result<Self, HlxError> {
        serde_json::from_slice(bytes)
            .map_err(|e| HlxError::json_error(e.to_string(), ""))
    }
}
pub struct HlxWriter<W: Write + Seek> {
    writer: W,
    header: Option<HlxHeader>,
    schema: Option<Arc<Schema>>,
    batches: Vec<RecordBatch>,
    compression_enabled: bool,
}
impl<W: Write + Seek> HlxWriter<W> {
    pub fn new(writer: W) -> Self {
        Self {
            writer,
            header: None,
            schema: None,
            batches: Vec::new(),
            compression_enabled: false,
        }
    }
    pub fn with_compression(mut self, enabled: bool) -> Self {
        self.compression_enabled = enabled;
        self
    }
    pub fn with_schema(mut self, schema: Arc<Schema>) -> Self {
        self.schema = Some(schema);
        self
    }
    pub fn add_batch(&mut self, batch: RecordBatch) -> Result<(), HlxError> {
        if self.schema.is_none() {
            self.schema = Some(batch.schema().clone());
        } else if **self.schema.as_ref().unwrap() != *batch.schema().as_ref() {
            return Err(
                HlxError::validation_error(
                    "Schema mismatch between batches",
                    "All batches must have the same schema",
                ),
            );
        }
        self.batches.push(batch);
        Ok(())
    }
    pub fn finalize(&mut self) -> Result<(), HlxError> {
        if self.batches.is_empty() {
            return Err(
                HlxError::validation_error(
                    "No data to write",
                    "At least one record batch is required",
                ),
            );
        }
        let schema = self.schema.as_ref().unwrap().clone();
        let total_rows = self.batches.iter().map(|b| b.num_rows()).sum::<usize>() as u64;
        let mut metadata = HashMap::new();
        metadata
            .insert(
                "format_version".to_string(),
                serde_json::Value::String("1.0".to_string()),
            );
        metadata
            .insert(
                "created_at".to_string(),
                serde_json::Value::String(chrono::Utc::now().to_rfc3339()),
            );
        let mut header = HlxHeader::new(schema.as_ref(), metadata)
            .with_compression(self.compression_enabled)
            .with_row_count(total_rows);
        if let Some(first_batch) = self.batches.first() {
            let preview_count = std::cmp::min(10, first_batch.num_rows());
            let preview_rows = extract_preview_rows(first_batch, preview_count);
            header = header.with_preview(preview_rows);
        }
        self.writer.write_all(HLX_MAGIC)?;
        let header_bytes = header.to_json_bytes()?;
        let header_len = header_bytes.len() as u32;
        self.writer.write_all(&header_len.to_le_bytes())?;
        self.writer.write_all(&header_bytes)?;
        let options = IpcWriteOptions::default();
        let mut stream_writer = StreamWriter::try_new_with_options(&mut self.writer, schema.as_ref(), options)
            .map_err(|e| HlxError::io_error(format!("Failed to create Arrow IPC stream writer: {}", e), "Check Arrow IPC format support"))?;
        for batch in &self.batches {
            stream_writer.write(batch)
                .map_err(|e| HlxError::io_error(format!("Failed to write Arrow record batch: {}", e), "Check Arrow record batch format"))?;
        }
        stream_writer.finish()
            .map_err(|e| HlxError::io_error(format!("Failed to finalize Arrow IPC stream: {}", e), "Check stream finalization"))?;
        if let Some(preview_rows) = &header.preview_rows {
            if !preview_rows.is_empty() {
                self.writer.write_all(HLX_FOOTER_MAGIC)?;
                let preview_jsonl = preview_rows
                    .iter()
                    .map(|row| serde_json::to_string(row).unwrap_or_default())
                    .collect::<Vec<_>>()
                    .join("\n");
                let footer_len = preview_jsonl.len() as u32;
                self.writer.write_all(&footer_len.to_le_bytes())?;
                self.writer.write_all(preview_jsonl.as_bytes())?;
            }
        }
        Ok(())
    }
}
fn extract_preview_rows(batch: &RecordBatch, count: usize) -> Vec<serde_json::Value> {
    let mut preview_rows = Vec::new();
    for row_idx in 0..count {
        let mut row_json = serde_json::Map::new();
        for (field_idx, field) in batch.schema().fields().iter().enumerate() {
            if let Some(array) = batch
                .column(field_idx)
                .as_any()
                .downcast_ref::<StringArray>()
            {
                if array.is_valid(row_idx) {
                    let value = array.value(row_idx);
                    row_json
                        .insert(
                            field.name().clone(),
                            serde_json::Value::String(value.to_string()),
                        );
                }
            } else if let Some(array) = batch
                .column(field_idx)
                .as_any()
                .downcast_ref::<Float64Array>()
            {
                if array.is_valid(row_idx) {
                    let value = array.value(row_idx);
                    row_json
                        .insert(
                            field.name().clone(),
                            serde_json::Value::Number(
                                serde_json::Number::from_f64(value)
                                    .unwrap_or(serde_json::Number::from(0)),
                            ),
                        );
                }
            } else if let Some(array) = batch
                .column(field_idx)
                .as_any()
                .downcast_ref::<Int64Array>()
            {
                if array.is_valid(row_idx) {
                    let value = array.value(row_idx);
                    row_json
                        .insert(
                            field.name().clone(),
                            serde_json::Value::Number(serde_json::Number::from(value)),
                        );
                }
            } else {
                if batch.column(field_idx).is_valid(row_idx) {
                    let value_str = format!("{:?}", batch.column(field_idx));
                    row_json
                        .insert(field.name().clone(), serde_json::Value::String(value_str));
                } else {
                    row_json.insert(field.name().clone(), serde_json::Value::Null);
                }
            }
        }
        preview_rows.push(serde_json::Value::Object(row_json));
    }
    preview_rows
}
pub struct HlxReader<R: std::io::Read + Seek> {
    reader: R,
    header: Option<HlxHeader>,
}
impl<R: std::io::Read + Seek> HlxReader<R> {
    pub fn new(reader: R) -> Self {
        Self { reader, header: None }
    }
    pub fn read_header(&mut self) -> Result<&HlxHeader, HlxError> {
        if self.header.is_some() {
            return Ok(self.header.as_ref().unwrap());
        }
        let mut magic = [0u8; 4];
        self.reader.read_exact(&mut magic)?;
        if magic != *HLX_MAGIC {
            return Err(
                HlxError::validation_error(
                    "Invalid HLX magic number",
                    "File does not appear to be a valid HLX file",
                ),
            );
        }
        let mut header_len_bytes = [0u8; 4];
        self.reader.read_exact(&mut header_len_bytes)?;
        let header_len = u32::from_le_bytes(header_len_bytes) as usize;
        let mut header_bytes = vec![0u8; header_len];
        self.reader.read_exact(&mut header_bytes)?;
        let header: HlxHeader = HlxHeader::from_json_bytes(&header_bytes)?;
        self.header = Some(header);
        Ok(self.header.as_ref().unwrap())
    }
    pub fn get_preview(&mut self) -> Result<Option<Vec<serde_json::Value>>, HlxError> {
        let header = self.read_header()?;
        if let Some(preview_rows) = &header.preview_rows {
            return Ok(Some(preview_rows.clone()));
        }
        if let Some(footer_rows) = self.read_footer()? {
            return Ok(Some(footer_rows));
        }
        Ok(None)
    }
    fn read_footer(&mut self) -> Result<Option<Vec<serde_json::Value>>, HlxError> {
        let file_size = self.reader.seek(SeekFrom::End(0))?;
        if file_size < 8 {
            return Ok(None);
        }
        self.reader.seek(SeekFrom::End(-8))?;
        let mut footer_header = [0u8; 8];
        self.reader.read_exact(&mut footer_header)?;
        let magic = &footer_header[0..4];
        if magic != *HLX_FOOTER_MAGIC {
            return Ok(None);
        }
        let footer_len = u32::from_le_bytes(footer_header[4..8].try_into().unwrap())
            as usize;
        self.reader.seek(SeekFrom::End(-8 - footer_len as i64))?;
        let mut footer_bytes = vec![0u8; footer_len];
        self.reader.read_exact(&mut footer_bytes)?;
        let footer_jsonl = String::from_utf8(footer_bytes)
            .map_err(|_| HlxError::validation_error("Invalid UTF-8 in footer", ""))?;
        let rows: Vec<serde_json::Value> = footer_jsonl
            .lines()
            .filter(|line| !line.trim().is_empty())
            .map(|line| serde_json::from_str(line).unwrap_or(serde_json::Value::Null))
            .collect();
        Ok(Some(rows))
    }
    pub fn get_schema(&mut self) -> Result<&Schema, HlxError> {
        let header = self.read_header()?;
        Err(
            HlxError::validation_error(
                "Schema deserialization not implemented",
                "Use Arrow's native IPC reader for full functionality",
            ),
        )
    }
}
pub fn value_to_arrow_array(
    field: &Field,
    values: Vec<crate::value::Value>,
) -> Result<ArrayRef, HlxError> {
    match field.data_type() {
        DataType::Utf8 => {
            let string_values: Vec<Option<String>> = values
                .into_iter()
                .map(|v| match v {
                    crate::value::Value::String(s) => Some(s),
                    _ => Some(v.to_string()),
                })
                .collect();
            let array = StringArray::from(string_values);
            Ok(Arc::new(array))
        }
        DataType::Float64 => {
            let float_values: Vec<Option<f64>> = values
                .into_iter()
                .map(|v| match v {
                    crate::value::Value::Number(n) => Some(n),
                    _ => v.to_string().parse().ok(),
                })
                .collect();
            let array = Float64Array::from(float_values);
            Ok(Arc::new(array))
        }
        DataType::Int64 => {
            let int_values: Vec<Option<i64>> = values
                .into_iter()
                .map(|v| match v {
                    crate::value::Value::Number(n) => Some(n as i64),
                    _ => v.to_string().parse().ok(),
                })
                .collect();
            let array = Int64Array::from(int_values);
            Ok(Arc::new(array))
        }
        _ => {
            let string_values: Vec<Option<String>> = values
                .into_iter()
                .map(|v| Some(v.to_string()))
                .collect();
            let array = StringArray::from(string_values);
            Ok(Arc::new(array))
        }
    }
}
pub fn create_arrow_schema(fields: Vec<(&str, DataType)>) -> Schema {
    let arrow_fields: Vec<Field> = fields
        .into_iter()
        .map(|(name, data_type)| Field::new(name, data_type, true))
        .collect();
    Schema::new(arrow_fields)
}
#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Cursor;
    #[test]
    fn test_hlx_header_serialization() {
        let mut metadata = HashMap::new();
        metadata
            .insert("test".to_string(), serde_json::Value::String("value".to_string()));
        let schema = create_arrow_schema(
            vec![("name", DataType::Utf8), ("age", DataType::Int64),],
        );
        let header = HlxHeader::new(schema, metadata)
            .with_compression(true)
            .with_row_count(100);
        let json_bytes = header.to_json_bytes().unwrap();
        let deserialized: HlxHeader = HlxHeader::from_json_bytes(&json_bytes).unwrap();
        assert_eq!(header.row_count, deserialized.row_count);
        assert!(header.is_compressed());
    }
    #[test]
    fn test_hlx_writer_basic() {
        let buffer = Vec::new();
        let cursor = Cursor::new(buffer);
        let mut writer = HlxWriter::new(cursor);
        assert!(writer.header.is_none());
    }
}