//! DataFusion `FileFormat`, `FileSource`, and `FileOpener` implementations
//! for the `json` streaming scan.
//!
//! The opener runs an `rjiter::RJiter` walker inside
//! `tokio::task::spawn_blocking`. The async chunk stream from
//! `object_store` is bridged into a sync `std::io::Read` via
//! `tokio_util::io::SyncIoBridge`. The walker emits RecordBatches
//! through an mpsc channel at the session `batch_size` granularity.

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

use async_trait::async_trait;
use datafusion::arrow::array::{Array, RecordBatch};
use datafusion::arrow::datatypes::{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, StreamExt};
use futures::TryStreamExt;
use object_store::{ObjectMeta, ObjectStore, ObjectStoreExt};
use parquet_variant_compute::{VariantArrayBuilder, VariantType};
use tokio::sync::mpsc;
use tokio_util::io::{StreamReader, SyncIoBridge};

use super::path::Segment;
use super::plan::NodePlan;
use super::walker;
use crate::udf::{normalize_variant_struct, variant_data_type};
use crate::variant_walker::WalkError;

/// Default rjiter working-buffer size.
const DEFAULT_RJITER_BUFFER_SIZE: usize = 1 << 20; // 1 MiB

/// Hard cap on RecordBatch row count for `json`, regardless of
/// DataFusion's session `batch_size` (which defaults to 8192).
///
/// Variant rows are much fatter than typical tabular rows — a single
/// CloudTrail-shape Variant can easily be 5–10 KiB. At the session
/// default of 8192 rows per batch, that's 40–80 MiB in a single
/// RecordBatch; multiplied by the number of concurrent scan workers
/// and the batches queued downstream (RepartitionExec, aggregation
/// input buffers) this dominates peak RSS. Capping at 1024 rows gives
/// 5–10 MiB per batch — small enough that downstream queuing is tame,
/// large enough that per-batch overhead (record-batch construction,
/// schema fields, projection) is amortized.
const MAX_JSON_BATCH_ROWS: usize = 1024;

/// Fixed output schema: one non-nullable column `value: Variant`.
pub fn json_schema() -> SchemaRef {
    Arc::new(Schema::new(vec![Field::new(
        "value",
        variant_data_type(),
        false,
    )
    .with_extension_type(VariantType)]))
}

/// DataFusion `FileFormat` for the `json` scan.
#[derive(Clone)]
pub struct JsonFormat {
    /// Path segments retained alongside the compiled plan so the
    /// analyzer rule can extend the path on a parent visit without
    /// decoding `NodePlan` back into segments.
    segments: Arc<[Segment]>,
    plan: Arc<NodePlan>,
    compression: FileCompressionType,
    rjiter_buffer_size: usize,
    /// Lowered path, retained for display only.
    path_display: Arc<str>,
}

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

impl JsonFormat {
    pub fn new(
        segments: Vec<Segment>,
        compression: FileCompressionType,
        parser_buffer_size: Option<usize>,
        path_display: impl Into<Arc<str>>,
    ) -> Self {
        let plan = NodePlan::compile(&segments);
        Self {
            segments: Arc::from(segments),
            plan: Arc::new(plan),
            compression,
            rjiter_buffer_size: parser_buffer_size.unwrap_or(DEFAULT_RJITER_BUFFER_SIZE),
            path_display: path_display.into(),
        }
    }

    pub fn segments(&self) -> &[Segment] {
        &self.segments
    }

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

    pub fn path_display(&self) -> &str {
        &self.path_display
    }
}

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

    fn get_ext(&self) -> String {
        "".into()
    }

    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(json_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(JsonSource {
            table_schema,
            plan: Arc::clone(&self.plan),
            compression: self.compression,
            rjiter_buffer_size: self.rjiter_buffer_size,
            batch_size: None,
            metrics: ExecutionPlanMetricsSet::new(),
            projection,
        })
    }
}

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

