use crate::Result;
use crate::readers::{
    OnceReader, RecordingBufRead, ReplayBufRead, batch_size::resolve_batch_size,
    dataframe_from_stream, emits_single_batch, finish_stream, stream_record_batches,
};
use autark_dataframe::DataFrame;
use autark_enhanced_reader::{autoread, autoread_to_bytes};
use serde_json::Value;
use std::io::{BufRead, BufReader, Cursor};
use std::sync::Arc;

use arrow::datatypes::Schema;
use arrow_json::{
    ReaderBuilder,
    reader::{infer_json_schema, infer_json_schema_from_seekable},
};

const UTF8_BOM: &[u8] = &[0xEF, 0xBB, 0xBF];
const DEFAULT_BATCH_SIZE: usize = 1 << 16;

#[inline]
fn strip_utf8_bom(bytes: &[u8]) -> &[u8] {
    bytes.strip_prefix(UTF8_BOM).unwrap_or(bytes)
}

#[inline]
fn first_non_whitespace(bytes: &[u8]) -> Option<u8> {
    bytes
        .iter()
        .copied()
        .find(|byte| !byte.is_ascii_whitespace())
}

#[inline]
fn maybe_parse_array_rows(bytes: &[u8]) -> Option<Vec<Value>> {
    (first_non_whitespace(bytes) == Some(b'['))
        .then(|| serde_json::from_slice::<Vec<Value>>(bytes).ok())
        .flatten()
}

fn to_ndjson(rows: Vec<Value>, capacity_hint: usize) -> Result<Vec<u8>> {
    rows.into_iter().try_fold(
        Vec::with_capacity(capacity_hint),
        |mut acc, row| -> Result<Vec<u8>> {
            if !row.is_object() {
                return Err(crate::Error::Reader(
                    "JSON array input must contain only objects".to_string(),
                ));
            }
            serde_json::to_writer(&mut acc, &row)
                .map_err(|err| crate::Error::Reader(format!("invalid JSON row: {err}")))?;
            acc.push(b'\n');
            Ok(acc)
        },
    )
}

pub struct JsonReader {
    source: JsonSource,
    batch_size: usize,
    batch_size_explicit: bool,
    schema: std::sync::OnceLock<Arc<Schema>>,
    normalized_bytes: std::sync::OnceLock<Vec<u8>>,
}

enum JsonSource {
    Uri(String),
    Inline(String),
}

#[derive(Clone, Copy, PartialEq, Eq)]
enum JsonMode {
    Ndjson,
    Array,
}

impl JsonReader {
    pub fn new(uri: &str) -> Result<JsonReader> {
        Ok(JsonReader {
            source: JsonSource::Uri(uri.to_string()),
            batch_size: DEFAULT_BATCH_SIZE,
            batch_size_explicit: false,
            schema: std::sync::OnceLock::new(),
            normalized_bytes: std::sync::OnceLock::new(),
        })
    }

    pub fn from_json_string(json: String) -> Result<JsonReader> {
        Ok(JsonReader {
            source: JsonSource::Inline(json),
            batch_size: DEFAULT_BATCH_SIZE,
            batch_size_explicit: false,
            schema: std::sync::OnceLock::new(),
            normalized_bytes: std::sync::OnceLock::new(),
        })
    }

    pub fn with_batch_size(mut self, batch_size: usize) -> JsonReader {
        self.batch_size = batch_size;
        self.batch_size_explicit = true;
        self
    }

    fn ensure_batch_size(&mut self) -> Result<()> {
        if self.batch_size_explicit {
            return Ok(());
        }

        self.batch_size = resolve_batch_size(DEFAULT_BATCH_SIZE);
        Ok(())
    }

    fn bytes(&self) -> Result<Vec<u8>> {
        match &self.source {
            JsonSource::Uri(uri) => autoread_to_bytes(uri),
            JsonSource::Inline(json) => Ok(json.as_bytes().to_vec()),
        }
    }

    fn ndjson_reader(&self) -> Result<Box<dyn BufRead + Send>> {
        match &self.source {
            JsonSource::Uri(uri) => autoread(uri),
            JsonSource::Inline(json) => Ok(Box::new(BufReader::new(Cursor::new(
                json.as_bytes().to_vec(),
            )))),
        }
    }

    fn mode(&self) -> Result<JsonMode> {
        match &self.source {
            JsonSource::Inline(json) => Ok(
                match first_non_whitespace(strip_utf8_bom(json.as_bytes())) {
                    Some(b'[') => JsonMode::Array,
                    _ => JsonMode::Ndjson,
                },
            ),
            JsonSource::Uri(uri) => {
                let mut reader = autoread(uri)?;
                Ok(match first_non_whitespace_in_reader(&mut reader)? {
                    Some(b'[') => JsonMode::Array,
                    _ => JsonMode::Ndjson,
                })
            }
        }
    }

    fn normalize_bytes(&self) -> Result<Vec<u8>> {
        let bytes = self.bytes()?;
        match maybe_parse_array_rows(strip_utf8_bom(&bytes)) {
            Some(rows) => to_ndjson(rows, bytes.len()),
            None => Ok(bytes),
        }
    }

