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

from typing import Any, Callable, Dict, List, Optional, Sequence, TypeAlias, Union

import pyarrow

ArrowTypeLike: TypeAlias = Union[pyarrow.DataType, pyarrow.Field, str]
InputFieldsLike: TypeAlias = Union[ArrowTypeLike, Sequence[ArrowTypeLike]]
VolatilityLike: TypeAlias = Union[str, Any]

class DataField:
    def name(self) -> str: ...
    def field_type(self) -> str: ...
    def is_nullable(self) -> bool: ...
    def description(self) -> Optional[str]: ...

class TableSchema:
    def fields(self) -> List[DataField]: ...
    def partition_keys(self) -> List[str]: ...
    def primary_keys(self) -> List[str]: ...
    def options(self) -> Dict[str, str]: ...
    def comment(self) -> Optional[str]: ...

class Split:
    def __init__(self, state: bytes) -> None: ...
    def row_count(self) -> int: ...
    # Java SplitSerializer v1 binary: a DataSplit (v8), or an IndexedSplit when the split has row ranges.
    def serialize(self) -> bytes: ...

class Plan:
    def splits(self) -> List[Split]: ...
    def __len__(self) -> int: ...

class TableScan:
    def plan(self) -> Plan: ...

class TableRead:
    def read(self, splits: Sequence[Split]) -> List[pyarrow.RecordBatch]: ...

class ReadBuilder:
    def with_projection(self, columns: List[str]) -> "ReadBuilder": ...
    def with_case_sensitive(self, case_sensitive: bool) -> "ReadBuilder":
        """
        Set whether column-name matching (projection and predicate column
        resolution) is case-sensitive. Defaults to ``True`` (exact match).

        Projection resolution is lazy, so this is order-independent with
        ``with_projection``. Predicates built via ``with_filter`` capture case
        sensitivity when they are constructed, so ``with_case_sensitive`` must be
        set before ``with_filter`` for the filter to honor it.
        """
        ...
    def with_limit(self, limit: int) -> "ReadBuilder": ...
    def with_filter(self, predicate: dict) -> "ReadBuilder": ...
    def new_scan(self) -> TableScan: ...
    def new_read(self) -> "TableRead": ...

# ---- #285: observability ----
class Snapshot:
    def id(self) -> int: ...
    def commit_time_ms(self) -> int: ...
    def total_record_count(self) -> Optional[int]: ...
    def delta_record_count(self) -> Optional[int]: ...
    def commit_kind(self) -> str: ...

class Tag:
    def name(self) -> str: ...
    def snapshot_id(self) -> int: ...

class PartitionStat:
    def partition(self) -> Dict[str, str]: ...
    def record_count(self) -> int: ...
    def file_count(self) -> int: ...
    def total_size_bytes(self) -> int: ...

class Table:
    def identifier(self) -> str: ...
    def location(self) -> str: ...
    def schema(self) -> TableSchema: ...
    def new_read_builder(self, options: Optional[Dict[str, str]] = None) -> ReadBuilder: ...
    def new_write_builder(self) -> "WriteBuilder": ...
    def latest_snapshot(self) -> Optional[Snapshot]:
        """
        Warning: This method blocks on a DataFusion runtime.
        Calling this from an active asyncio event loop will result in a panic.
        """
        ...
    def list_snapshots(self) -> List[Snapshot]:
        """
        Returns all snapshots ordered newest first (descending by ID).

        Warning: This method blocks on a DataFusion runtime.
        Calling this from an active asyncio event loop will result in a panic.
        """
        ...
    def list_tags(self) -> List[Tag]:
        """
        Warning: This method blocks on a DataFusion runtime.
        Calling this from an active asyncio event loop will result in a panic.
        """
        ...
    def list_partitions(self) -> List[Dict[str, str]]:
        """
        Warning: This method blocks on a DataFusion runtime.
        Calling this from an active asyncio event loop will result in a panic.
        """
        ...
    def partition_stats(self) -> List[PartitionStat]:
        """
        Warning: This method blocks on a DataFusion runtime.
        Calling this from an active asyncio event loop will result in a panic.
        """
        ...

class CommitMessage: ...

class TableWrite:
    def write_arrow(self, batch: pyarrow.RecordBatch) -> None: ...
    def prepare_commit(self) -> List[CommitMessage]: ...

class TableCommit:
    def commit(self, messages: Sequence[CommitMessage]) -> None: ...

class WriteBuilder:
    def new_write(self) -> TableWrite: ...
    def new_commit(self) -> TableCommit: ...

class PaimonCatalog:
    def __init__(self, catalog_options: Dict[str, str]) -> None: ...
    def __datafusion_catalog_provider__(self, session: Any) -> object: ...
    def list_databases(self) -> List[str]: ...
    def list_tables(self, database_name: str) -> List[str]: ...
    def get_table(self, identifier: str) -> Table: ...

class PythonScalarUDF:
    def __init__(
        self,
        name: str,
        func: Callable[..., pyarrow.Array],
        input_fields: InputFieldsLike,
        return_field: ArrowTypeLike,
        volatility: VolatilityLike,
    ) -> None: ...
    @staticmethod
    def udf(
        func: Callable[..., pyarrow.Array],
        input_fields: InputFieldsLike,
        return_field: ArrowTypeLike,
        volatility: VolatilityLike,
        name: Optional[str] = None,
    ) -> "PythonScalarUDF": ...
    @property
    def name(self) -> str: ...

def udf(
    func: Callable[..., pyarrow.Array],
    input_fields: InputFieldsLike,
    return_field: ArrowTypeLike,
    volatility: VolatilityLike,
    name: Optional[str] = None,
) -> PythonScalarUDF:
    """
    Create a scalar UDF.

    This mirrors DataFusion Python's function-style API:
    ``udf(func, input_fields, return_field, volatility, name)``.
    ``input_fields`` and ``return_field`` accept PyArrow DataType or Field
    values. String type names remain accepted for compatibility.
    """
    ...

class SQLContext:
    def __init__(self) -> None: ...
    def register_catalog(
        self, catalog_name: str, catalog_options: Dict[str, str]
    ) -> None: ...
    def set_current_catalog(self, catalog_name: str) -> None: ...
    def set_current_database(self, database_name: str) -> None: ...
    def register_batch(self, name: str, batch: pyarrow.RecordBatch) -> None: ...
    def register_udf(self, udf: PythonScalarUDF) -> None: ...
    def sql(self, sql: str) -> List[pyarrow.RecordBatch]: ...
