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

"""Full-text read to read index files."""

from abc import ABC, abstractmethod
from concurrent.futures import wait
from typing import List

from pypaimon.globalindex.full_text_search import FullTextSearch
from pypaimon.globalindex.global_index_meta import GlobalIndexIOMeta
from pypaimon.globalindex.global_index_result import GlobalIndexResult
from pypaimon.globalindex.offset_global_index_reader import OffsetGlobalIndexReader
from pypaimon.globalindex.vector_search_result import DictBasedScoredIndexResult
from pypaimon.table.source.full_text_search_split import FullTextSearchSplit
from pypaimon.table.source.full_text_scan import FullTextScanPlan


class FullTextRead(ABC):
    """Full-text read to read index files."""

    def read_plan(self, plan: FullTextScanPlan) -> GlobalIndexResult:
        return self.read(plan.splits())

    @abstractmethod
    def read(self, splits: List[FullTextSearchSplit]) -> GlobalIndexResult:
        pass


class FullTextReadImpl(FullTextRead):
    """Implementation for FullTextRead."""

    def __init__(
        self,
        table: 'FileStoreTable',
        limit: int,
        text_column: 'DataField',
        query_text: str,
        query_operator: str = "or"
    ):
        self._table = table
        self._limit = limit
        self._text_column = text_column
        self._query_text = query_text
        self._query_operator = query_operator

    def read(self, splits: List[FullTextSearchSplit]) -> GlobalIndexResult:
        if not splits:
            return GlobalIndexResult.create_empty()

        futures = [
            self._eval(
                split.row_range_start, split.row_range_end,
                split.full_text_index_files
            )
            for split in splits
        ]

        wait(futures)

        merged_scores = {}
        for future in futures:
            split_result = future.result()
            if split_result is not None:
                score_getter = split_result.score_getter()
                for row_id in split_result.results():
                    if row_id not in merged_scores:
                        merged_scores[row_id] = score_getter(row_id)

        return DictBasedScoredIndexResult(merged_scores).top_k(self._limit)

    def _eval(self, row_range_start, row_range_end, full_text_index_files):
        index_io_meta_list = []
        for index_file in full_text_index_files:
            meta = index_file.global_index_meta
            assert meta is not None
            index_io_meta_list.append(
                GlobalIndexIOMeta(
                    file_name=index_file.file_name,
                    file_size=index_file.file_size,
                    metadata=meta.index_meta,
                    external_path=index_file.external_path,
                )
            )

        index_type = full_text_index_files[0].index_type
        index_path = self._table.path_factory().global_index_path_factory().index_path()
        file_io = self._table.file_io

        reader = _create_full_text_reader(
            index_type, file_io, index_path,
            index_io_meta_list
        )

        full_text_search = FullTextSearch(
            query_text=self._query_text,
            limit=self._limit,
            field_name=self._text_column.name,
            query_operator=self._query_operator
        )

        offset_reader = OffsetGlobalIndexReader(reader, row_range_start, row_range_end)
        future = offset_reader.visit_full_text_search(full_text_search)
        future.add_done_callback(lambda _: reader.close())
        return future


def _create_full_text_reader(index_type, file_io, index_path, index_io_meta_list):
    """Create a global index reader for full-text search."""
    from pypaimon.globalindex.tantivy.tantivy_full_text_global_index_reader import (
        TANTIVY_FULLTEXT_IDENTIFIER,
    )
    if index_type == TANTIVY_FULLTEXT_IDENTIFIER:
        from pypaimon.globalindex.tantivy.tantivy_full_text_global_index_reader import (
            TantivyFullTextGlobalIndexReader,
        )
        return TantivyFullTextGlobalIndexReader(file_io, index_path, index_io_meta_list)
    raise ValueError(f"Unsupported full-text index type: '{index_type}'")
