//! Structural sync walker using `rjiter`.
//!
//! Sync recursion directly over the JSON tree, driving `VariantBuilder`
//! with no intermediate state, no event queue, no byte tee. String
//! bytes live in rjiter's caller-owned buffer and are copied exactly
//! once (into VariantBuilder's own storage).
//!
//! The actual leaf walker (the per-token Variant emission) lives in
//! [`super::variant_walker`] and is shared with the `parse_json` UDF.
//! This file owns the *plan-driven* descent — top-of-document, field
//! selection, array fanout — that is specific to the
//! `json_explode` scan path.
//!
//! The caller is expected to run this inside
//! `tokio::task::spawn_blocking` and bridge any async byte stream into
//! a `std::io::Read` via `tokio_util::io::SyncIoBridge` (see
//! `super::reader`).

use std::io::Read as StdRead;

use embedded_io::Read as EmbeddedRead;
use embedded_io::Write as EmbeddedWrite;
use parquet_variant_compute::VariantArrayBuilder;
use rjiter::jiter::Peek;
use rjiter::{Error as RJiterError, RJiter};

use super::plan::NodePlan;
use crate::variant_walker::{write_subtree_to_array, WalkError};

/// Row sink abstraction for the rjiter walker.
///
/// The walker appends rows into `builder()`. After each row is
/// completed (one leaf → one row), the walker calls [`Self::after_row`].
/// The sink can use this hook to swap in a fresh builder whenever it
/// has accumulated `batch_size` rows, letting the caller flush a
/// RecordBatch mid-walk without ever buffering the whole file.
pub trait RowSink {
    /// The builder the walker currently writes into.
    fn builder(&mut self) -> &mut VariantArrayBuilder;

    /// Called exactly once per emitted row, right after the row has
    /// been appended to the current builder.
    fn after_row(&mut self) -> Result<(), WalkError>;
}

/// Trivial sink that just wraps a `&mut VariantArrayBuilder` with a
/// no-op `after_row`. Used by tests.
#[cfg(test)]
pub struct PassthroughSink<'a>(pub &'a mut VariantArrayBuilder);

#[cfg(test)]
impl<'a> RowSink for PassthroughSink<'a> {
    fn builder(&mut self) -> &mut VariantArrayBuilder {
        self.0
    }

    fn after_row(&mut self) -> Result<(), WalkError> {
        Ok(())
    }
}

/// Drive the walker to completion over a sync byte reader. Intended
/// to be called inside `tokio::task::spawn_blocking`.
///
/// `reader` supplies the full byte stream (post-decompression). The
/// walker holds a caller-owned scratch buffer of size
/// `parser_buffer_size` for rjiter's working state.
pub fn walk<R: EmbeddedRead, S: RowSink>(
    reader: &mut R,
    buf: &mut [u8],
    plan: &NodePlan,
    sink: &mut S,
) -> Result<(), WalkError> {
    let mut rjiter = RJiter::new(reader, buf);
    loop {
        match rjiter.peek() {
            Ok(peek) => {
                descend_top(&mut rjiter, plan, sink, peek)?;
            }
            Err(e) if is_eof(&e) => return Ok(()),
            Err(e) => return Err(e.into()),
        }
    }
}

/// Adapter so callers can supply a standard `std::io::Read` reader.
pub struct StdReadAdapter<R: StdRead>(pub R);

impl<R: StdRead> embedded_io::ErrorType for StdReadAdapter<R> {
    type Error = std::io::Error;
}

impl<R: StdRead> EmbeddedRead for StdReadAdapter<R> {
    fn read(&mut self, buf: &mut [u8]) -> Result<usize, Self::Error> {
        self.0.read(buf)
    }
}

fn is_eof(err: &RJiterError) -> bool {
    use rjiter::error::ErrorType;
    use rjiter::jiter::JsonErrorType;
    matches!(
        err.error_type,
        ErrorType::JsonError(JsonErrorType::EofWhileParsingValue),
    )
}

