tellaro-query-language 1.3.8

//! File operations for TQL.
//!
//! This module provides file-based query operations including:
//! - Reading and querying JSON, JSONL, and CSV files
//! - Folder batch processing with pattern matching
//! - Format auto-detection
//! - Streaming support for large files

use crate::error::{Result, TqlError};
use crate::evaluator::TqlEvaluator;
use crate::parser::{AstNode, TqlParser};
use crate::stats_evaluator::{AggregationSpec, StatsEvaluator, StatsQuery};
use glob::glob;
use rayon::prelude::*;
use serde_json::Value as JsonValue;
use std::fs::File;
use std::io::{BufRead, BufReader};
use std::path::Path;

/// Supported file formats
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum FileFormat {
    /// Single JSON object or array
    Json,
    /// JSON Lines (newline-delimited JSON)
    JsonL,
    /// Comma-separated values
    Csv,
    /// Auto-detect format from file extension
    Auto,
}

impl std::str::FromStr for FileFormat {
    type Err = std::convert::Infallible;

    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
        Ok(match s.to_lowercase().as_str() {
            "json" => FileFormat::Json,
            "jsonl" | "ndjson" => FileFormat::JsonL,
            "csv" => FileFormat::Csv,
            _ => FileFormat::Auto,
        })
    }
}

impl FileFormat {
    /// Detect format from file extension
    pub fn from_path(path: &Path) -> Self {
        match path.extension().and_then(|e| e.to_str()) {
            Some("json") => FileFormat::Json,
            Some("jsonl") | Some("ndjson") => FileFormat::JsonL,
            Some("csv") => FileFormat::Csv,
            _ => FileFormat::Json, // Default to JSON
        }
    }
}

/// CSV parsing configuration
#[derive(Debug, Clone)]
pub struct CsvConfig {
    /// Field delimiter character
    pub delimiter: u8,
    /// Whether the first row contains headers
    pub has_headers: bool,
    /// Custom header names (overrides first row if has_headers is true)
    pub custom_headers: Option<Vec<String>>,
}

impl Default for CsvConfig {
    fn default() -> Self {
        Self {
            delimiter: b',',
            has_headers: true,
            custom_headers: None,
        }
    }
}

impl CsvConfig {
    /// Create config with custom delimiter
    pub fn with_delimiter(mut self, delimiter: char) -> Self {
        self.delimiter = delimiter as u8;
        self
    }

    /// Create config without headers
    pub fn without_headers(mut self) -> Self {
        self.has_headers = false;
        self
    }

    /// Create config with custom headers
    pub fn with_headers(mut self, headers: Vec<String>) -> Self {
        self.custom_headers = Some(headers);
        self
    }
}

/// File query operations
pub struct FileOps {
    parser: TqlParser,
    evaluator: TqlEvaluator,
    stats_evaluator: StatsEvaluator,
}

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

impl FileOps {
    /// Create new FileOps instance
    pub fn new() -> Self {
        Self {
            parser: TqlParser::new(),
            evaluator: TqlEvaluator::new(),
            stats_evaluator: StatsEvaluator::new(),
        }
    }

    /// Read records from a file
    pub fn read_file(
        &self,
        path: &Path,
        format: FileFormat,
        csv_config: &CsvConfig,
    ) -> Result<Vec<JsonValue>> {
        let format = if format == FileFormat::Auto {
            FileFormat::from_path(path)
        } else {
            format
        };

        match format {
            FileFormat::Json => self.read_json(path),
            FileFormat::JsonL => self.read_jsonl(path),
            FileFormat::Csv => self.read_csv(path, csv_config),
            FileFormat::Auto => self.read_json(path),
        }
    }

    /// Read JSON file (single object or array)
    fn read_json(&self, path: &Path) -> Result<Vec<JsonValue>> {
        let file = File::open(path).map_err(|e| {
            TqlError::ExecutionError(format!("Failed to open file {}: {}", path.display(), e))
        })?;

        let reader = BufReader::new(file);
        let value: JsonValue = serde_json::from_reader(reader).map_err(|e| {
            TqlError::ExecutionError(format!(
                "Failed to parse JSON from {}: {}",
                path.display(),
                e
            ))
        })?;

        match value {
            JsonValue::Array(arr) => Ok(arr),
            obj @ JsonValue::Object(_) => Ok(vec![obj]),
            _ => Err(TqlError::ExecutionError(format!(
                "JSON file {} must contain an object or array",
                path.display()
            ))),
        }
    }

