//! Reader-agnostic JSON → Variant walker.
//!
//! Two callers parse JSON straight into a `VariantArrayBuilder` /
//! `ObjectBuilder` / `ListBuilder` without going through
//! `serde_json::Value`:
//!
//! 1. The streaming `json_explode` scan in [`super::walker`], which
//!    drives `RJiter` over a chunked `embedded_io::Read` byte stream.
//! 2. The `parse_json` UDF in
//!    [`crate::udf::json_to_variant`], which drives `Jiter` directly
//!    over a fully-buffered `&[u8]`.
//!
//! Both backends expose the same parser surface (`peek`, `known_*`,
//! `array_step`, `next_key`). The [`JsonReader`] trait abstracts that
//! surface so the leaf walker can be written once. Error types are
//! flattened to [`WalkError`] at the trait boundary so the walker
//! itself never has to think about which backend it has.
//!
//! The walker is intentionally a thin recursive descent — one match
//! per JSON token, write straight into the builder. String slices are
//! copied exactly once (into the builder's storage):
//!
//! - **Values** flow straight from the reader into the builder; the
//!   borrow ends synchronously inside the `append_value`/`insert` call
//!   before any further reader access, so no intermediate `String` is
//!   needed.
//! - **Keys** must outlive a recursive value walk, so each
//!   `write_object_body` frame copies the key into a single `String`
//!   scratch reused across all keys at that nesting level. Allocation
//!   count is `O(nesting_depth)`, not `O(total_keys)`.

use parquet_variant::{
    BuilderSpecificState, ListBuilder, ObjectBuilder, Variant, VariantBuilderExt,
};
use parquet_variant_compute::VariantArrayBuilder;
use rjiter::jiter::{JiterError, NumberAny, NumberInt, Peek};
use rjiter::Error as RJiterError;

/// Error returned by the JSON → Variant walker.
///
/// `Parse` flattens both `JiterError` and `RJiterError` into a single
/// string-rendered variant so the walker code does not have to be
/// generic over the underlying error type. Walker callers that need
/// to introspect the original error (e.g. [`super::walker::walk`]'s
/// EOF detection) operate on the raw reader error *before* it
/// crosses into the walker.
#[derive(Debug, thiserror::Error)]
pub enum WalkError {
    #[error("json parse error: {0}")]
    Parse(String),
    #[error("unexpected JSON token while building variant: {0:?}")]
    Unexpected(Peek),
    #[error("sink error: {0}")]
    Sink(String),
}

impl From<RJiterError> for WalkError {
    fn from(e: RJiterError) -> Self {
        WalkError::Parse(format!("{:?}", e))
    }
}

impl From<JiterError> for WalkError {
    fn from(e: JiterError) -> Self {
        WalkError::Parse(format!("{:?}", e))
    }
}

/// Backend-agnostic streaming JSON reader.
///
/// Implemented for `jiter::Jiter<'_>` (in-memory `&[u8]`) and
/// `rjiter::RJiter<'_, R>` (chunked `embedded_io::Read`). Both crates
/// expose the same parser surface; this trait lets the walker be
/// written once over either.
///
/// `&str` returned by [`JsonReader::known_str`],
/// [`JsonReader::known_object`], and [`JsonReader::next_key`] is
/// borrowed from internal parser state and may be invalidated by the
/// next call. Callers that need a value beyond the very next reader
/// call must copy it; the walker copies object keys into a per-level
/// scratch buffer and consumes string values synchronously.
pub trait JsonReader {
    fn peek(&mut self) -> Result<Peek, WalkError>;
    fn known_null(&mut self) -> Result<(), WalkError>;
    fn known_bool(&mut self, peek: Peek) -> Result<bool, WalkError>;
    fn known_str(&mut self) -> Result<&str, WalkError>;
    fn known_number(&mut self, peek: Peek) -> Result<NumberAny, WalkError>;
    fn known_object(&mut self) -> Result<Option<&str>, WalkError>;
    fn next_key(&mut self) -> Result<Option<&str>, WalkError>;
    fn known_array(&mut self) -> Result<Option<Peek>, WalkError>;
    fn array_step(&mut self) -> Result<Option<Peek>, WalkError>;
}

