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

package org.apache.datafusion;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.channels.Channels;
import java.util.ArrayList;
import java.util.List;

import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.ipc.ArrowStreamWriter;
import org.apache.arrow.vector.ipc.ReadChannel;
import org.apache.arrow.vector.ipc.message.MessageSerializer;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.Schema;

/**
 * A DataFusion session context.
 *
 * <p>Instances are <strong>not thread-safe</strong>. Concurrent calls to any of {@link #sql},
 * {@link #registerParquet}, or {@link #close} from different threads can produce a use-after-free
 * on the native side. Callers must externally synchronize, or confine each context to a single
 * thread.
 */
public final class SessionContext implements AutoCloseable {
  static {
    NativeLibraryLoader.loadLibrary();
  }

  private long nativeHandle;

  public SessionContext() {
    this.nativeHandle = createSessionContext();
    if (this.nativeHandle == 0) {
      throw new RuntimeException("Failed to create native SessionContext");
    }
  }

  SessionContext(byte[] optionsBytes) {
    this.nativeHandle = createSessionContextWithOptions(optionsBytes);
    if (this.nativeHandle == 0) {
      throw new RuntimeException("Failed to create native SessionContext");
    }
  }

  /** Start configuring a {@link SessionContext}. */
  public static SessionContextBuilder builder() {
    return new SessionContextBuilder();
  }

  /**
   * Parse and plan {@code query}, returning a lazy {@link DataFrame}. The query is not executed
   * until {@link DataFrame#collect} is called.
   */
  public DataFrame sql(String query) {
    if (nativeHandle == 0) {
      throw new IllegalStateException("SessionContext is closed");
    }
    long dfHandle = createDataFrame(nativeHandle, query);
    return new DataFrame(dfHandle);
  }

  /**
   * Decode a DataFusion-Proto {@code LogicalPlanNode} and return a lazy {@link DataFrame}. The plan
   * is not executed until {@link DataFrame#collect} is called.
   *
   * <p>The bytes must be a serialized {@code datafusion.LogicalPlanNode} (see {@code
   * org.apache.datafusion.protobuf.LogicalPlanNode}).
   *
   * @throws RuntimeException if the bytes are not a valid {@code LogicalPlanNode} or if logical
   *     planning fails.
   */
  public DataFrame fromProto(byte[] planBytes) {
    if (nativeHandle == 0) {
      throw new IllegalStateException("SessionContext is closed");
    }
    long dfHandle = createDataFrameFromProto(nativeHandle, planBytes);
    return new DataFrame(dfHandle);
  }

  /**
   * Snapshot the session's memory pool: bytes currently held and the peak observed since this
   * session was created. Thread-safe; can be polled while queries run.
   *
   * <p>For multi-tenant attribution, place each tenant in its own {@link SessionContext}. Within a
   * single session, the snapshot is the sum across all in-flight queries -- there is no
   * per-DataFrame breakdown today.
   *
   * <p>The session's {@code MemoryPool} is wrapped transparently with a tracking adapter at
   * construction time; the wrapper layers on top of whatever pool {@link
   * SessionContextBuilder#memoryLimit(long, double)} produced (or DataFusion's default unbounded
   * pool) and does not change pool semantics (limits, eviction, spilling).
   *
   * <p><b>What this counts:</b> bytes reserved against the {@code MemoryPool} -- operator state for
   * sorts, hash joins, aggregates, repartition buffers, and anything else that uses DataFusion's
   * {@code MemoryReservation} machinery during execution.
   *
   * <p><b>What this does <i>not</i> count:</b> memory held outside the pool, including record-batch
   * buffers materialised by {@link DataFrame#cache()} (stored in an in-memory {@code MemTable} as
   * plain {@code Vec<RecordBatch>} with no reservation), record-batch buffers that have crossed the
   * FFI boundary into Arrow's Java allocator, and JVM-side allocations. Operator-level reservations
   * are released as the plan unwinds, so a query that runs to completion typically returns {@code
   * currentBytes} to ~0 even if the result set is large.
   *
   * @throws IllegalStateException if this context is closed.
   * @throws RuntimeException if the native side has not registered a tracker for this handle
   *     (should not happen in practice -- tracker registration is done by the constructor).
   */
  public MemoryUsage memoryUsage() {
    if (nativeHandle == 0) {
      throw new IllegalStateException("SessionContext is closed");
    }
    long[] values = memoryUsageNative(nativeHandle);
    return new MemoryUsage(values[0], values[1]);
  }