impl FileSource for JsonSource {
    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(
                "json FileSource::create_file_opener called before with_batch_size".to_string(),
            )
        })?;
        let inner: Arc<dyn FileOpener> = Arc::new(JsonOpener {
            object_store,
            plan: Arc::clone(&self.plan),
            compression: self.compression,
            rjiter_buffer_size: self.rjiter_buffer_size,
            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.min(MAX_JSON_BATCH_ROWS));
        Arc::new(copy)
    }

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

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

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

    fn supports_repartitioning(&self) -> bool {
        // Each file decodes independently, so we can distribute files
        // across partitions. We do NOT support byte-range splitting
        // within a file — a file is one logical JSON stream.
        // Byte-range safety is enforced by overriding `repartitioned`
        // below so the default size-balanced partitioner (which emits
        // mid-file ranges) is never used.
        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(
            "json",
            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 JsonOpener {
    object_store: Arc<dyn ObjectStore>,
    plan: Arc<NodePlan>,
    compression: FileCompressionType,
    rjiter_buffer_size: usize,
    batch_size: usize,
    schema: SchemaRef,
}

impl FileOpener for JsonOpener {
    fn open(&self, partitioned_file: PartitionedFile) -> Result<FileOpenFuture> {
        let object_store = Arc::clone(&self.object_store);
        let plan = Arc::clone(&self.plan);
        let compression = self.compression;
        let rjiter_buffer_size = self.rjiter_buffer_size;
        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 = compression
                .convert_stream(byte_stream.map_err(DataFusionError::from).boxed())?
                .map_err(std::io::Error::other)
                .boxed();

            // Yields one RecordBatch per ~batch_size rows, rather than
            // buffering the whole file into one terminal batch. Peak
            // memory is bounded by one batch in flight + walker scratch.
            let stream = rjiter_stream(decompressed, plan, rjiter_buffer_size, batch_size, schema);
            Ok(stream.boxed())
        }))
    }
}

/// Build a RecordBatch from an accumulated VariantArrayBuilder, swapping
/// in a fresh one. Returns `None` if there are no rows (empty batch).
fn flush_builder(
    builder: &mut VariantArrayBuilder,
    batch_size: usize,
    schema: &SchemaRef,
) -> Result<Option<RecordBatch>, DataFusionError> {
    let fresh = VariantArrayBuilder::new(batch_size);
    let done = std::mem::replace(builder, fresh);
    let variant_array = done.build();
    let struct_arr = normalize_variant_struct(variant_array.into_inner());
    if struct_arr.len() == 0 {
        return Ok(None);
    }
    let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(struct_arr)])
        .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?;
    Ok(Some(batch))
}

/// Drive the rjiter walker by bridging the async chunk stream into a
/// sync `std::io::Read`, running the walker inside `spawn_blocking`,
/// and piping batches out via an mpsc channel. Yields batches as the
/// blocking walker produces them.
///
/// # Runtime mechanics (why the sync/async crossings look unusual)
///
/// This function lives at an async → sync → async boundary:
///
/// 1. **DataFusion hands us a `BoxStream<Bytes>`** (async) — the
///    decompressed object-store stream.
/// 2. **rjiter is sync**. To feed it from the async stream we wrap the
///    stream in a `StreamReader` (which impls `tokio::io::AsyncRead`)
///    and then in a `SyncIoBridge` (which turns that into
///    `std::io::Read`). `SyncIoBridge` captures a tokio runtime
///    handle at construction time and uses it to drive the inner
///    async reader — that is why this whole thing must run inside
///    `tokio::task::spawn_blocking`, which guarantees a runtime is
///    entered on the worker thread.
/// 3. **The walker hands batches back out** via a `tokio::sync::mpsc`
///    channel. Because the walker is running on a blocking thread, it
///    uses `blocking_send` to push batches back to the async side while
///    preserving bounded backpressure.
///
/// The sync/async crossings are still load-bearing, but there is now
/// only one runtime story: `SyncIoBridge` uses tokio to drive reads,
/// and the outbound channel is also tokio-native.
fn rjiter_stream(
    decompressed: futures::stream::BoxStream<'static, Result<bytes::Bytes, std::io::Error>>,
    plan: Arc<NodePlan>,
    rjiter_buffer_size: usize,
    batch_size: usize,
    schema: SchemaRef,
) -> impl futures::Stream<Item = Result<RecordBatch, DataFusionError>> + Send + 'static {
    let (tx, rx) = mpsc::channel::<Result<RecordBatch, DataFusionError>>(2);

    let async_read = StreamReader::new(decompressed);
    // Clone of the sender so that if the walker's blocking thread panics,
    // we still have a handle to push the error onto the stream instead of
    // letting tokio silently absorb the panic (which would close the
    // channel and look to DataFusion like a clean end-of-stream —
    // i.e. silent data truncation).
    let panic_tx = tx.clone();
    tokio::task::spawn_blocking(move || {
        let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
            run_rjiter_blocking(async_read, plan, rjiter_buffer_size, batch_size, schema, tx)
        }));
        match outcome {
            // Receiver disconnect (query cancellation, downstream operator
            // dropped the stream) is the only reason any of the sends in
            // `run_rjiter_blocking` can fail — there's nothing to tell
            // because the consumer is already gone.
            Ok(Ok(())) => {}
            Ok(Err(WalkerChannelClosed)) => {}
            // Panic inside the walker. Don't let tokio swallow it silently:
            // turn it into a DataFusionError and push onto the stream so
            // DataFusion surfaces it instead of ending the scan cleanly.
            // If even this push fails, the receiver is already gone, and
            // there is no caller left to tell.
            Err(panic) => {
                let msg = panic
                    .downcast_ref::<&'static str>()
                    .map(|s| (*s).to_string())
                    .or_else(|| panic.downcast_ref::<String>().cloned())
                    .unwrap_or_else(|| "<non-string panic payload>".to_string());
                let err = DataFusionError::External(format!("json walker panicked: {msg}").into());
                match send_to_channel(&panic_tx, Err(err)) {
                    Ok(()) => {}
                    Err(WalkerChannelClosed) => {}
                }
            }
        }
    });

    stream::unfold(rx, |mut rx| async move {
        rx.recv().await.map(|item| (item, rx))
    })
}