impl JsonReader for rjiter::jiter::Jiter<'_> {
    fn peek(&mut self) -> Result<Peek, WalkError> {
        rjiter::jiter::Jiter::peek(self).map_err(Into::into)
    }
    fn known_null(&mut self) -> Result<(), WalkError> {
        rjiter::jiter::Jiter::known_null(self).map_err(Into::into)
    }
    fn known_bool(&mut self, peek: Peek) -> Result<bool, WalkError> {
        rjiter::jiter::Jiter::known_bool(self, peek).map_err(Into::into)
    }
    fn known_str(&mut self) -> Result<&str, WalkError> {
        rjiter::jiter::Jiter::known_str(self).map_err(Into::into)
    }
    fn known_number(&mut self, peek: Peek) -> Result<NumberAny, WalkError> {
        rjiter::jiter::Jiter::known_number(self, peek).map_err(Into::into)
    }
    fn known_object(&mut self) -> Result<Option<&str>, WalkError> {
        rjiter::jiter::Jiter::known_object(self).map_err(Into::into)
    }
    fn next_key(&mut self) -> Result<Option<&str>, WalkError> {
        rjiter::jiter::Jiter::next_key(self).map_err(Into::into)
    }
    fn known_array(&mut self) -> Result<Option<Peek>, WalkError> {
        rjiter::jiter::Jiter::known_array(self).map_err(Into::into)
    }
    fn array_step(&mut self) -> Result<Option<Peek>, WalkError> {
        rjiter::jiter::Jiter::array_step(self).map_err(Into::into)
    }
}

impl<R: embedded_io::Read> JsonReader for rjiter::RJiter<'_, R> {
    fn peek(&mut self) -> Result<Peek, WalkError> {
        rjiter::RJiter::peek(self).map_err(Into::into)
    }
    fn known_null(&mut self) -> Result<(), WalkError> {
        rjiter::RJiter::known_null(self).map_err(Into::into)
    }
    fn known_bool(&mut self, peek: Peek) -> Result<bool, WalkError> {
        rjiter::RJiter::known_bool(self, peek).map_err(Into::into)
    }
    fn known_str(&mut self) -> Result<&str, WalkError> {
        rjiter::RJiter::known_str(self).map_err(Into::into)
    }
    fn known_number(&mut self, peek: Peek) -> Result<NumberAny, WalkError> {
        rjiter::RJiter::known_number(self, peek).map_err(Into::into)
    }
    fn known_object(&mut self) -> Result<Option<&str>, WalkError> {
        rjiter::RJiter::known_object(self).map_err(Into::into)
    }
    fn next_key(&mut self) -> Result<Option<&str>, WalkError> {
        rjiter::RJiter::next_key(self).map_err(Into::into)
    }
    fn known_array(&mut self) -> Result<Option<Peek>, WalkError> {
        rjiter::RJiter::known_array(self).map_err(Into::into)
    }
    fn array_step(&mut self) -> Result<Option<Peek>, WalkError> {
        rjiter::RJiter::array_step(self).map_err(Into::into)
    }
}

/// Walk one JSON value at the current reader position, appending it
/// as a single Variant row into `out`. The reader must be positioned
/// at the start of the value (caller has already called `peek` and
/// is passing it in).
pub fn write_subtree_to_array<R: JsonReader>(
    rjiter: &mut R,
    out: &mut VariantArrayBuilder,
    peek: Peek,
) -> Result<(), WalkError> {
    match peek {
        Peek::Null => {
            rjiter.known_null()?;
            out.append_value(Variant::Null);
        }
        Peek::True | Peek::False => {
            let b = rjiter.known_bool(peek)?;
            out.append_value(b);
        }
        Peek::String => {
            // Borrow ends at the end of this call; no further reader
            // access until after append_value copies into the builder.
            let s = rjiter.known_str()?;
            out.append_value(s);
        }
        Peek::Array => {
            let mut list = out.new_list();
            write_list_body(rjiter, &mut list)?;
            list.finish();
        }
        Peek::Object => {
            let mut obj = out.new_object();
            let mut key_scratch = String::with_capacity(64);
            write_object_body(rjiter, &mut obj, &mut key_scratch)?;
            obj.finish();
        }
        p if p.is_num() => {
            let n = rjiter.known_number(peek)?;
            append_number_to_sink(out, n);
        }
        other => return Err(WalkError::Unexpected(other)),
    }
    Ok(())
}