    /// Read JSONL file (newline-delimited JSON)
    fn read_jsonl(&self, path: &Path) -> Result<Vec<JsonValue>> {
        let file = File::open(path).map_err(|e| {
            TqlError::ExecutionError(format!("Failed to open file {}: {}", path.display(), e))
        })?;

        let reader = BufReader::new(file);
        let mut records = Vec::new();

        for (line_num, line) in reader.lines().enumerate() {
            let line = line.map_err(|e| {
                TqlError::ExecutionError(format!(
                    "Failed to read line {} from {}: {}",
                    line_num + 1,
                    path.display(),
                    e
                ))
            })?;

            let trimmed = line.trim();
            if trimmed.is_empty() {
                continue;
            }

            let value: JsonValue = serde_json::from_str(trimmed).map_err(|e| {
                TqlError::ExecutionError(format!(
                    "Failed to parse JSON at line {} in {}: {}",
                    line_num + 1,
                    path.display(),
                    e
                ))
            })?;

            records.push(value);
        }

        Ok(records)
    }

    /// Read CSV file
    fn read_csv(&self, path: &Path, config: &CsvConfig) -> Result<Vec<JsonValue>> {
        let file = File::open(path).map_err(|e| {
            TqlError::ExecutionError(format!("Failed to open file {}: {}", path.display(), e))
        })?;

        let mut rdr = csv::ReaderBuilder::new()
            .delimiter(config.delimiter)
            .has_headers(config.has_headers)
            .from_reader(file);

        let headers: Vec<String> = if let Some(ref custom) = config.custom_headers {
            custom.clone()
        } else if config.has_headers {
            rdr.headers()
                .map_err(|e| {
                    TqlError::ExecutionError(format!(
                        "Failed to read CSV headers from {}: {}",
                        path.display(),
                        e
                    ))
                })?
                .iter()
                .map(|s| s.to_string())
                .collect()
        } else {
            Vec::new()
        };

        let mut records = Vec::new();

        for (row_num, result) in rdr.records().enumerate() {
            let record = result.map_err(|e| {
                TqlError::ExecutionError(format!(
                    "Failed to read CSV row {} from {}: {}",
                    row_num + 1,
                    path.display(),
                    e
                ))
            })?;

            let mut obj = serde_json::Map::new();

            for (i, field) in record.iter().enumerate() {
                let key = if i < headers.len() {
                    headers[i].clone()
                } else {
                    format!("col{}", i)
                };

                // Try to parse as number or boolean, fall back to string
                let value = if let Ok(n) = field.parse::<i64>() {
                    JsonValue::Number(n.into())
                } else if let Ok(n) = field.parse::<f64>() {
                    JsonValue::Number(serde_json::Number::from_f64(n).unwrap_or_else(|| 0.into()))
                } else if field.eq_ignore_ascii_case("true") {
                    JsonValue::Bool(true)
                } else if field.eq_ignore_ascii_case("false") {
                    JsonValue::Bool(false)
                } else if field.is_empty() || field.eq_ignore_ascii_case("null") {
                    JsonValue::Null
                } else {
                    JsonValue::String(field.to_string())
                };

                obj.insert(key, value);
            }

            records.push(JsonValue::Object(obj));
        }

        Ok(records)
    }

    /// Query a file with a TQL query
    pub fn query_file(
        &self,
        path: &Path,
        query: &str,
        format: FileFormat,
        csv_config: &CsvConfig,
    ) -> Result<Vec<JsonValue>> {
        let records = self.read_file(path, format, csv_config)?;
        let ast = self.parser.parse(query)?;

        let results = self.evaluator.filter(&ast, &records)?;
        Ok(results.into_iter().cloned().collect())
    }

    /// Query a file with enrichment (applies mutators to results)
    pub fn query_file_enriched(
        &self,
        path: &Path,
        query: &str,
        format: FileFormat,
        csv_config: &CsvConfig,
    ) -> Result<Vec<JsonValue>> {
        let records = self.read_file(path, format, csv_config)?;
        let ast = self.parser.parse(query)?;
        self.evaluator.filter_and_enrich(&ast, &records)
    }