  /**
   * Snapshot operational counters from the underlying Tokio runtime: worker count, busy time, queue
   * depth, etc. Thread-safe; can be polled while queries run.
   *
   * <p>The runtime is process-wide rather than per-session because the JNI library drives a single
   * shared multi-threaded Tokio runtime. The {@link SessionContext} handle is checked only to
   * ensure the caller still has a live session; the values returned are not session-specific.
   *
   * <p>Requires the {@code runtime-metrics} Cargo feature on the native crate (off by default).
   * Rebuild with:
   *
   * <pre>{@code
   * RUSTFLAGS="--cfg tokio_unstable" cargo build --features runtime-metrics
   * }</pre>
   *
   * <p>If invoked against a native binary built without the feature, this method throws {@link
   * RuntimeException} with a message pointing at the rebuild command.
   *
   * @throws IllegalStateException if this context is closed.
   * @throws RuntimeException if the native crate was built without the {@code runtime-metrics}
   *     feature.
   */
  public RuntimeStats runtimeStats() {
    if (nativeHandle == 0) {
      throw new IllegalStateException("SessionContext is closed");
    }
    long[] s = runtimeStatsNative(nativeHandle);
    return new RuntimeStats(
        (int) s[0], s[1], s[2], s[3], s[4], s[5], s[6], s[7], s[8], s[9], s[10]);
  }

  /**
   * Return the Arrow {@link Schema} of a registered table. Transferred via Arrow IPC; no {@link
   * org.apache.arrow.memory.BufferAllocator} is required because a schema carries no buffer data.
   *
   * @throws RuntimeException if {@code tableName} is not registered in this context.
   */
  public Schema tableSchema(String tableName) {
    if (nativeHandle == 0) {
      throw new IllegalStateException("SessionContext is closed");
    }
    byte[] ipcBytes = tableSchemaIpc(nativeHandle, tableName);
    try {
      return MessageSerializer.deserializeSchema(
          new ReadChannel(Channels.newChannel(new ByteArrayInputStream(ipcBytes))));
    } catch (IOException e) {
      throw new RuntimeException("Failed to deserialize IPC schema", e);
    }
  }

  /**
   * Read the current value of a {@code datafusion.*} config key. The key must be one DataFusion
   * recognises (see {@link SessionContextBuilder#setOption(String, String)} for examples and the
   * upstream configuration reference for the full list).
   *
   * <p>{@code datafusion.runtime.*} keys (memory limit, temp directory, cache sizes, etc) are not
   * yet supported by this getter and will throw. Use the typed {@link
   * SessionContextBuilder#memoryLimit(long, double)} and {@link
   * SessionContextBuilder#tempDirectory(String)} setters at construction time instead. Round-trip
   * support for the runtime subtree is tracked as a follow-up.
   *
   * @return the current value as a string, or {@code null} if the key is recognised but has no
   *     value set and no default.
   * @throws IllegalArgumentException if {@code key} is {@code null}.
   * @throws RuntimeException if the key is not recognised by DataFusion or is in the {@code
   *     datafusion.runtime.*} subtree.
   * @throws IllegalStateException if this context is closed.
   */
  public String getOption(String key) {
    if (nativeHandle == 0) {
      throw new IllegalStateException("SessionContext is closed");
    }
    if (key == null) {
      throw new IllegalArgumentException("getOption key must be non-null");
    }
    return getOptionNative(nativeHandle, key);
  }

  public void registerCsv(String name, String path) {
    registerCsv(name, path, new CsvReadOptions());
  }

  /**
   * Register a CSV file (or directory of CSV files) as a table with the supplied {@link
   * CsvReadOptions}.
   *
   * @throws RuntimeException if registration fails (path not found, schema inference error, etc.).
   */
  public void registerCsv(String name, String path, CsvReadOptions options) {
    if (nativeHandle == 0) {
      throw new IllegalStateException("SessionContext is closed");
    }
    registerCsvWithOptions(
        nativeHandle,
        name,
        path,
        options.toBytes(),
        options.schema() != null ? serializeSchemaIpc(options.schema()) : null);
  }

  /** Read a CSV file as a {@link DataFrame} without registering it. */
  public DataFrame readCsv(String path) {
    return readCsv(path, new CsvReadOptions());
  }

  /**
   * Read a CSV file as a {@link DataFrame} with the supplied {@link CsvReadOptions}.
   *
   * @throws RuntimeException if the read fails.
   */
  public DataFrame readCsv(String path, CsvReadOptions options) {
    if (nativeHandle == 0) {
      throw new IllegalStateException("SessionContext is closed");
    }
    long dfHandle =
        readCsvWithOptions(
            nativeHandle,
            path,
            options.toBytes(),
            options.schema() != null ? serializeSchemaIpc(options.schema()) : null);
    return new DataFrame(dfHandle);
  }

