//! DataFusion `FileFormat`, `FileSource`, and `FileOpener` for the
//! `lines` scan. Reads bytes line-by-line and emits each line as one
//! non-nullable `Utf8View` row in a column named `line`.
//!
//! Pure async: the decompressed `BoxStream<Bytes>` is consumed
//! directly, lines are split on `\n` (CRLF normalized), and full
//! batches are emitted whenever `batch_size` rows have accumulated.
//! An unterminated final line is emitted as a row, matching CSV.

use std::any::Any;
use std::fmt::{self, Formatter};
use std::sync::Arc;

use async_trait::async_trait;
use datafusion::arrow::array::builder::StringViewBuilder;
use datafusion::arrow::array::{Array, RecordBatch};
use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use datafusion::catalog::Session;
use datafusion::common::{GetExt, Result, Statistics};
use datafusion::datasource::file_format::file_compression_type::FileCompressionType;
use datafusion::datasource::file_format::FileFormat;
use datafusion::datasource::listing::PartitionedFile;
use datafusion::datasource::physical_plan::{
    FileOpenFuture, FileOpener, FileScanConfig, FileScanConfigBuilder, FileSource,
};
use datafusion::datasource::source::DataSourceExec;
use datafusion::datasource::table_schema::TableSchema;
use datafusion::error::DataFusionError;
use datafusion::physical_expr::LexOrdering;
use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet;
use datafusion::physical_plan::{DisplayFormatType, ExecutionPlan};
use datafusion_datasource::projection::{ProjectionOpener, SplitProjection};
use datafusion_physical_expr::projection::ProjectionExprs;
use futures::stream::{self, BoxStream, StreamExt};
use futures::TryStreamExt;
use object_store::{ObjectMeta, ObjectStore, ObjectStoreExt};

/// Fixed output schema: one non-nullable column `line: Utf8View`.
pub fn lines_schema() -> SchemaRef {
    Arc::new(Schema::new(vec![Field::new(
        "line",
        DataType::Utf8View,
        false,
    )]))
}

/// DataFusion `FileFormat` for the `lines` scan.
#[derive(Clone)]
pub struct LinesFormat {
    compression: FileCompressionType,
    /// File-extension hint (e.g. `".log"`); empty matches all files.
    file_extension: Arc<str>,
}

impl fmt::Debug for LinesFormat {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        f.debug_struct("LinesFormat")
            .field("compression", &self.compression)
            .field("file_extension", &self.file_extension)
            .finish()
    }
}

impl LinesFormat {
    pub fn new(compression: FileCompressionType, file_extension: impl Into<Arc<str>>) -> Self {
        Self {
            compression,
            file_extension: file_extension.into(),
        }
    }

    pub fn compression(&self) -> FileCompressionType {
        self.compression
    }
}

#[async_trait]
impl FileFormat for LinesFormat {
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn get_ext(&self) -> String {
        self.file_extension.to_string()
    }

    fn get_ext_with_compression(
        &self,
        file_compression_type: &FileCompressionType,
    ) -> Result<String> {
        Ok(file_compression_type.get_ext())
    }

    fn compression_type(&self) -> Option<FileCompressionType> {
        Some(self.compression)
    }

    async fn infer_schema(
        &self,
        _state: &dyn Session,
        _store: &Arc<dyn ObjectStore>,
        _objects: &[ObjectMeta],
    ) -> Result<SchemaRef> {
        Ok(lines_schema())
    }

    async fn infer_stats(
        &self,
        _state: &dyn Session,
        _store: &Arc<dyn ObjectStore>,
        table_schema: SchemaRef,
        _object: &ObjectMeta,
    ) -> Result<Statistics> {
        Ok(Statistics::new_unknown(&table_schema))
    }

    async fn create_physical_plan(
        &self,
        _state: &dyn Session,
        conf: FileScanConfig,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        let conf = FileScanConfigBuilder::from(conf)
            .with_file_compression_type(self.compression)
            .build();
        Ok(DataSourceExec::from_data_source(conf))
    }