    /// Execute stats query on a file
    pub fn query_file_stats(
        &self,
        path: &Path,
        query: &str,
        format: FileFormat,
        csv_config: &CsvConfig,
    ) -> Result<JsonValue> {
        let records = self.read_file(path, format, csv_config)?;
        let ast = self.parser.parse(query)?;

        self.evaluate_stats_query(&records, &ast, query)
    }

    /// Query multiple files in a folder
    pub fn query_folder(
        &self,
        folder_path: &Path,
        query: &str,
        pattern: &str,
        format: FileFormat,
        csv_config: &CsvConfig,
        parallel: bool,
    ) -> Result<Vec<JsonValue>> {
        let glob_pattern = folder_path.join(pattern);
        let glob_pattern = glob_pattern.to_string_lossy();

        let paths: Vec<_> = glob(&glob_pattern)
            .map_err(|e| {
                TqlError::ExecutionError(format!("Invalid glob pattern '{}': {}", pattern, e))
            })?
            .filter_map(|entry| entry.ok())
            .filter(|path| path.is_file())
            .collect();

        if paths.is_empty() {
            return Ok(Vec::new());
        }

        let ast = self.parser.parse(query)?;

        if parallel {
            let results: Result<Vec<Vec<JsonValue>>> = paths
                .par_iter()
                .map(|path| {
                    let records = self.read_file(path, format, csv_config)?;
                    let filtered = self.evaluator.filter(&ast, &records)?;
                    Ok(filtered.into_iter().cloned().collect())
                })
                .collect();

            Ok(results?.into_iter().flatten().collect())
        } else {
            let mut all_results = Vec::new();
            for path in paths {
                let records = self.read_file(&path, format, csv_config)?;
                let filtered = self.evaluator.filter(&ast, &records)?;
                all_results.extend(filtered.into_iter().cloned());
            }
            Ok(all_results)
        }
    }

    /// Query folder with stats aggregation
    pub fn query_folder_stats(
        &self,
        folder_path: &Path,
        query: &str,
        pattern: &str,
        format: FileFormat,
        csv_config: &CsvConfig,
    ) -> Result<JsonValue> {
        let glob_pattern = folder_path.join(pattern);
        let glob_pattern = glob_pattern.to_string_lossy();

        let paths: Vec<_> = glob(&glob_pattern)
            .map_err(|e| {
                TqlError::ExecutionError(format!("Invalid glob pattern '{}': {}", pattern, e))
            })?
            .filter_map(|entry| entry.ok())
            .filter(|path| path.is_file())
            .collect();

        let mut all_records = Vec::new();
        for path in paths {
            let records = self.read_file(&path, format, csv_config)?;
            all_records.extend(records);
        }

        let ast = self.parser.parse(query)?;
        self.evaluate_stats_query(&all_records, &ast, query)
    }

    /// Streaming file reader - returns an iterator over records
    pub fn stream_file(
        &self,
        path: &Path,
        format: FileFormat,
    ) -> Result<Box<dyn Iterator<Item = Result<JsonValue>>>> {
        let format = if format == FileFormat::Auto {
            FileFormat::from_path(path)
        } else {
            format
        };

        match format {
            FileFormat::JsonL => {
                let file = File::open(path).map_err(|e| {
                    TqlError::ExecutionError(format!(
                        "Failed to open file {}: {}",
                        path.display(),
                        e
                    ))
                })?;
                let reader = BufReader::new(file);
                let path_str = path.display().to_string();

                Ok(Box::new(reader.lines().enumerate().filter_map(
                    move |(line_num, line)| match line {
                        Ok(line) => {
                            let trimmed = line.trim();
                            if trimmed.is_empty() {
                                return None;
                            }
                            Some(serde_json::from_str(trimmed).map_err(|e| {
                                TqlError::ExecutionError(format!(
                                    "Failed to parse JSON at line {} in {}: {}",
                                    line_num + 1,
                                    path_str,
                                    e
                                ))
                            }))
                        }
                        Err(e) => Some(Err(TqlError::ExecutionError(format!(
                            "Failed to read line {} from {}: {}",
                            line_num + 1,
                            path_str,
                            e
                        )))),
                    },
                )))
            }
            _ => Err(TqlError::ExecutionError(
                "Streaming is only supported for JSONL format".to_string(),
            )),
        }
    }