  public void registerJson(String name, String path) {
    registerJson(name, path, new NdJsonReadOptions());
  }

  /**
   * Register a newline-delimited JSON file (or directory of NDJSON files) as a table with the
   * supplied {@link NdJsonReadOptions}.
   *
   * @throws RuntimeException if registration fails (path not found, schema inference error, etc.).
   */
  public void registerJson(String name, String path, NdJsonReadOptions options) {
    if (nativeHandle == 0) {
      throw new IllegalStateException("SessionContext is closed");
    }
    if (name == null) {
      throw new IllegalArgumentException("registerJson name must be non-null");
    }
    if (path == null) {
      throw new IllegalArgumentException("registerJson path must be non-null");
    }
    if (options == null) {
      throw new IllegalArgumentException("registerJson options must be non-null");
    }
    registerJsonWithOptions(
        nativeHandle,
        name,
        path,
        options.toBytes(),
        options.schema() != null ? serializeSchemaIpc(options.schema()) : null);
  }

  /** Read a newline-delimited JSON file as a {@link DataFrame} without registering it. */
  public DataFrame readJson(String path) {
    return readJson(path, new NdJsonReadOptions());
  }

  /**
   * Read a newline-delimited JSON file as a {@link DataFrame} with the supplied {@link
   * NdJsonReadOptions}.
   *
   * @throws RuntimeException if the read fails.
   */
  public DataFrame readJson(String path, NdJsonReadOptions options) {
    if (nativeHandle == 0) {
      throw new IllegalStateException("SessionContext is closed");
    }
    if (path == null) {
      throw new IllegalArgumentException("readJson path must be non-null");
    }
    if (options == null) {
      throw new IllegalArgumentException("readJson options must be non-null");
    }
    long dfHandle =
        readJsonWithOptions(
            nativeHandle,
            path,
            options.toBytes(),
            options.schema() != null ? serializeSchemaIpc(options.schema()) : null);
    return new DataFrame(dfHandle);
  }

  public void registerParquet(String name, String path) {
    registerParquet(name, path, new ParquetReadOptions());
  }

  /**
   * Register a parquet file as a table with the supplied {@link ParquetReadOptions}.
   *
   * @throws RuntimeException if registration fails (path not found, schema mismatch, etc.).
   */
  public void registerParquet(String name, String path, ParquetReadOptions options) {
    if (nativeHandle == 0) {
      throw new IllegalStateException("SessionContext is closed");
    }
    registerParquetWithOptions(
        nativeHandle,
        name,
        path,
        options.toBytes(),
        options.schema() != null ? serializeSchemaIpc(options.schema()) : null);
  }

  /** Read a parquet file as a {@link DataFrame} without registering it. */
  public DataFrame readParquet(String path) {
    return readParquet(path, new ParquetReadOptions());
  }

  /**
   * Read a parquet file as a {@link DataFrame} with the supplied {@link ParquetReadOptions}.
   *
   * @throws RuntimeException if the read fails.
   */
  public DataFrame readParquet(String path, ParquetReadOptions options) {
    if (nativeHandle == 0) {
      throw new IllegalStateException("SessionContext is closed");
    }
    long dfHandle =
        readParquetWithOptions(
            nativeHandle,
            path,
            options.toBytes(),
            options.schema() != null ? serializeSchemaIpc(options.schema()) : null);
    return new DataFrame(dfHandle);
  }

  /** Register an Arrow IPC file (or directory of Arrow IPC files) as a table. */
  public void registerArrow(String name, String path) {
    registerArrow(name, path, new ArrowReadOptions());
  }

  /**
   * Register an Arrow IPC file (or directory of Arrow IPC files) as a table with the supplied
   * {@link ArrowReadOptions}.
   *
   * @throws IllegalArgumentException if any of {@code name}, {@code path}, or {@code options} is
   *     {@code null}.
   * @throws RuntimeException if registration fails (path not found, schema mismatch, etc.).
   */
  public void registerArrow(String name, String path, ArrowReadOptions options) {
    if (nativeHandle == 0) {
      throw new IllegalStateException("SessionContext is closed");
    }
    if (name == null) {
      throw new IllegalArgumentException("registerArrow name must be non-null");
    }
    if (path == null) {
      throw new IllegalArgumentException("registerArrow path must be non-null");
    }
    if (options == null) {
      throw new IllegalArgumentException("registerArrow options must be non-null");
    }
    registerArrowWithOptions(
        nativeHandle,
        name,
        path,
        options.toBytes(),
        options.schema() != null ? serializeSchemaIpc(options.schema()) : null);
  }