    fn normalized_bytes_cached(&self) -> Result<&[u8]> {
        if let Some(bytes) = self.normalized_bytes.get() {
            return Ok(bytes.as_slice());
        }

        let normalized = self.normalize_bytes()?;
        let _ = self.normalized_bytes.set(normalized);
        self.normalized_bytes
            .get()
            .map(Vec::as_slice)
            .ok_or_else(|| {
                crate::Error::Reader("failed to cache normalized JSON bytes".to_string())
            })
    }

    fn infer_schema(&self) -> Result<Arc<Schema>> {
        match self.mode()? {
            JsonMode::Array => {
                let bytes = self.normalized_bytes_cached()?;
                let mut reader = BufReader::new(Cursor::new(bytes));
                let (schema, _inferred_rows) =
                    infer_json_schema_from_seekable(&mut reader, Some(1e6 as usize))?;
                Ok(Arc::new(schema))
            }
            JsonMode::Ndjson => {
                let (schema, _inferred_rows) =
                    infer_json_schema(self.ndjson_reader()?, Some(1e6 as usize))?;
                Ok(Arc::new(schema))
            }
        }
    }

    fn schema_cached(&self) -> Result<Arc<Schema>> {
        if let Some(schema) = self.schema.get() {
            return Ok(schema.clone());
        }
        let schema = self.infer_schema()?;
        let _ = self.schema.set(schema.clone());
        Ok(schema)
    }

    fn cache_schema(&self, schema: Arc<Schema>) -> Arc<Schema> {
        let _ = self.schema.set(schema.clone());
        self.schema.get().cloned().unwrap_or(schema)
    }

    fn open_ndjson_streaming_reader(
        &self,
    ) -> Result<arrow_json::Reader<ReplayBufRead<Box<dyn BufRead + Send>>>> {
        let source = self.ndjson_reader()?;
        let mut recording = RecordingBufRead::new(source);
        let (schema, _inferred_rows) = infer_json_schema(&mut recording, Some(1e6 as usize))?;
        let schema = self.cache_schema(Arc::new(schema));
        let (source, prefix) = recording.into_parts();

        ReaderBuilder::new(schema)
            .with_batch_size(self.batch_size.max(1))
            .build(ReplayBufRead::new(prefix, source))
            .map_err(Into::into)
    }
}

impl OnceReader for JsonReader {
    fn read(&mut self) -> Result<DataFrame> {
        dataframe_from_stream(self)
    }

    fn schema(&self) -> Result<Arc<Schema>> {
        self.schema_cached()
    }

    fn read_stream(
        &mut self,
        sink: &mut dyn FnMut(arrow::array::RecordBatch) -> Result<()>,
    ) -> Result<()> {
        self.ensure_batch_size()?;
        let emit_single_batch = emits_single_batch(self.batch_size);
        let mode = self.mode()?;
        let mut saw_any = false;
        let mut batch_ix = 0usize;
        match mode {
            JsonMode::Array => {
                let schema = self.schema_cached()?;
                let bytes = self.normalized_bytes_cached()?;
                let mut reader = ReaderBuilder::new(schema)
                    .with_batch_size(self.batch_size.max(1))
                    .build(BufReader::new(Cursor::new(bytes)))?;
                stream_record_batches(
                    "json",
                    &mut batch_ix,
                    &mut saw_any,
                    emit_single_batch,
                    &mut reader,
                    crate::Error::from,
                    sink,
                )?;
            }
            JsonMode::Ndjson => {
                if let Some(schema) = self.schema.get() {
                    let mut reader = ReaderBuilder::new(schema.clone())
                        .with_batch_size(self.batch_size.max(1))
                        .build(self.ndjson_reader()?)?;
                    stream_record_batches(
                        "json",
                        &mut batch_ix,
                        &mut saw_any,
                        emit_single_batch,
                        &mut reader,
                        crate::Error::from,
                        sink,
                    )?;
                } else {
                    let mut reader = self.open_ndjson_streaming_reader()?;
                    stream_record_batches(
                        "json",
                        &mut batch_ix,
                        &mut saw_any,
                        emit_single_batch,
                        &mut reader,
                        crate::Error::from,
                        sink,
                    )?;
                }
            }
        }

        finish_stream(saw_any)
    }
}

fn first_non_whitespace_in_reader(reader: &mut dyn BufRead) -> Result<Option<u8>> {
    let mut saw_bom_prefix = 0usize;
    loop {
        let buf = reader.fill_buf()?;
        if buf.is_empty() {
            return Ok(None);
        }

        let mut ix = 0usize;
        if saw_bom_prefix < UTF8_BOM.len() {
            while saw_bom_prefix < UTF8_BOM.len()
                && ix < buf.len()
                && buf[ix] == UTF8_BOM[saw_bom_prefix]
            {
                saw_bom_prefix += 1;
                ix += 1;
            }

            if saw_bom_prefix < UTF8_BOM.len() && saw_bom_prefix > 0 {
                saw_bom_prefix = UTF8_BOM.len();
            }
        }

        while ix < buf.len() {
            let byte = buf[ix];
            if !byte.is_ascii_whitespace() {
                return Ok(Some(byte));
            }
            ix += 1;
        }

        reader.consume(ix);
    }
}