    fn file_source(&self, table_schema: TableSchema) -> Arc<dyn FileSource> {
        let projection = SplitProjection::unprojected(&table_schema);
        Arc::new(LinesSource {
            table_schema,
            compression: self.compression,
            batch_size: None,
            metrics: ExecutionPlanMetricsSet::new(),
            projection,
        })
    }
}

#[derive(Clone)]
struct LinesSource {
    table_schema: TableSchema,
    compression: FileCompressionType,
    batch_size: Option<usize>,
    metrics: ExecutionPlanMetricsSet,
    projection: SplitProjection,
}

impl FileSource for LinesSource {
    fn create_file_opener(
        &self,
        object_store: Arc<dyn ObjectStore>,
        _base_config: &FileScanConfig,
        _partition: usize,
    ) -> Result<Arc<dyn FileOpener>> {
        let batch_size = self.batch_size.ok_or_else(|| {
            DataFusionError::Internal(
                "lines FileSource::create_file_opener called before with_batch_size".to_string(),
            )
        })?;
        let inner: Arc<dyn FileOpener> = Arc::new(LinesOpener {
            object_store,
            compression: self.compression,
            batch_size,
            schema: self.table_schema.file_schema().clone(),
        });
        let opener = ProjectionOpener::try_new(
            self.projection.clone(),
            inner,
            self.table_schema.file_schema(),
        )?;
        Ok(opener)
    }

    fn as_any(&self) -> &dyn Any {
        self
    }

    fn table_schema(&self) -> &TableSchema {
        &self.table_schema
    }

    fn with_batch_size(&self, batch_size: usize) -> Arc<dyn FileSource> {
        let mut copy = self.clone();
        copy.batch_size = Some(batch_size);
        Arc::new(copy)
    }

    fn metrics(&self) -> &ExecutionPlanMetricsSet {
        &self.metrics
    }

    fn file_type(&self) -> &str {
        "lines"
    }

    fn fmt_extra(&self, _t: DisplayFormatType, _f: &mut Formatter) -> fmt::Result {
        Ok(())
    }

    fn supports_repartitioning(&self) -> bool {
        // Each file decodes independently; we do NOT split a single file
        // mid-stream. File-granularity-only repartitioning is enforced
        // by overriding `repartitioned` below.
        true
    }

    fn repartitioned(
        &self,
        target_partitions: usize,
        repartition_file_min_size: usize,
        _output_ordering: Option<LexOrdering>,
        config: &FileScanConfig,
    ) -> Result<Option<FileScanConfig>> {
        crate::file_scan_common::lpt_repartition(
            "lines",
            target_partitions,
            repartition_file_min_size,
            config,
        )
    }

    fn projection(&self) -> Option<&ProjectionExprs> {
        Some(&self.projection.source)
    }

    fn try_pushdown_projection(
        &self,
        projection: &ProjectionExprs,
    ) -> Result<Option<Arc<dyn FileSource>>> {
        let mut source = self.clone();
        let merged = self.projection.source.try_merge(projection)?;
        source.projection = SplitProjection::new(self.table_schema.file_schema(), &merged);
        Ok(Some(Arc::new(source)))
    }
}

struct LinesOpener {
    object_store: Arc<dyn ObjectStore>,
    compression: FileCompressionType,
    batch_size: usize,
    schema: SchemaRef,
}

impl FileOpener for LinesOpener {
    fn open(&self, partitioned_file: PartitionedFile) -> Result<FileOpenFuture> {
        let object_store = Arc::clone(&self.object_store);
        let compression = self.compression;
        let batch_size = self.batch_size;
        let schema = Arc::clone(&self.schema);

        Ok(Box::pin(async move {
            let get = object_store
                .get(&partitioned_file.object_meta.location)
                .await?;
            let byte_stream = get.into_stream().map_err(std::io::Error::other);

            let decompressed: BoxStream<
                'static,
                std::result::Result<bytes::Bytes, std::io::Error>,
            > = compression
                .convert_stream(byte_stream.map_err(DataFusionError::from).boxed())?
                .map_err(std::io::Error::other)
                .boxed();

            Ok(lines_stream(decompressed, batch_size, schema).boxed())
        }))
    }
}