  /** Read an Arrow IPC file as a {@link DataFrame} without registering it. */
  public DataFrame readArrow(String path) {
    return readArrow(path, new ArrowReadOptions());
  }

  /**
   * Read an Arrow IPC file as a {@link DataFrame} with the supplied {@link ArrowReadOptions}.
   *
   * @throws IllegalArgumentException if {@code path} or {@code options} is {@code null}.
   * @throws RuntimeException if the read fails.
   */
  public DataFrame readArrow(String path, ArrowReadOptions options) {
    if (nativeHandle == 0) {
      throw new IllegalStateException("SessionContext is closed");
    }
    if (path == null) {
      throw new IllegalArgumentException("readArrow path must be non-null");
    }
    if (options == null) {
      throw new IllegalArgumentException("readArrow options must be non-null");
    }
    long dfHandle =
        readArrowWithOptions(
            nativeHandle,
            path,
            options.toBytes(),
            options.schema() != null ? serializeSchemaIpc(options.schema()) : null);
    return new DataFrame(dfHandle);
  }

  /** Register an Avro file (or directory of Avro files) as a table. */
  public void registerAvro(String name, String path) {
    registerAvro(name, path, new AvroReadOptions());
  }

  /**
   * Register an Avro file (or directory of Avro files) as a table with the supplied {@link
   * AvroReadOptions}.
   *
   * @throws IllegalArgumentException if any of {@code name}, {@code path}, or {@code options} is
   *     {@code null}.
   * @throws RuntimeException if registration fails (path not found, schema mismatch, etc.).
   */
  public void registerAvro(String name, String path, AvroReadOptions options) {
    if (nativeHandle == 0) {
      throw new IllegalStateException("SessionContext is closed");
    }
    if (name == null) {
      throw new IllegalArgumentException("registerAvro name must be non-null");
    }
    if (path == null) {
      throw new IllegalArgumentException("registerAvro path must be non-null");
    }
    if (options == null) {
      throw new IllegalArgumentException("registerAvro options must be non-null");
    }
    registerAvroWithOptions(
        nativeHandle,
        name,
        path,
        options.toBytes(),
        options.schema() != null ? serializeSchemaIpc(options.schema()) : null);
  }

  /** Read an Avro file as a {@link DataFrame} without registering it. */
  public DataFrame readAvro(String path) {
    return readAvro(path, new AvroReadOptions());
  }

  /**
   * Read an Avro file as a {@link DataFrame} with the supplied {@link AvroReadOptions}.
   *
   * @throws IllegalArgumentException if {@code path} or {@code options} is {@code null}.
   * @throws RuntimeException if the read fails.
   */
  public DataFrame readAvro(String path, AvroReadOptions options) {
    if (nativeHandle == 0) {
      throw new IllegalStateException("SessionContext is closed");
    }
    if (path == null) {
      throw new IllegalArgumentException("readAvro path must be non-null");
    }
    if (options == null) {
      throw new IllegalArgumentException("readAvro options must be non-null");
    }
    long dfHandle =
        readAvroWithOptions(
            nativeHandle,
            path,
            options.toBytes(),
            options.schema() != null ? serializeSchemaIpc(options.schema()) : null);
    return new DataFrame(dfHandle);
  }

  /**
   * Register a Java-implemented scalar UDF. After registration, the function can be invoked by SQL
   * via the UDF's name or referenced in DataFusion plans deserialised with {@link #fromProto}.
   *
   * <p>The UDF is registered with an exact signature: the runtime will reject calls whose argument
   * types do not match the declared {@link ScalarFunction#argFields()} exactly.
   *
   * @throws RuntimeException if registration fails (e.g., name already registered with an
   *     incompatible signature, schema serialisation failure).
   */
  public void registerUdf(ScalarUdf udf) {
    if (nativeHandle == 0) {
      throw new IllegalStateException("SessionContext is closed");
    }
    java.util.Objects.requireNonNull(udf, "udf");
    ScalarFunction impl = udf.impl();
    String name = udf.name();
    Volatility volatility = udf.volatility();
    List<Field> fields = new ArrayList<>(udf.argFields().size() + 1);
    fields.add(udf.returnField());
    fields.addAll(udf.argFields());
    Schema signatureSchema = new Schema(fields);
    byte[] signatureBytes = serializeSchemaIpc(signatureSchema);
    registerScalarUdf(nativeHandle, name, signatureBytes, volatility.code(), impl);
  }

