################################################################################
#  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 pypaimon.read.split import Split


class QueryAuthSplit(Split):

    def __init__(self, split: Split, auth_result):
        self._split = split
        self._auth_result = auth_result

    @property
    def split(self) -> Split:
        return self._split

    @property
    def auth_result(self):
        return self._auth_result

    @property
    def row_count(self) -> int:
        return self._split.row_count

    @property
    def files(self):
        return self._split.files

    @property
    def partition(self):
        return self._split.partition

    @property
    def bucket(self) -> int:
        return self._split.bucket

    @property
    def is_streaming(self):
        return getattr(self._split, 'is_streaming', False)

    @property
    def raw_convertible(self) -> bool:
        return self._split.raw_convertible

    def merged_row_count(self):
        if self._auth_result.filter:
            return None
        return self._split.merged_row_count()

    def __getattr__(self, name):
        if name.startswith('_'):
            raise AttributeError(name)
        return getattr(self._split, name)


def query_auth_columns(read_type, predicate=None, partition_predicate=None, columns=None):
    select = [f.name for f in read_type] if read_type else None
    if select is not None:
        from pypaimon.read.push_down_utils import predicate_field_names
        for query_filter in (predicate, partition_predicate):
            if query_filter is not None:
                names = predicate_field_names(query_filter) - set(select)
                if columns is not None:
                    # other names are leaf aliases of projected columns, which are already selected
                    names &= columns
                select.extend(sorted(names))
    return select


def resolve_auth_result(query_auth_fn, read_type, predicate=None, partition_predicate=None,
                        columns=None):
    if query_auth_fn is None:
        return None
    select = query_auth_columns(read_type, predicate, partition_predicate, columns)
    result = query_auth_fn(select)
    if result is None or not result.has_restrictions:
        return None
    return result


def wrap_plan_with_auth(auth_result, plan):
    if auth_result is not None:
        return auth_result.convert_plan(plan)
    return plan