// 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 std::collections::HashMap;
use std::ffi::{CStr, CString};
use std::sync::Arc;

use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec;
use datafusion_ffi::table_provider::FFI_TableProvider;
use iceberg::TableIdent;
use iceberg::io::FileIOBuilder;
use iceberg::table::StaticTable;
use iceberg_datafusion::table::IcebergStaticTableProvider;
use iceberg_storage_opendal::OpenDalResolvingStorageFactory;
use pyo3::exceptions::{PyRuntimeError, PyValueError};
use pyo3::prelude::*;
use pyo3::types::{PyAny, PyCapsule};

use crate::runtime::runtime;

// pyo3 0.28's CapsuleName only exposes `unsafe fn as_cstr() -> &CStr`,
// so we accept &CStr to allow direct comparison without UTF-8 validation.
pub(crate) fn validate_pycapsule(capsule: &Bound<PyCapsule>, name: &CStr) -> PyResult<()> {
    let capsule_name = capsule.name()?;
    match capsule_name {
        None => Err(PyValueError::new_err(
            "Expected PyCapsule to have name set.",
        )),
        Some(capsule_name) if unsafe { capsule_name.as_cstr() } != name => {
            Err(PyValueError::new_err("PyCapsule name mismatch"))
        }
        _ => Ok(()),
    }
}

pub(crate) fn ffi_logical_codec_from_pycapsule(
    obj: Bound<PyAny>,
) -> PyResult<FFI_LogicalExtensionCodec> {
    let attr_name = "__datafusion_logical_extension_codec__";
    let capsule = if obj.hasattr(attr_name)? {
        obj.getattr(attr_name)?.call0()?
    } else {
        obj
    };

    let capsule_name = c"datafusion_logical_extension_codec";
    let capsule = capsule.cast::<PyCapsule>()?;
    validate_pycapsule(capsule, capsule_name)?;

    let ptr = capsule.pointer_checked(Some(capsule_name))?;
    let codec = unsafe { &*(ptr.as_ptr() as *const FFI_LogicalExtensionCodec) };
    Ok(codec.clone())
}

#[pyclass(name = "IcebergDataFusionTable")]
pub struct PyIcebergDataFusionTable {
    inner: Arc<IcebergStaticTableProvider>,
}

#[pymethods]
impl PyIcebergDataFusionTable {
    #[new]
    fn new(
        identifier: Vec<String>,
        metadata_location: String,
        file_io_properties: Option<HashMap<String, String>>,
    ) -> PyResult<Self> {
        let runtime = runtime();

        let provider = runtime.block_on(async {
            let table_ident = TableIdent::from_strs(identifier)
                .map_err(|e| PyRuntimeError::new_err(format!("Invalid table identifier: {e}")))?;

            let factory = Arc::new(OpenDalResolvingStorageFactory::new());

            let mut builder = FileIOBuilder::new(factory);

            if let Some(props) = file_io_properties {
                builder = builder.with_props(props);
            }

            let file_io = builder.build();

            let static_table =
                StaticTable::from_metadata_file(&metadata_location, table_ident, file_io)
                    .await
                    .map_err(|e| {
                        PyRuntimeError::new_err(format!("Failed to load static table: {e}"))
                    })?;

            let table = static_table.into_table();

            IcebergStaticTableProvider::try_new_from_table(table)
                .await
                .map_err(|e| {
                    PyRuntimeError::new_err(format!("Failed to create table provider: {e}"))
                })
        })?;

        Ok(Self {
            inner: Arc::new(provider),
        })
    }

    fn __datafusion_table_provider__<'py>(
        &self,
        py: Python<'py>,
        session: Bound<'py, PyAny>,
    ) -> PyResult<Bound<'py, PyCapsule>> {
        let capsule_name = CString::new("datafusion_table_provider").unwrap();

        let logical_codec = ffi_logical_codec_from_pycapsule(session)?;

        let ffi_provider = FFI_TableProvider::new_with_ffi_codec(
            self.inner.clone(),
            false,
            Some(runtime()),
            logical_codec,
        );

        PyCapsule::new(py, ffi_provider, Some(capsule_name))
    }
}

pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
    let this = PyModule::new(py, "datafusion")?;

    this.add_class::<PyIcebergDataFusionTable>()?;

    m.add_submodule(&this)?;
    py.import("sys")?
        .getattr("modules")?
        .set_item("pyiceberg_core.datafusion", this)?;

    Ok(())
}
