// 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.

//! Apply [`ObjectStoreRegistration`] entries from the session-options proto to a
//! [`SessionContext`]'s [`RuntimeEnv`].
//!
//! Each backend (`s3` / `gcs` / `http`) is gated behind its own Cargo feature;
//! the corresponding arm of the proto `oneof` returns a clear error if the
//! feature was not enabled at build time. The default build of this crate
//! enables all three so Java callers don't have to think about features.

use std::sync::Arc;

use datafusion::prelude::SessionContext;
use url::Url;

use crate::proto_gen::object_store_registration::Backend;
use crate::proto_gen::ObjectStoreRegistration;
use datafusion_jni_common::errors::JniResult;

#[cfg(feature = "object-store-gcp")]
use crate::proto_gen::GcsOptions;
#[cfg(feature = "object-store-http")]
use crate::proto_gen::HttpOptions;
#[cfg(feature = "object-store-aws")]
use crate::proto_gen::S3Options;

/// Apply every registration in `regs`, in order, to the context's `RuntimeEnv`.
/// If two registrations resolve to the same URL, the later one wins (matching
/// upstream `RuntimeEnv::register_object_store`).
pub(crate) fn apply_registrations(
    ctx: &SessionContext,
    regs: &[ObjectStoreRegistration],
) -> JniResult<()> {
    for reg in regs {
        let backend = reg
            .backend
            .as_ref()
            .ok_or("ObjectStoreRegistration.backend is required")?;
        let (url, store) = build_store(reg.url.as_deref(), backend)?;
        ctx.runtime_env().register_object_store(&url, store);
    }
    Ok(())
}

#[allow(unused_variables)] // `url_override` is unused on builds with no features
fn build_store(
    url_override: Option<&str>,
    backend: &Backend,
) -> JniResult<(Url, Arc<dyn object_store::ObjectStore>)> {
    match backend {
        #[cfg(feature = "object-store-aws")]
        Backend::S3(opts) => build_s3(url_override, opts),
        #[cfg(not(feature = "object-store-aws"))]
        Backend::S3(_) => Err(
            "object-store-aws Cargo feature is not enabled in this build of datafusion-jni".into(),
        ),

        #[cfg(feature = "object-store-gcp")]
        Backend::Gcs(opts) => build_gcs(url_override, opts),
        #[cfg(not(feature = "object-store-gcp"))]
        Backend::Gcs(_) => Err(
            "object-store-gcp Cargo feature is not enabled in this build of datafusion-jni".into(),
        ),

        #[cfg(feature = "object-store-http")]
        Backend::Http(opts) => build_http(url_override, opts),
        #[cfg(not(feature = "object-store-http"))]
        Backend::Http(_) => Err(
            "object-store-http Cargo feature is not enabled in this build of datafusion-jni".into(),
        ),
    }
}