fn write_list_body<R: JsonReader, S: BuilderSpecificState>(
    rjiter: &mut R,
    list: &mut ListBuilder<'_, S>,
) -> Result<(), WalkError> {
    let mut next = rjiter.known_array()?;
    while let Some(element_peek) = next {
        write_list_element(rjiter, list, element_peek)?;
        next = rjiter.array_step()?;
    }
    Ok(())
}

fn write_list_element<R: JsonReader, S: BuilderSpecificState>(
    rjiter: &mut R,
    list: &mut ListBuilder<'_, S>,
    peek: Peek,
) -> Result<(), WalkError> {
    match peek {
        Peek::Null => {
            rjiter.known_null()?;
            list.append_value(Variant::Null);
        }
        Peek::True | Peek::False => {
            let b = rjiter.known_bool(peek)?;
            list.append_value(b);
        }
        Peek::String => {
            let s = rjiter.known_str()?;
            list.append_value(s);
        }
        Peek::Array => {
            let mut child = list.new_list();
            write_list_body(rjiter, &mut child)?;
            child.finish();
        }
        Peek::Object => {
            let mut child = list.new_object();
            let mut key_scratch = String::with_capacity(64);
            write_object_body(rjiter, &mut child, &mut key_scratch)?;
            child.finish();
        }
        p if p.is_num() => {
            let n = rjiter.known_number(peek)?;
            append_number_to_list(list, n);
        }
        other => return Err(WalkError::Unexpected(other)),
    }
    Ok(())
}

/// `key_scratch` is reused across every key at this object level —
/// the borrowed `&str` from `known_object`/`next_key` can only live
/// until the next reader call, so we copy into the scratch and pass
/// `scratch.as_str()` down to the value walker. One `String`
/// allocation per nesting level, not per key.
fn write_object_body<R: JsonReader, S: BuilderSpecificState>(
    rjiter: &mut R,
    obj: &mut ObjectBuilder<'_, S>,
    key_scratch: &mut String,
) -> Result<(), WalkError> {
    let mut next = rjiter.known_object()?;
    while let Some(key) = next {
        key_scratch.clear();
        key_scratch.push_str(key);
        let value_peek = rjiter.peek()?;
        write_object_field(rjiter, obj, key_scratch.as_str(), value_peek)?;
        next = rjiter.next_key()?;
    }
    Ok(())
}

fn write_object_field<R: JsonReader, S: BuilderSpecificState>(
    rjiter: &mut R,
    obj: &mut ObjectBuilder<'_, S>,
    key: &str,
    peek: Peek,
) -> Result<(), WalkError> {
    match peek {
        Peek::Null => {
            rjiter.known_null()?;
            obj.insert(key, Variant::Null);
        }
        Peek::True | Peek::False => {
            let b = rjiter.known_bool(peek)?;
            obj.insert(key, b);
        }
        Peek::String => {
            let s = rjiter.known_str()?;
            obj.insert(key, s);
        }
        Peek::Array => {
            let mut list = obj.new_list(key);
            write_list_body(rjiter, &mut list)?;
            list.finish();
        }
        Peek::Object => {
            let mut child = obj.new_object(key);
            let mut child_key_scratch = String::with_capacity(64);
            write_object_body(rjiter, &mut child, &mut child_key_scratch)?;
            child.finish();
        }
        p if p.is_num() => {
            let n = rjiter.known_number(peek)?;
            insert_number_to_object(obj, key, n);
        }
        other => return Err(WalkError::Unexpected(other)),
    }
    Ok(())
}

/// Coerce a parser `NumberAny` to the best Variant-compatible Rust
/// representation.
///
/// - Fits i64 → i64.
/// - BigInt that round-trips through f64 → f64 (matches
///   `parquet_variant_json::from_json::variant_from_number`). Precision
///   may be lost above 2^53, which is the same tradeoff the official
///   Variant-JSON crate makes (it has an explicit TODO for decimal
///   support). Returning f64 keeps the value numeric for downstream
///   queries even when the integer is beyond i64 range.
/// - BigInt whose string representation won't parse as f64 (rare —
///   numbers with hundreds of digits) → **the original digit string**,
///   stored as a Variant string. No data loss.
/// - f64 → f64.
fn number_to_variant(n: NumberAny) -> VariantNumber {
    match n {
        NumberAny::Int(NumberInt::Int(i)) => VariantNumber::Int(i),
        NumberAny::Int(NumberInt::BigInt(big)) => {
            let s = big.to_string();
            match s.parse::<f64>() {
                Ok(f) if f.is_finite() => VariantNumber::Float(f),
                _ => VariantNumber::BigIntString(s),
            }
        }
        NumberAny::Float(f) => VariantNumber::Float(f),
    }
}