    /// Helper to evaluate stats query
    fn evaluate_stats_query(
        &self,
        records: &[JsonValue],
        ast: &AstNode,
        query: &str,
    ) -> Result<JsonValue> {
        use crate::parser::QueryWithStatsNode;

        match ast {
            AstNode::StatsExpr(stats_node) => self.evaluate_stats_node(records, stats_node),
            AstNode::QueryWithStats(QueryWithStatsNode { filter, stats }) => {
                let filtered = self.evaluator.filter(filter, records)?;
                let owned_records: Vec<JsonValue> = filtered.iter().map(|&r| r.clone()).collect();
                self.evaluate_stats_node(&owned_records, stats)
            }
            _ => Err(TqlError::SyntaxError {
                message: "Query does not contain stats expressions".to_string(),
                position: None,
                query: Some(query.to_string()),
                suggestions: vec!["Use '| stats' to add aggregations".to_string()],
            }),
        }
    }

    /// Helper to evaluate stats node
    fn evaluate_stats_node(
        &self,
        records: &[JsonValue],
        stats_node: &crate::parser::StatsNode,
    ) -> Result<JsonValue> {
        use crate::parser::{Aggregation, GroupBy};

        let aggregations: Vec<AggregationSpec> = stats_node
            .aggregations
            .iter()
            .map(|agg: &Aggregation| AggregationSpec {
                function: agg.function.clone(),
                field: agg.field.clone().unwrap_or_else(|| "*".to_string()),
                alias: agg.alias.clone(),
                params: std::collections::HashMap::new(),
            })
            .collect();

        let group_by: Vec<String> = stats_node
            .group_by
            .iter()
            .map(|gb: &GroupBy| gb.field.clone())
            .collect();

        let stats_query = StatsQuery {
            aggregations,
            group_by,
        };

        self.stats_evaluator.evaluate_stats(records, &stats_query)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Write;
    use tempfile::TempDir;

    fn create_test_file(dir: &TempDir, name: &str, content: &str) -> std::path::PathBuf {
        let path = dir.path().join(name);
        let mut file = File::create(&path).unwrap();
        file.write_all(content.as_bytes()).unwrap();
        path
    }

    #[test]
    fn test_read_json_array() {
        let dir = TempDir::new().unwrap();
        let path = create_test_file(
            &dir,
            "data.json",
            r#"[{"name": "Alice", "age": 30}, {"name": "Bob", "age": 25}]"#,
        );

        let ops = FileOps::new();
        let records = ops
            .read_file(&path, FileFormat::Json, &CsvConfig::default())
            .unwrap();

        assert_eq!(records.len(), 2);
        assert_eq!(records[0]["name"], "Alice");
        assert_eq!(records[1]["name"], "Bob");
    }

    #[test]
    fn test_read_json_object() {
        let dir = TempDir::new().unwrap();
        let path = create_test_file(&dir, "data.json", r#"{"name": "Alice", "age": 30}"#);

        let ops = FileOps::new();
        let records = ops
            .read_file(&path, FileFormat::Json, &CsvConfig::default())
            .unwrap();

        assert_eq!(records.len(), 1);
        assert_eq!(records[0]["name"], "Alice");
    }

    #[test]
    fn test_read_jsonl() {
        let dir = TempDir::new().unwrap();
        let path = create_test_file(
            &dir,
            "data.jsonl",
            r#"{"name": "Alice", "age": 30}
{"name": "Bob", "age": 25}
{"name": "Charlie", "age": 35}"#,
        );

        let ops = FileOps::new();
        let records = ops
            .read_file(&path, FileFormat::JsonL, &CsvConfig::default())
            .unwrap();

        assert_eq!(records.len(), 3);
        assert_eq!(records[0]["name"], "Alice");
        assert_eq!(records[2]["name"], "Charlie");
    }

    #[test]
    fn test_read_csv() {
        let dir = TempDir::new().unwrap();
        let path = create_test_file(
            &dir,
            "data.csv",
            "name,age,active\nAlice,30,true\nBob,25,false",
        );

        let ops = FileOps::new();
        let records = ops
            .read_file(&path, FileFormat::Csv, &CsvConfig::default())
            .unwrap();

        assert_eq!(records.len(), 2);
        assert_eq!(records[0]["name"], "Alice");
        assert_eq!(records[0]["age"], 30);
        assert_eq!(records[0]["active"], true);
        assert_eq!(records[1]["name"], "Bob");
        assert_eq!(records[1]["active"], false);
    }

    #[test]
    fn test_query_file() {
        let dir = TempDir::new().unwrap();
        let path = create_test_file(
            &dir,
            "data.json",
            r#"[
                {"name": "Alice", "age": 30, "status": "active"},
                {"name": "Bob", "age": 25, "status": "inactive"},
                {"name": "Charlie", "age": 35, "status": "active"}
            ]"#,
        );

        let ops = FileOps::new();
        let results = ops
            .query_file(
                &path,
                "status eq 'active'",
                FileFormat::Auto,
                &CsvConfig::default(),
            )
            .unwrap();

        assert_eq!(results.len(), 2);
        assert!(results.iter().all(|r| r["status"] == "active"));
    }

    #[test]
    fn test_query_file_with_comparison() {
        let dir = TempDir::new().unwrap();
        let path = create_test_file(
            &dir,
            "data.json",
            r#"[
                {"name": "Alice", "age": 30},
                {"name": "Bob", "age": 25},
                {"name": "Charlie", "age": 35}
            ]"#,
        );

        let ops = FileOps::new();
        let results = ops
            .query_file(&path, "age > 28", FileFormat::Auto, &CsvConfig::default())
            .unwrap();

        assert_eq!(results.len(), 2);
        assert!(results.iter().all(|r| r["age"].as_i64().unwrap() > 28));
    }