#[cfg(feature = "object-store-aws")]
fn build_s3(
    url_override: Option<&str>,
    opts: &S3Options,
) -> JniResult<(Url, Arc<dyn object_store::ObjectStore>)> {
    use object_store::aws::{AmazonS3Builder, AmazonS3ConfigKey};

    if opts.bucket.is_empty() {
        return Err("S3Options.bucket is required".into());
    }

    // Seed from the standard AWS env (AWS_ACCESS_KEY_ID, AWS_DEFAULT_REGION,
    // AWS_WEB_IDENTITY_TOKEN_FILE, ECS task creds, etc.) before applying
    // explicit Java overrides. Without this, callers running on EC2/ECS/EKS
    // who omit Java-side credentials and expect the SDK default-credential
    // chain would get an empty builder and fail at request time.
    //
    // Walk env explicitly (instead of AmazonS3Builder::from_env()) so we can
    // skip endpoint keys when the Java caller set `endpoint(...)`.
    // AmazonS3Builder::build() picks `s3_endpoint` over `endpoint`
    // (object_store 0.13 aws/builder.rs:1209), so a stray
    // AWS_ENDPOINT_URL_S3 in the JVM env would otherwise silently override
    // the explicit Java endpoint -- breaking MinIO/R2/etc. registrations
    // run out of a process with global AWS env vars.
    let java_has_endpoint = opts.endpoint.is_some();
    let mut b = AmazonS3Builder::default();
    for (os_key, os_value) in std::env::vars_os() {
        let (Some(key), Some(value)) = (os_key.to_str(), os_value.to_str()) else {
            continue;
        };
        if !key.starts_with("AWS_") {
            continue;
        }
        if java_has_endpoint
            && (key == "AWS_ENDPOINT" || key == "AWS_ENDPOINT_URL" || key == "AWS_ENDPOINT_URL_S3")
        {
            continue;
        }
        // Skip AWS_BUCKET / AWS_BUCKET_NAME: the Java caller's bucket is the
        // identity of this registration (it's also the host portion of the
        // registry-key URL we return). An env-supplied bucket would overwrite
        // it here while the URL still says `s3://<opts.bucket>`, so a query
        // for the Java bucket would be served by the env bucket. Apply the
        // Java bucket explicitly *after* the loop so it can't be clobbered.
        if key == "AWS_BUCKET" || key == "AWS_BUCKET_NAME" {
            continue;
        }
        if let Ok(config_key) = key.to_ascii_lowercase().parse() {
            b = b.with_config(config_key, value);
        }
    }
    b = b.with_bucket_name(&opts.bucket);
    if let Some(ref v) = opts.region {
        b = b.with_region(v);
    }
    if let Some(ref v) = opts.endpoint {
        b = b.with_endpoint(v);
    }
    if let Some(ref v) = opts.access_key_id {
        b = b.with_access_key_id(v);
    }
    if let Some(ref v) = opts.secret_access_key {
        b = b.with_secret_access_key(v);
    }
    if let Some(ref v) = opts.session_token {
        b = b.with_token(v);
    }
    if let Some(v) = opts.allow_http {
        b = b.with_allow_http(v);
    }
    if let Some(v) = opts.skip_signature {
        b = b.with_skip_signature(v);
    }
    if let Some(v) = opts.imdsv1_fallback {
        // AmazonS3Builder::with_imdsv1_fallback() only sets to `true` (no
        // boolean arg). To honor explicit `imdsv1Fallback(false)` and
        // override an env-seeded `AWS_IMDSV1_FALLBACK=true`, write through
        // with_config(...) which always sets the field, both directions.
        b = b.with_config(AmazonS3ConfigKey::ImdsV1Fallback, v.to_string());
    }

    let store = b.build()?;
    let url = parse_url(url_override, format!("s3://{}", opts.bucket))?;
    Ok((url, Arc::new(store)))
}