enum VariantNumber {
    Int(i64),
    Float(f64),
    BigIntString(String),
}

fn append_number_to_sink(sink: &mut VariantArrayBuilder, n: NumberAny) {
    match number_to_variant(n) {
        VariantNumber::Int(i) => sink.append_value(i),
        VariantNumber::Float(f) => sink.append_value(f),
        VariantNumber::BigIntString(s) => sink.append_value(s.as_str()),
    }
}

fn append_number_to_list<S: BuilderSpecificState>(list: &mut ListBuilder<'_, S>, n: NumberAny) {
    match number_to_variant(n) {
        VariantNumber::Int(i) => list.append_value(i),
        VariantNumber::Float(f) => list.append_value(f),
        VariantNumber::BigIntString(s) => list.append_value(s.as_str()),
    }
}

fn insert_number_to_object<S: BuilderSpecificState>(
    obj: &mut ObjectBuilder<'_, S>,
    key: &str,
    n: NumberAny,
) {
    match number_to_variant(n) {
        VariantNumber::Int(i) => obj.insert(key, i),
        VariantNumber::Float(f) => obj.insert(key, f),
        VariantNumber::BigIntString(s) => obj.insert(key, s.as_str()),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use embedded_io::Read as EmbeddedRead;
    use parquet_variant_compute::VariantArray;
    use rjiter::jiter::Jiter;
    use rjiter::RJiter;
    use std::io::{Cursor, Read as StdRead};

    /// Local std::io::Read → embedded_io::Read adapter so the test
    /// doesn't depend on the json::walker module (which is gated
    /// behind the `executor` feature).
    struct StdReadAdapter<R: StdRead>(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 variant_from_jiter(input: &str) -> VariantArray {
        let mut j = Jiter::new(input.as_bytes());
        let peek = j.peek().expect("peek ok");
        let mut builder = VariantArrayBuilder::new(1);
        write_subtree_to_array(&mut j, &mut builder, peek).expect("walk ok");
        builder.build()
    }

    fn variant_from_rjiter(input: &str) -> VariantArray {
        let mut reader = StdReadAdapter(Cursor::new(input.as_bytes().to_vec()));
        let mut buf = vec![0u8; 4096];
        let mut rjiter = RJiter::new(&mut reader, &mut buf);
        let peek = rjiter.peek().expect("peek ok");
        let mut builder = VariantArrayBuilder::new(1);
        write_subtree_to_array(&mut rjiter, &mut builder, peek).expect("walk ok");
        builder.build()
    }

    /// Both readers must produce binary-identical variant rows for
    /// the same JSON. This is the parity guarantee the parse_json
    /// rewrite relies on: the UDF switches reader implementations
    /// without changing the on-disk shape.
    fn assert_parity(input: &str) {
        let a = variant_from_jiter(input);
        let b = variant_from_rjiter(input);
        assert_eq!(a.len(), b.len(), "row count differs for: {input}");
        let av = a.value(0);
        let bv = b.value(0);
        // Variant equality is defined; compare via debug for richer
        // diffs on failure.
        assert_eq!(format!("{:?}", av), format!("{:?}", bv), "for: {input}");
    }

    #[test]
    fn scalar_string() {
        // Bare scalars are wrapped in an array so the streaming
        // RJiter parser sees a clear end-of-value sentinel (`]`).
        // Both readers consume bare scalars correctly when followed
        // by a structural token or whitespace; they only diverge on
        // a bare scalar at EOF, which is not a shape parse_json (or
        // any production caller) ever encounters.
        assert_parity(r#"["hello"]"#);
    }

    #[test]
    fn scalar_int() {
        assert_parity(r#"[42]"#);
    }

    #[test]
    fn scalar_negative_int() {
        assert_parity(r#"[-7]"#);
    }

    #[test]
    fn scalar_float() {
        assert_parity(r#"[3.14]"#);
    }

    #[test]
    fn scalar_true() {
        assert_parity(r#"[true]"#);
    }

    #[test]
    fn scalar_false() {
        assert_parity(r#"[false]"#);
    }

    #[test]
    fn scalar_null() {
        assert_parity(r#"[null]"#);
    }

    /// `parse_json` passes a *complete* JSON document straight into
    /// `Jiter`, including bare scalar inputs like `"42"` or
    /// `"\"hi\""`. Cover those shapes against the Jiter backend
    /// specifically — they're a real production input shape (e.g.
    /// `parse_json("42")` in a Hamelin query).
    #[test]
    fn jiter_handles_bare_scalars() {
        for s in [
            r#"42"#,
            r#"-7"#,
            r#"3.14"#,
            r#""hello""#,
            r#"true"#,
            r#"false"#,
            r#"null"#,
        ] {
            let arr = variant_from_jiter(s);
            assert_eq!(arr.len(), 1, "should produce one row for: {s}");
            // Just walking through write_subtree_to_array without
            // panicking is the assertion. The shape parity with
            // append_json is covered separately by the existing
            // json_to_variant.rs scalar tests.
        }
    }

    #[test]
    fn small_object() {
        assert_parity(r#"{"a":1,"b":"x"}"#);
    }

    #[test]
    fn small_array() {
        assert_parity(r#"[1,2,3]"#);
    }

    #[test]
    fn nested_object_with_arrays() {
        assert_parity(r#"{"a":[1,2,{"b":"x"}],"c":null}"#);
    }

    #[test]
    fn empty_object() {
        assert_parity(r#"{}"#);
    }

    #[test]
    fn empty_array() {
        assert_parity(r#"[]"#);
    }

    /// A representative inner full_log payload from a Suricata
    /// stream_tcp event. This is the kind of nested object the
    /// parse_json hot path actually sees in production.
    #[test]
    fn suricata_like_payload() {
        let json = r#"{
            "timestamp": "2024-09-19T12:34:56.789Z",
            "flow_id": 1234567890123456,
            "event_type": "alert",
            "src_ip": "10.0.0.1",
            "src_port": 49152,
            "dest_ip": "10.0.0.2",
            "dest_port": 443,
            "proto": "TCP",
            "alert": {
                "action": "allowed",
                "gid": 1,
                "signature_id": 2027865,
                "rev": 3,
                "signature": "ET MALWARE Possible Malicious Activity",
                "category": "A Network Trojan was detected",
                "severity": 1
            },
            "tags": ["malware", "stream", "tcp"],
            "stream_tcp": {
                "tcp_flags": "ap",
                "tcp_flags_ts": "ap",
                "tcp_flags_tc": "a",
                "syn": false,
                "fin": false,
                "psh": true,
                "ack": true,
                "urg": false,
                "state": "established"
            }
        }"#;
        assert_parity(json);
    }

    /// Read fields back through the Variant API to confirm the
    /// walker produces a structurally-correct object — not just one
    /// that round-trips against itself.
    #[test]
    fn jiter_object_field_lookup() {
        let arr = variant_from_jiter(r#"{"name":"test","value":42,"nested":{"k":"v"}}"#);
        let v = arr.value(0);
        let obj = match v {
            Variant::Object(o) => o,
            other => panic!("expected object, got {:?}", other),
        };
        assert!(obj.get("name").is_some());
        assert!(obj.get("value").is_some());
        let nested = obj.get("nested").expect("nested present");
        match nested {
            Variant::Object(_) => {}
            other => panic!("expected nested object, got {:?}", other),
        }
    }

    /// Big integer fallback: BigInts that round-trip through f64 land
    /// as Float; BigInts that overflow f64 (parse to ±INFINITY) fall
    /// back to a string carrying the original digits, not Infinity.
    #[test]
    fn jiter_bigint_falls_back_to_float_or_string() {
        let huge_int = "1".to_string() + &"0".repeat(500);
        let input = format!(r#"[42, 1208925819614629174706176, {}]"#, huge_int);
        let arr = variant_from_jiter(&input);
        assert_eq!(arr.len(), 1);
        let v = arr.value(0);
        let list = match v {
            Variant::List(list) => list,
            other => panic!("expected list, got {:?}", other),
        };
        assert_eq!(list.len(), 3);
        match list.get(2).expect("third element") {
            Variant::String(s) => assert_eq!(s, huge_int.as_str()),
            other => panic!(
                "expected string fallback for f64-overflow BigInt, got {:?}",
                other
            ),
        }
    }
}