/// Async-native lines reader. Drains the chunk stream, splits on `\n`
/// (CRLF normalized), and yields a `RecordBatch` whenever `batch_size`
/// rows have accumulated. Trailing partial line is emitted at EOF.
fn lines_stream(
    decompressed: BoxStream<'static, std::result::Result<bytes::Bytes, std::io::Error>>,
    batch_size: usize,
    schema: SchemaRef,
) -> impl futures::Stream<Item = Result<RecordBatch, DataFusionError>> + Send + 'static {
    struct State {
        upstream: BoxStream<'static, std::result::Result<bytes::Bytes, std::io::Error>>,
        /// Bytes after the last `\n` we've seen so far — an in-progress line.
        leftover: Vec<u8>,
        builder: StringViewBuilder,
        rows: usize,
        schema: SchemaRef,
        batch_size: usize,
        eof_seen: bool,
    }

    let init = State {
        upstream: decompressed,
        leftover: Vec::with_capacity(256),
        builder: StringViewBuilder::with_capacity(batch_size),
        rows: 0,
        schema,
        batch_size,
        eof_seen: false,
    };

    stream::unfold(init, |mut st| async move {
        if st.eof_seen && st.rows == 0 {
            return None;
        }

        // Pull chunks and slice them into rows until the batch fills or we hit EOF.
        while st.rows < st.batch_size && !st.eof_seen {
            // Drain any newlines already buffered in `leftover` from a
            // prior batch-fill mid-chunk. Without this, a chunk tail
            // carried over from the previous outer iteration may contain
            // multiple `\n` bytes, and the for-loop below would treat
            // the whole prefix as one line — silently merging rows.
            match drain_leftover_lines(
                &mut st.leftover,
                &mut st.builder,
                &mut st.rows,
                st.batch_size,
            ) {
                Ok(()) => {}
                Err(e) => return Some((Err(e), st)),
            }
            if st.rows >= st.batch_size {
                break;
            }

            match st.upstream.next().await {
                Some(Ok(chunk)) => {
                    let mut start = 0usize;
                    for (i, &b) in chunk.iter().enumerate() {
                        if b != b'\n' {
                            continue;
                        }
                        // Found a line terminator. The line is leftover ++
                        // chunk[start..i], minus a single trailing \r.
                        let cr_in_chunk = i > start && chunk[i - 1] == b'\r';
                        let chunk_end = if cr_in_chunk { i - 1 } else { i };
                        let line_bytes: &[u8] = if st.leftover.is_empty() {
                            &chunk[start..chunk_end]
                        } else {
                            st.leftover.extend_from_slice(&chunk[start..chunk_end]);
                            // CR may have ended up as the last byte of leftover
                            // (CR was in the previous chunk, LF in this one).
                            // Only pop when CR wasn't already stripped from the
                            // chunk side — otherwise a content \r at the end of
                            // leftover would be silently dropped.
                            if !cr_in_chunk && st.leftover.last() == Some(&b'\r') {
                                st.leftover.pop();
                            }
                            &st.leftover[..]
                        };
                        match std::str::from_utf8(line_bytes) {
                            Ok(s) => st.builder.append_value(s),
                            Err(e) => {
                                return Some((Err(DataFusionError::External(Box::new(e))), st));
                            }
                        }
                        st.leftover.clear();
                        st.rows += 1;
                        start = i + 1;
                        if st.rows >= st.batch_size {
                            // Carry the rest of this chunk into leftover. The
                            // tail may contain more `\n` bytes; the
                            // `drain_leftover_lines` step above will split them
                            // out at the start of the next outer iteration.
                            st.leftover.extend_from_slice(&chunk[start..]);
                            start = chunk.len();
                            break;
                        }
                    }
                    // Bytes after the final `\n` (or the whole chunk if none) carry over.
                    if start < chunk.len() {
                        st.leftover.extend_from_slice(&chunk[start..]);
                    }
                }
                Some(Err(e)) => {
                    return Some((Err(DataFusionError::IoError(e)), st));
                }
                None => {
                    st.eof_seen = true;
                    // Drain any remaining buffered newlines before the
                    // trailing-partial-line flush below.
                    match drain_leftover_lines(
                        &mut st.leftover,
                        &mut st.builder,
                        &mut st.rows,
                        st.batch_size,
                    ) {
                        Ok(()) => {}
                        Err(e) => return Some((Err(e), st)),
                    }
                    // Flush trailing partial line as a row, matching CSV.
                    if !st.leftover.is_empty() {
                        let line = std::mem::take(&mut st.leftover);
                        match std::str::from_utf8(&line) {
                            Ok(s) => {
                                st.builder.append_value(s);
                                st.rows += 1;
                            }
                            Err(e) => {
                                return Some((Err(DataFusionError::External(Box::new(e))), st));
                            }
                        }
                    }
                }
            }
        }

        if st.rows == 0 {
            return None;
        }
        match build_batch(&mut st.builder, &st.schema) {
            Ok(batch) => {
                st.rows = 0;
                Some((Ok(batch), st))
            }
            Err(e) => Some((Err(e), st)),
        }
    })
}