fn descend_top<R: EmbeddedRead, S: RowSink>(
    rjiter: &mut RJiter<'_, R>,
    plan: &NodePlan,
    sink: &mut S,
    peek: Peek,
) -> Result<(), WalkError> {
    match plan {
        NodePlan::Leaf => {
            write_subtree_to_array(rjiter, sink.builder(), peek)?;
            // One leaf = one emitted row. Let the sink decide whether
            // to flush a batch now.
            sink.after_row()?;
            Ok(())
        }
        NodePlan::Object { key, child } => descend_object_top(rjiter, key, child, sink, peek),
        NodePlan::Array { child } => descend_array_top(rjiter, child, sink, peek),
    }
}

fn descend_object_top<R: EmbeddedRead, S: RowSink>(
    rjiter: &mut RJiter<'_, R>,
    descend_key: &str,
    child: &NodePlan,
    sink: &mut S,
    peek: Peek,
) -> Result<(), WalkError> {
    if peek != Peek::Object {
        skip_top(rjiter, peek)?;
        return Ok(());
    }
    let mut next = rjiter.known_object()?;
    while let Some(key) = next {
        let key = key.to_string();
        let child_peek = rjiter.peek()?;
        if key == descend_key {
            descend_top(rjiter, child, sink, child_peek)?;
        } else {
            skip_top(rjiter, child_peek)?;
        }
        next = rjiter.next_key()?;
    }
    Ok(())
}

fn descend_array_top<R: EmbeddedRead, S: RowSink>(
    rjiter: &mut RJiter<'_, R>,
    child: &NodePlan,
    sink: &mut S,
    peek: Peek,
) -> Result<(), WalkError> {
    if peek != Peek::Array {
        skip_top(rjiter, peek)?;
        return Ok(());
    }
    let mut next = rjiter.known_array()?;
    while let Some(element_peek) = next {
        descend_top(rjiter, child, sink, element_peek)?;
        next = rjiter.array_step()?;
    }
    Ok(())
}

struct DiscardWriter;

impl embedded_io::ErrorType for DiscardWriter {
    type Error = std::io::Error;
}

impl EmbeddedWrite for DiscardWriter {
    fn write(&mut self, buf: &[u8]) -> Result<usize, Self::Error> {
        Ok(buf.len())
    }

    fn flush(&mut self) -> Result<(), Self::Error> {
        Ok(())
    }
}