  /**
   * Register a Java-implemented {@link TableProvider} under {@code name}. SQL queries that
   * reference {@code name} call back into {@code provider} to fetch batches.
   *
   * <p>{@link TableProvider#schema()} is called once here, on the calling thread, and cached on the
   * native side. {@link TableProvider#scan(org.apache.arrow.memory.BufferAllocator)} is called once
   * per query that touches the table, on a Tokio worker thread; it must return a fresh, independent
   * {@link org.apache.arrow.vector.ipc.ArrowReader} on every call, with its buffers allocated from
   * the {@link org.apache.arrow.memory.BufferAllocator} the framework supplies.
   *
   * <p>This is the Java counterpart to DataFusion's Rust {@code SessionContext::register_table}.
   *
   * @throws IllegalArgumentException if {@code name} or {@code provider} is {@code null}.
   * @throws IllegalStateException if {@code provider.schema()} returns {@code null}, or this
   *     context is closed.
   * @throws RuntimeException if native registration fails.
   */
  public void registerTable(String name, TableProvider provider) {
    if (nativeHandle == 0) {
      throw new IllegalStateException("SessionContext is closed");
    }
    if (name == null) {
      throw new IllegalArgumentException("registerTable name must be non-null");
    }
    if (provider == null) {
      throw new IllegalArgumentException("registerTable provider must be non-null");
    }
    Schema schema = provider.schema();
    if (schema == null) {
      throw new IllegalStateException("TableProvider.schema returned null");
    }
    byte[] schemaIpc = serializeSchemaIpc(schema);
    registerTableNative(nativeHandle, name, schemaIpc, provider);
  }

  private static byte[] serializeSchemaIpc(Schema schema) {
    ByteArrayOutputStream baos = new ByteArrayOutputStream();
    try (BufferAllocator allocator = new RootAllocator();
        VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator);
        ArrowStreamWriter writer = new ArrowStreamWriter(root, null, Channels.newChannel(baos))) {
      writer.start();
      writer.end();
    } catch (IOException e) {
      throw new RuntimeException("Failed to serialize Arrow schema for JNI", e);
    }
    return baos.toByteArray();
  }

  @Override
  public void close() {
    if (nativeHandle != 0) {
      closeSessionContext(nativeHandle);
      nativeHandle = 0;
    }
  }

  private static native long createSessionContext();

  private static native long createSessionContextWithOptions(byte[] optionsBytes);

  private static native long createDataFrame(long handle, String sql);

  private static native long createDataFrameFromProto(long handle, byte[] planBytes);

  private static native byte[] tableSchemaIpc(long handle, String tableName);

  private static native String getOptionNative(long handle, String key);

  private static native long[] memoryUsageNative(long handle);

  private static native long[] runtimeStatsNative(long handle);

  private static native void registerParquetWithOptions(
      long handle, String name, String path, byte[] optionsBytes, byte[] schemaIpcBytes);

  private static native long readParquetWithOptions(
      long handle, String path, byte[] optionsBytes, byte[] schemaIpcBytes);

  private static native void registerCsvWithOptions(
      long handle, String name, String path, byte[] optionsBytes, byte[] schemaIpcBytes);

  private static native long readCsvWithOptions(
      long handle, String path, byte[] optionsBytes, byte[] schemaIpcBytes);

  private static native void registerArrowWithOptions(
      long handle, String name, String path, byte[] optionsBytes, byte[] schemaIpcBytes);

  private static native long readArrowWithOptions(
      long handle, String path, byte[] optionsBytes, byte[] schemaIpcBytes);

  private static native void registerAvroWithOptions(
      long handle, String name, String path, byte[] optionsBytes, byte[] schemaIpcBytes);

  private static native long readAvroWithOptions(
      long handle, String path, byte[] optionsBytes, byte[] schemaIpcBytes);

  private static native void registerJsonWithOptions(
      long handle, String name, String path, byte[] optionsBytes, byte[] schemaIpcBytes);

  private static native long readJsonWithOptions(
      long handle, String path, byte[] optionsBytes, byte[] schemaIpcBytes);

  private static native void closeSessionContext(long handle);

  private static native void registerScalarUdf(
      long handle, String name, byte[] signatureSchemaBytes, byte volatility, ScalarFunction impl);

  private static native void registerTableNative(
      long handle, String name, byte[] schemaIpcBytes, TableProvider provider);
}
