// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

use datafusion::common::config::JsonOptions;
use datafusion::common::parsers::CompressionTypeVariant;
use datafusion::dataframe::{DataFrame, DataFrameWriteOptions};
use datafusion::datasource::file_format::file_compression_type::FileCompressionType;
use datafusion::error::DataFusionError;
use datafusion::prelude::JsonReadOptions;
use datafusion::prelude::SessionContext;
use jni::objects::{JByteArray, JClass, JString};
use jni::sys::jlong;
use jni::JNIEnv;
use prost::Message;

use crate::proto_gen::{
    FileCompressionType as ProtoFileCompressionType, JsonWriteOptionsProto, NdJsonReadOptionsProto,
};
use crate::runtime;
use crate::schema::decode_optional_schema;
use datafusion_jni_common::errors::{try_unwrap_or_throw, JniResult};

fn with_json_options(
    env: &mut JNIEnv,
    options_bytes: JByteArray,
    schema_ipc_bytes: JByteArray,
    f: impl FnOnce(JsonReadOptions) -> JniResult,
) -> JniResult {
    let bytes: Vec = env.convert_byte_array(&options_bytes)?;
    let p = NdJsonReadOptionsProto::decode(bytes.as_slice())?;

    let schema = decode_optional_schema(env, schema_ipc_bytes)?;

    let compression = match p.file_compression_type() {
        ProtoFileCompressionType::Unspecified => {
            return Err("NdJsonReadOptionsProto.file_compression_type is UNSPECIFIED".into());
        }
        ProtoFileCompressionType::Uncompressed => FileCompressionType::UNCOMPRESSED,
        ProtoFileCompressionType::Gzip => FileCompressionType::GZIP,
        ProtoFileCompressionType::Bzip2 => FileCompressionType::BZIP2,
        ProtoFileCompressionType::Xz => FileCompressionType::XZ,
        ProtoFileCompressionType::Zstd => FileCompressionType::ZSTD,
    };

    let file_ext = p.file_extension;
    let mut opts = JsonReadOptions::default()
        .file_extension(&file_ext)
        .file_compression_type(compression);

    if let Some(n) = p.schema_infer_max_records {
        opts = opts.schema_infer_max_records(n as usize);
    }
    if let Some(ref s) = schema {
        opts = opts.schema(s);
    }

    f(opts)
}

#[no_mangle]
pub extern "system" fn Java_org_apache_datafusion_SessionContext_registerJsonWithOptions,
    _class: JClass,
    path: JString,
    schema_ipc_bytes: JByteArray(
    mut env: JNIEnv,
    handle: jlong,
    path: JString,
    schema_ipc_bytes: JByteArray(
    mut env: JNIEnv,
    handle: jlong,
    path: JString,
) {
    try_unwrap_or_throw(&mut env, (), |env| -> JniResult {
        if handle == 0 {
            return Err("DataFrame handle is null".into());
        }
        let df = unsafe { &*(handle as *const DataFrame) }.clone();
        let path: String = env.get_string(&path)?.into();
        let bytes: Vec = env.convert_byte_array(&options_bytes)?;
        let p = JsonWriteOptionsProto::decode(bytes.as_slice())?;

        // Decode the file_compression_type field eagerly so an unknown wire
        // value surfaces as a clear error rather than a silent default.
        let compression = if p.file_compression_type.is_some() {
            Some(proto_compression_to_variant(p.file_compression_type())?)
        } else {
            None
        };

        // When the caller left `singleFileOutput` unset, force directory output (`false`)
        // rather than leaving DataFusion in `Automatic` mode. Automatic mode treats paths
        // with an extension (e.g. `out.json`) as single-file targets, which would silently
        // contradict the documented "directory unless overridden" default and surprise any
        // caller that hands writeJson a `.json` path.
        let mut write_opts = DataFrameWriteOptions::new()
            .with_single_file_output(p.single_file_output.unwrap_or(false));
        if !p.partition_cols.is_empty() {
            write_opts = write_opts.with_partition_by(p.partition_cols.clone());
        }

        // Build JsonOptions only when a writer-side knob is set, so the
        // DataFusion default is preserved when the caller passes
        // `new JsonWriteOptions()`. JsonOptions has no fluent setters --
        // the fields are public, so use struct-update syntax (same
        // idiom we use for ArrowReadOptions / AvroReadOptions).
        let writer_opts: Option = compression.map(|c| JsonOptions {
            compression: c,
            ..JsonOptions::default()
        });

        runtime().block_on(df.write_json(&path, write_opts, writer_opts))?;
        Ok(())
    })
}