fn skip_top<R: EmbeddedRead>(rjiter: &mut RJiter<'_, R>, peek: Peek) -> Result<(), WalkError> {
    match peek {
        Peek::Null => {
            rjiter.known_null()?;
        }
        Peek::True | Peek::False => {
            let _ = rjiter.known_bool(peek)?;
        }
        Peek::String => {
            let mut sink = DiscardWriter;
            rjiter.write_long_bytes(&mut sink)?;
        }
        Peek::Array => {
            let mut next = rjiter.known_array()?;
            while let Some(element_peek) = next {
                skip_top(rjiter, element_peek)?;
                next = rjiter.array_step()?;
            }
        }
        Peek::Object => {
            let mut next = rjiter.known_object()?;
            while next.is_some() {
                let value_peek = rjiter.peek()?;
                skip_top(rjiter, value_peek)?;
                next = rjiter.next_key()?;
            }
        }
        p if p.is_num() => {
            let _ = rjiter.known_number(peek)?;
        }
        other => return Err(WalkError::Unexpected(other)),
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use parquet_variant_compute::VariantArray;
    use std::io::Cursor;

    fn cloudtrail_plan() -> NodePlan {
        use super::super::path::Segment;
        NodePlan::compile(&[Segment::Field("Records".into()), Segment::ArrayElements])
    }

    fn run_with_buf(input: &str, plan: NodePlan, buf_size: usize) -> VariantArray {
        let mut reader = StdReadAdapter(Cursor::new(input.as_bytes().to_vec()));
        let mut buf = vec![0u8; buf_size];
        let mut builder = VariantArrayBuilder::new(16);
        let mut sink = PassthroughSink(&mut builder);
        walk(&mut reader, &mut buf, &plan, &mut sink).expect("walk ok");
        builder.build()
    }

    fn run(input: &str, plan: NodePlan) -> VariantArray {
        run_with_buf(input, plan, 4096)
    }

    /// End-to-end: a JSON with a fits-i64, a too-big-for-i64-fits-f64,
    /// and a overflows-f64 integer in the same array. Walker must produce
    /// three rows without erroring regardless of magnitude.
    #[test]
    fn bigint_numbers_fall_back_gracefully() {
        // 2^80 is beyond i64 (2^63) but well within f64 (~1.8e308);
        // "1" followed by 500 zeros is 10^500, beyond f64::MAX entirely.
        let huge_int = "1".to_string() + &"0".repeat(500);
        let input = format!(
            r#"{{"Records":[{{"v": 42}},{{"v": 1208925819614629174706176}},{{"v": {}}}]}}"#,
            huge_int,
        );
        let arr = run(&input, cloudtrail_plan());
        assert_eq!(
            arr.len(),
            3,
            "all three rows should survive numeric fallback"
        );
    }

    #[test]
    fn cloudtrail_two_records() {
        let arr = run(r#"{"Records":[{"a":1},{"a":2}]}"#, cloudtrail_plan());
        assert_eq!(arr.len(), 2);
    }

    #[test]
    fn missing_path_emits_zero_rows() {
        let arr = run(r#"{}"#, cloudtrail_plan());
        assert_eq!(arr.len(), 0);
    }

    #[test]
    fn wrong_type_at_descent_emits_zero_rows() {
        let arr = run(r#"{"Records": 1}"#, cloudtrail_plan());
        assert_eq!(arr.len(), 0);
    }

    #[test]
    fn empty_array_emits_zero_rows() {
        let arr = run(r#"{"Records": []}"#, cloudtrail_plan());
        assert_eq!(arr.len(), 0);
    }

    #[test]
    fn null_descent_value_emits_zero_rows() {
        let arr = run(r#"{"Records": null}"#, cloudtrail_plan());
        assert_eq!(arr.len(), 0);
    }

    #[test]
    fn null_leaf_emits_variant_null_row() {
        let arr = run(r#"{"Records": [null]}"#, cloudtrail_plan());
        assert_eq!(arr.len(), 1);
        assert!(!arr.is_null(0));
    }

    #[test]
    fn multiple_top_level_documents() {
        let arr = run(
            r#"{"Records":[{"a":1}]}  {"Records":[{"a":2},{"a":3}]}"#,
            cloudtrail_plan(),
        );
        assert_eq!(arr.len(), 3);
    }

    #[test]
    fn nested_otlp_plan() {
        use super::super::path::Segment;
        let plan = NodePlan::compile(&[
            Segment::Field("resourceLogs".into()),
            Segment::ArrayElements,
            Segment::Field("scopeLogs".into()),
            Segment::ArrayElements,
            Segment::Field("logRecords".into()),
            Segment::ArrayElements,
        ]);
        let input = r#"
            {
              "resourceLogs": [
                {
                  "scopeLogs": [
                    {"logRecords": [{"i":1},{"i":2}]},
                    {"logRecords": [{"i":3}]}
                  ]
                },
                {
                  "scopeLogs": [
                    {"logRecords": [{"i":4}]}
                  ]
                }
              ]
            }
        "#;
        let arr = run(input, plan);
        assert_eq!(arr.len(), 4);
    }

    #[test]
    fn many_top_level_documents_with_tiny_buffer() {
        let doc = r#"{"Records":[{"a":1},{"a":2},{"a":3}]}"#;
        let input = std::iter::repeat_n(doc, 2_000)
            .collect::<Vec<_>>()
            .join("\n");
        let arr = run_with_buf(&input, cloudtrail_plan(), 32);
        assert_eq!(arr.len(), 6_000);
    }

    #[test]
    fn skipping_large_sibling_field_survives_many_refills() {
        let big_skip = format!(
            "[{}]",
            std::iter::repeat_n(r#"{"noise":[1,2,3,4,5,6,7,8]}"#, 40)
                .collect::<Vec<_>>()
                .join(",")
        );
        let doc = format!(r#"{{"skip":{big_skip},"Records":[{{"a":1}},{{"a":2}}]}}"#);
        let input = std::iter::repeat_n(doc.as_str(), 512)
            .collect::<Vec<_>>()
            .join("\n");
        let arr = run_with_buf(&input, cloudtrail_plan(), 64);
        assert_eq!(arr.len(), 1_024);
    }
}