fn build_batch(
    builder: &mut StringViewBuilder,
    schema: &SchemaRef,
) -> Result<RecordBatch, DataFusionError> {
    let arr = builder.finish();
    debug_assert!(arr.len() > 0);
    RecordBatch::try_new(schema.clone(), vec![Arc::new(arr)])
        .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))
}

/// Pop complete lines off the front of `leftover`, appending each as a
/// row, until either no `\n` remains or the batch is full. CRLF
/// endings are normalized by stripping one trailing `\r` per line.
///
/// Required because, when a batch fills mid-chunk, the chunk tail is
/// dumped wholesale into `leftover` — and that tail can contain
/// further newlines. Without this drain, the next outer iteration
/// would treat the tail as one contiguous prefix and merge several
/// real lines into one.
fn drain_leftover_lines(
    leftover: &mut Vec<u8>,
    builder: &mut StringViewBuilder,
    rows: &mut usize,
    batch_size: usize,
) -> Result<(), DataFusionError> {
    while *rows < batch_size {
        let Some(nl) = leftover.iter().position(|&b| b == b'\n') else {
            break;
        };
        let end = if nl > 0 && leftover[nl - 1] == b'\r' {
            nl - 1
        } else {
            nl
        };
        let s = std::str::from_utf8(&leftover[..end])
            .map_err(|e| DataFusionError::External(Box::new(e)))?;
        builder.append_value(s);
        *rows += 1;
        leftover.drain(..=nl);
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use bytes::Bytes;
    use datafusion::arrow::array::StringViewArray;
    use futures::stream::TryStreamExt;
    use tokio::runtime::Runtime;

    /// Collect every `line` cell from every batch yielded by
    /// `lines_stream` over the supplied chunks.
    fn collect_lines(chunks: Vec<&'static [u8]>, batch_size: usize) -> Vec<String> {
        let rt = Runtime::new().expect("tokio runtime");
        rt.block_on(async move {
            let upstream = stream::iter(
                chunks
                    .into_iter()
                    .map(|c| Ok::<_, std::io::Error>(Bytes::from_static(c))),
            )
            .boxed();
            let stream = lines_stream(upstream, batch_size, lines_schema());
            let batches: Vec<RecordBatch> = stream.try_collect().await.expect("stream ok");
            let mut out = Vec::new();
            for batch in batches {
                let arr = batch
                    .column(0)
                    .as_any()
                    .downcast_ref::<StringViewArray>()
                    .expect("Utf8View column");
                for i in 0..arr.len() {
                    out.push(arr.value(i).to_string());
                }
            }
            out
        })
    }

    /// Regression test: with `batch_size = 1` the batch fills on the
    /// very first `\n` of the chunk. Before the leftover-drain fix the
    /// rest of the chunk was buffered with embedded `\n` bytes intact,
    /// causing the next `\n` in the next chunk to merge several real
    /// lines into one. With the fix all three lines must come through.
    #[test]
    fn batch_fill_mid_chunk_does_not_drop_lines() {
        let lines = collect_lines(vec![b"line0\nline1\nline2\n"], 1);
        assert_eq!(lines, vec!["line0", "line1", "line2"]);
    }

    /// Same shape but with `batch_size = 2` — a batch fills after
    /// `line1` and the trailing `line2\n` is carried over as
    /// leftover. The drain must split it out on the next iteration.
    #[test]
    fn batch_fill_with_trailing_complete_line() {
        let lines = collect_lines(vec![b"line0\nline1\nline2\n"], 2);
        assert_eq!(lines, vec!["line0", "line1", "line2"]);
    }

    /// CRLF line endings inside the carried-over chunk tail must
    /// still be normalized to bare lines.
    #[test]
    fn batch_fill_mid_chunk_strips_cr_in_drain() {
        let lines = collect_lines(vec![b"a\r\nb\r\nc\r\n"], 1);
        assert_eq!(lines, vec!["a", "b", "c"]);
    }

    /// Trailing partial line at EOF still emits a row.
    #[test]
    fn unterminated_final_line_is_emitted() {
        let lines = collect_lines(vec![b"a\nb"], 16);
        assert_eq!(lines, vec!["a", "b"]);
    }

    /// CR at the end of one chunk and LF at the start of the next is
    /// still treated as a single CRLF line ending.
    #[test]
    fn cr_lf_split_across_chunks() {
        let lines = collect_lines(vec![b"a\r", b"\nb\n"], 16);
        assert_eq!(lines, vec!["a", "b"]);
    }

    /// Direct repro of the original bug across two chunks. With
    /// `batch_size = 1`, the first chunk yields `line0`, then dumps
    /// `b"line1\nline2\n"` into leftover. The next chunk arrives with
    /// fresh content. Without the leftover-drain at the top of the
    /// loop, the for-loop's first `\n` (in `line3\n`) would merge
    /// `line1\nline2\n` + `line3` into ONE row.
    #[test]
    fn batch_fill_mid_chunk_then_next_chunk_does_not_merge_lines() {
        let lines = collect_lines(vec![b"line0\nline1\nline2\n", b"line3\n"], 1);
        assert_eq!(lines, vec!["line0", "line1", "line2", "line3"]);
    }

    /// Leftover holds multiple complete lines AND a partial-line tail
    /// when the next chunk arrives. The drain emits the complete
    /// lines, the chunk loop stitches the partial-line tail with the
    /// chunk prefix.
    #[test]
    fn leftover_with_complete_lines_and_partial_tail() {
        // First chunk: batch_size=1, fills after `a`. Tail
        // `b\nc\npartial` lands in leftover. Next chunk: `_done\nd\n`.
        // Expected: a, b, c, partial_done, d.
        let lines = collect_lines(vec![b"a\nb\nc\npartial", b"_done\nd\n"], 1);
        assert_eq!(lines, vec!["a", "b", "c", "partial_done", "d"]);
    }

    /// Many lines across multiple batches with `batch_size = 2`.
    /// Catches off-by-one merging that would only show up across
    /// several batch flushes.
    #[test]
    fn many_lines_across_multiple_batches() {
        let lines = collect_lines(vec![b"a\nb\nc\nd\ne\nf\ng\nh\n"], 2);
        assert_eq!(lines, vec!["a", "b", "c", "d", "e", "f", "g", "h"]);
    }

    /// A literal `\r` in line content followed by a CRLF terminator,
    /// split such that the content `\r` lands in leftover and the
    /// terminator `\r\n` lands in the next chunk. Previously the CR
    /// was stripped from chunk_end AND popped from leftover, dropping
    /// a content byte and producing `hello` instead of `hello\r`.
    #[test]
    fn content_cr_followed_by_crlf_split_across_chunks() {
        let lines = collect_lines(vec![b"hello\r", b"\r\nnext\n"], 16);
        assert_eq!(lines, vec!["hello\r", "next"]);
    }
}