#[cfg(feature = "object-store-gcp")]
fn build_gcs(
    url_override: Option<&str>,
    opts: &GcsOptions,
) -> JniResult<(Url, Arc<dyn object_store::ObjectStore>)> {
    use object_store::gcp::GoogleCloudStorageBuilder;

    if opts.bucket.is_empty() {
        return Err("GcsOptions.bucket is required".into());
    }
    if opts.service_account_key.is_some() && opts.service_account_path.is_some() {
        return Err(
            "GcsOptions: service_account_key and service_account_path are mutually exclusive"
                .into(),
        );
    }

    // Seed from the standard GCS env (GOOGLE_BUCKET, GOOGLE_SERVICE_ACCOUNT,
    // etc.) before applying explicit Java overrides.
    //
    // Walk env explicitly (instead of GoogleCloudStorageBuilder::from_env())
    // so we can skip credential keys that conflict with Java-supplied
    // credentials. GoogleCloudStorageBuilder::build() rejects the combo of
    // service_account_path AND service_account_key with
    // ServiceAccountPathAndKeyProvided (object_store 0.13 gcp/builder.rs:525)
    // -- so if env has GOOGLE_SERVICE_ACCOUNT_KEY and Java sets
    // serviceAccountPath(...), the build would fail. Skip env credential
    // keys when Java provides any of the three credential fields.
    let java_has_credential = opts.service_account_key.is_some()
        || opts.service_account_path.is_some()
        || opts.application_credentials.is_some();
    let mut b = GoogleCloudStorageBuilder::default();
    for (os_key, os_value) in std::env::vars_os() {
        let (Some(key), Some(value)) = (os_key.to_str(), os_value.to_str()) else {
            continue;
        };
        if !key.starts_with("GOOGLE_") {
            continue;
        }
        if java_has_credential
            && (key == "GOOGLE_SERVICE_ACCOUNT"
                || key == "GOOGLE_SERVICE_ACCOUNT_PATH"
                || key == "GOOGLE_SERVICE_ACCOUNT_KEY"
                || key == "GOOGLE_APPLICATION_CREDENTIALS")
        {
            continue;
        }
        // Skip GOOGLE_BUCKET / GOOGLE_BUCKET_NAME: the Java caller's bucket is
        // the identity of this registration (matches the registry-key URL).
        // Apply the Java bucket explicitly *after* the loop so an env-supplied
        // bucket can't silently route the registered URL to a different
        // backend bucket.
        if key == "GOOGLE_BUCKET" || key == "GOOGLE_BUCKET_NAME" {
            continue;
        }
        if let Ok(config_key) = key.to_ascii_lowercase().parse() {
            b = b.with_config(config_key, value);
        }
    }
    b = b.with_bucket_name(&opts.bucket);
    if let Some(ref v) = opts.service_account_key {
        b = b.with_service_account_key(v);
    }
    if let Some(ref v) = opts.service_account_path {
        b = b.with_service_account_path(v);
    }
    if let Some(ref v) = opts.application_credentials {
        b = b.with_application_credentials(v);
    }

    let store = b.build()?;
    let url = parse_url(url_override, format!("gs://{}", opts.bucket))?;
    Ok((url, Arc::new(store)))
}

#[cfg(feature = "object-store-http")]
fn build_http(
    url_override: Option<&str>,
    opts: &HttpOptions,
) -> JniResult<(Url, Arc<dyn object_store::ObjectStore>)> {
    use object_store::http::HttpBuilder;

    let listing = url_override.ok_or(
        "HttpOptions: ObjectStoreRegistration.url is required for the HTTP backend (no scheme-default)",
    )?;
    let listing_url =
        Url::parse(listing).map_err(|e| format!("invalid HTTP URL {listing:?}: {e}"))?;

    // DataFusion's DefaultObjectStoreRegistry::get_url_key strips paths from
    // the registry key (it keeps only scheme + host:port), so a registration
    // base like `https://example.com/data/` ends up under key
    // `https://example.com`. If we hand HttpBuilder the same pathful URL,
    // HttpStore prepends `/data/` to every object path on top of the path
    // DataFusion already supplied -- a request for
    // `https://example.com/data/file.parquet` becomes
    // `https://example.com/data/data/file.parquet`. Strip the URL to
    // scheme + authority for both the store base and the registry key so
    // they match upstream's lookup semantics; the SQL-side URL still carries
    // the full path, and HttpStore appends it once.
    if !listing_url.path().is_empty() && listing_url.path() != "/" {
        return Err(format!(
            "HttpOptions: listing URL must be a host root (no path component); \
             got {listing:?}. DataFusion's object-store registry keys ignore the \
             path portion of a URL, so a pathful base would cause every object \
             read to double-prepend the path. Register one HTTP store per \
             scheme+host and let SQL paths carry the rest of the URL."
        )
        .into());
    }

    let mut b = HttpBuilder::new().with_url(listing);
    if let Some(v) = opts.allow_http {
        // allow_http lives on ClientOptions in object_store 0.13, not directly
        // on HttpBuilder; route through with_client_options.
        b = b.with_client_options(object_store::ClientOptions::new().with_allow_http(v));
    }

    let store = b.build()?;
    Ok((listing_url, Arc::new(store)))
}

fn parse_url(url_override: Option<&str>, default: String) -> JniResult<Url> {
    let s = url_override.unwrap_or(&default);
    Url::parse(s).map_err(|e| format!("invalid object store URL {s:?}: {e}").into())
}