/// Error type returned when the outbound mpsc receiver has been dropped
/// (the only way any of the walker's sends can fail).
#[derive(Debug, thiserror::Error)]
#[error("json walker output channel disconnected")]
struct WalkerChannelClosed;

fn run_rjiter_blocking(
    async_read: StreamReader<
        futures::stream::BoxStream<'static, Result<bytes::Bytes, std::io::Error>>,
        bytes::Bytes,
    >,
    plan: Arc<NodePlan>,
    rjiter_buffer_size: usize,
    batch_size: usize,
    schema: SchemaRef,
    tx: mpsc::Sender<Result<RecordBatch, DataFusionError>>,
) -> Result<(), WalkerChannelClosed> {
    let bridge = SyncIoBridge::new(async_read);
    let mut adapter = walker::StdReadAdapter(bridge);
    let mut buf = vec![0u8; rjiter_buffer_size];

    // Batching sink: flushes a batch every `batch_size` rows. When the
    // sink decides a batch is ready it block_on-sends into the mpsc
    // channel for the async side to pick up. Because the channel is
    // bounded (capacity 2), the walker back-pressures on fast
    // downstream consumption — we never buffer more than ~2 batches
    // across the async boundary.
    let mut sink = BatchingSink {
        builder: VariantArrayBuilder::new(batch_size),
        batch_size,
        rows_in_builder: 0,
        schema: schema.clone(),
        tx: &tx,
    };

    let result = walker::walk(&mut adapter, &mut buf, &plan, &mut sink);

    if let Err(e) = result {
        return send_to_channel(&tx, Err(DataFusionError::External(Box::new(e))));
    }

    // Flush the final partial batch (if any).
    let tail = flush_builder(&mut sink.builder, batch_size, &schema);
    match tail {
        Ok(Some(batch)) => send_to_channel(&tx, Ok(batch)),
        Ok(None) => Ok(()),
        Err(e) => send_to_channel(&tx, Err(e)),
    }
}

/// Blocking send into the outbound channel. Converts a
/// receiver-disconnected error into `WalkerChannelClosed`.
fn send_to_channel(
    tx: &mpsc::Sender<Result<RecordBatch, DataFusionError>>,
    item: Result<RecordBatch, DataFusionError>,
) -> Result<(), WalkerChannelClosed> {
    tx.blocking_send(item).map_err(|_| WalkerChannelClosed)
}

/// Row sink used by the rjiter blocking walker. Flushes a RecordBatch
/// into the outbound channel whenever the builder reaches `batch_size`
/// rows; always flushes at leaf boundaries (never mid-row).
struct BatchingSink<'tx> {
    builder: VariantArrayBuilder,
    batch_size: usize,
    rows_in_builder: usize,
    schema: SchemaRef,
    tx: &'tx mpsc::Sender<Result<RecordBatch, DataFusionError>>,
}

impl<'tx> walker::RowSink for BatchingSink<'tx> {
    fn builder(&mut self) -> &mut VariantArrayBuilder {
        &mut self.builder
    }

    fn after_row(&mut self) -> Result<(), WalkError> {
        self.rows_in_builder += 1;
        if self.batch_size != 0 && self.rows_in_builder >= self.batch_size {
            let maybe_batch = flush_builder(&mut self.builder, self.batch_size, &self.schema)
                .map_err(|e| WalkError::Sink(e.to_string()))?;
            self.rows_in_builder = 0;
            if let Some(batch) = maybe_batch {
                send_to_channel(self.tx, Ok(batch)).map_err(|e| WalkError::Sink(e.to_string()))?;
            }
        }
        Ok(())
    }
}
