################################################################################
#  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 Dict, List, Set

from pypaimon.common.predicate import Predicate
from pypaimon.common.predicate_builder import PredicateBuilder


def trim_and_transform_predicate(input_predicate: Predicate, all_fields: List[str], trimmed_keys: List[str]):
    new_predicate = trim_predicate_by_fields(input_predicate, trimmed_keys)
    part_to_index = {element: idx for idx, element in enumerate(trimmed_keys)}
    mapping: Dict[int, int] = {
        i: part_to_index.get(all_fields[i], -1)
        for i in range(len(all_fields))
    }
    return _change_index(new_predicate, mapping)


def trim_predicate_by_fields(input_predicate: Predicate, trimmed_keys: List[str]):
    if not input_predicate or not trimmed_keys:
        return None

    predicates: list[Predicate] = _split_and(input_predicate)
    predicates = [element for element in predicates if _get_all_fields(element).issubset(trimmed_keys)]
    return PredicateBuilder.and_predicates(predicates)


def _split_and(input_predicate: Predicate):
    if not input_predicate:
        return list()

    if input_predicate.method == 'and':
        return [p for element in (input_predicate.literals or []) for p in _split_and(element)]

    return [input_predicate]


def _change_index(input_predicate: Predicate, mapping: Dict[int, int]):
    if not input_predicate:
        return None

    if input_predicate.method == 'and' or input_predicate.method == 'or':
        predicates: list[Predicate] = input_predicate.literals
        new_predicates = [_change_index(element, mapping) for element in predicates]
        return input_predicate.new_literals(new_predicates)

    return input_predicate.new_index(mapping[input_predicate.index])


def _get_all_fields(predicate: Predicate) -> Set[str]:
    if predicate.field is not None:
        return {predicate.field}
    involved_fields = set()
    if predicate.literals:
        for sub_predicate in predicate.literals:
            involved_fields.update(_get_all_fields(sub_predicate))
    return involved_fields