    #[test]
    fn test_query_folder() {
        let dir = TempDir::new().unwrap();

        create_test_file(
            &dir,
            "data1.json",
            r#"[{"name": "Alice", "status": "active"}]"#,
        );
        create_test_file(
            &dir,
            "data2.json",
            r#"[{"name": "Bob", "status": "inactive"}, {"name": "Charlie", "status": "active"}]"#,
        );

        let ops = FileOps::new();
        let results = ops
            .query_folder(
                dir.path(),
                "status eq 'active'",
                "*.json",
                FileFormat::Auto,
                &CsvConfig::default(),
                false,
            )
            .unwrap();

        assert_eq!(results.len(), 2);
        assert!(results.iter().all(|r| r["status"] == "active"));
    }

    #[test]
    fn test_query_folder_parallel() {
        let dir = TempDir::new().unwrap();

        for i in 0..5 {
            create_test_file(
                &dir,
                &format!("data{}.json", i),
                &format!(r#"[{{"id": {}, "status": "active"}}]"#, i),
            );
        }

        let ops = FileOps::new();
        let results = ops
            .query_folder(
                dir.path(),
                "status eq 'active'",
                "*.json",
                FileFormat::Auto,
                &CsvConfig::default(),
                true,
            )
            .unwrap();

        assert_eq!(results.len(), 5);
    }

    #[test]
    fn test_format_auto_detection() {
        assert_eq!(
            FileFormat::from_path(Path::new("data.json")),
            FileFormat::Json
        );
        assert_eq!(
            FileFormat::from_path(Path::new("data.jsonl")),
            FileFormat::JsonL
        );
        assert_eq!(
            FileFormat::from_path(Path::new("data.ndjson")),
            FileFormat::JsonL
        );
        assert_eq!(
            FileFormat::from_path(Path::new("data.csv")),
            FileFormat::Csv
        );
        assert_eq!(
            FileFormat::from_path(Path::new("data.txt")),
            FileFormat::Json
        );
    }

    #[test]
    fn test_csv_custom_delimiter() {
        let dir = TempDir::new().unwrap();
        let path = create_test_file(&dir, "data.csv", "name;age\nAlice;30\nBob;25");

        let ops = FileOps::new();
        let config = CsvConfig::default().with_delimiter(';');
        let records = ops.read_file(&path, FileFormat::Csv, &config).unwrap();

        assert_eq!(records.len(), 2);
        assert_eq!(records[0]["name"], "Alice");
        assert_eq!(records[0]["age"], 30);
    }

    #[test]
    fn test_stream_jsonl() {
        let dir = TempDir::new().unwrap();
        let path = create_test_file(
            &dir,
            "data.jsonl",
            r#"{"id": 1}
{"id": 2}
{"id": 3}"#,
        );

        let ops = FileOps::new();
        let mut count = 0;
        for result in ops.stream_file(&path, FileFormat::JsonL).unwrap() {
            let record = result.unwrap();
            count += 1;
            assert!(record["id"].as_i64().unwrap() >= 1);
        }
        assert_eq!(count, 3);
    }
}