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

import unittest
from unittest import mock
from types import SimpleNamespace

from pypaimon.common.options.core_options import GlobalIndexSearchMode
from pypaimon.index.pk.primary_key_index_definitions import PrimaryKeyIndexDefinitions
from pypaimon.index.pk.primary_key_index_source_file import PrimaryKeyIndexSourceFile
from pypaimon.index.pk.primary_key_index_source_meta import PrimaryKeyIndexSourceMeta
from pypaimon.schema.data_types import AtomicType, DataField
from pypaimon.schema.table_schema import TableSchema
from pypaimon.table.source.full_text_search_builder import FullTextSearchBuilderImpl
from pypaimon.table.source.primary_key_full_text_read import PrimaryKeyFullTextRead
from pypaimon.table.source.primary_key_full_text_scan import PrimaryKeyFullTextScan
from pypaimon.table.source.primary_key_full_text_scan import (
    _current_payloads, _should_read_source as _full_text_should_read_source)
from pypaimon.table.source.primary_key_vector_read import (
    PrimaryKeyVectorRead, _allowed, _include_row_ids)
from pypaimon.table.source.primary_key_vector_scan import (
    PrimaryKeyVectorScan, _bucket_splits, _should_read_source)
from pypaimon.table.source.primary_key_scored_result import (
    PrimaryKeyScoredResult, PrimaryKeySearchPosition)
from pypaimon.table.source.vector_search_builder import VectorSearchBuilderImpl
from pypaimon.manifest.schema.data_file_meta import DataFileMeta
from pypaimon.read.split import DataSplit
from pypaimon.table.row.generic_row import GenericRow
from pypaimon.utils.range import Range


class PrimaryKeyIndexDefinitionsTest(unittest.TestCase):
    def test_definitions_follow_schema_order(self):
        schema = TableSchema(
            fields=[DataField(0, "id", AtomicType("INT")),
                    DataField(1, "name", AtomicType("STRING"))],
            options={"pk-btree.index.columns": "name",
                     "pk-bitmap.index.columns": "id",
                     "fields.name.pk-btree.index.options":
                         '{"block-size": "4096"}'},
        )

        definitions = PrimaryKeyIndexDefinitions.create(schema).definitions

        self.assertEqual(["id", "name"], [d.column for d in definitions])
        self.assertEqual(["bitmap", "btree"], [d.index_type for d in definitions])
        self.assertEqual(
            "4096", definitions[1].options.to_map()["btree-index.block-size"])

    def test_source_meta_uses_java_binary_layout(self):
        meta = PrimaryKeyIndexSourceMeta(
            2, [PrimaryKeyIndexSourceFile("数据-\0-😀.parquet", 7)])

        restored = PrimaryKeyIndexSourceMeta.deserialize(meta.serialize())

        self.assertEqual(meta, restored)

    def test_vector_and_full_text_definitions(self):
        schema = TableSchema(
            fields=[DataField(0, "embedding", AtomicType("VECTOR(3)")),
                    DataField(1, "content", AtomicType("STRING"))],
            options={
                "pk-vector.index.columns": "embedding",
                "fields.embedding.pk-vector.index.type": "vindex",
                "fields.embedding.pk-vector.index.options": '{"m": "16"}',
                "pk-full-text.index.columns": "content",
                "fields.content.pk-full-text.index.options": '{"analyzer": "english"}',
            })

        definitions = PrimaryKeyIndexDefinitions.create(schema).definitions

        self.assertEqual(["vindex", "full-text"], [d.index_type for d in definitions])
        self.assertEqual("16", definitions[0].options.to_map()["vindex.m"])
        self.assertEqual(
            "inner_product", definitions[0].options.to_map()["vindex.metric"])
        self.assertEqual(
            "english", definitions[1].options.to_map()["full-text.analyzer"])

    def test_builders_route_primary_key_searches(self):
        schema = TableSchema(
            fields=[DataField(0, "embedding", AtomicType("VECTOR(3)")),
                    DataField(1, "content", AtomicType("STRING"))],
            options={
                "pk-vector.index.columns": "embedding",
                "fields.embedding.pk-vector.index.type": "vindex",
                "pk-full-text.index.columns": "content",
            })
        table = SimpleNamespace(
            fields=schema.fields, table_schema=schema, partition_keys=[])

        vector = (VectorSearchBuilderImpl(table)
                  .with_vector_column("embedding")
                  .with_query_vector([1.0, 2.0, 3.0])
                  .with_limit(2))
        full_text = (FullTextSearchBuilderImpl(table)
                     .with_query("content", "paimon")
                     .with_limit(2))

        self.assertIsInstance(vector.new_vector_search_scan(), PrimaryKeyVectorScan)
        self.assertIsInstance(vector.new_vector_search_read(), PrimaryKeyVectorRead)
        self.assertIsInstance(full_text.new_full_text_scan(), PrimaryKeyFullTextScan)
        self.assertIsInstance(full_text.new_full_text_read(), PrimaryKeyFullTextRead)

    def test_primary_key_full_text_only_supports_fast_mode(self):
        fields = [DataField(0, "id", AtomicType("INT")),
                  DataField(1, "content", AtomicType("STRING"))]

        def builder(mode):
            schema = TableSchema(
                fields=fields,
                options={"pk-full-text.index.columns": "content",
                         "global-index.search-mode": mode})
            table = SimpleNamespace(
                fields=fields, table_schema=schema, partition_keys=[])
            return (FullTextSearchBuilderImpl(table)
                    .with_query("content", "paimon").with_limit(2))

        self.assertIsInstance(
            builder("fast").new_full_text_read(), PrimaryKeyFullTextRead)
        schema = TableSchema(
            fields=fields,
            options={
                "pk-full-text.index.columns": "content",
                "global-index.search-mode": "full",
                "full-text-index.search-mode": "fast",
            })
        table = SimpleNamespace(
            fields=fields, table_schema=schema, partition_keys=[])
        self.assertIsInstance(
            FullTextSearchBuilderImpl(table)
            .with_query("content", "paimon")
            .with_limit(2)
            .new_full_text_read(),
            PrimaryKeyFullTextRead,
        )
        for mode in ("full", "detail"):
            with self.subTest(mode=mode), self.assertRaisesRegex(
                    NotImplementedError, "only supports the FAST"):
                builder(mode).new_full_text_read()

    def test_primary_key_full_text_reader_cannot_bypass_fast_mode(self):
        field = DataField(1, "content", AtomicType("STRING"))
        for mode in ("full", "detail"):
            schema = TableSchema(
                fields=[field],
                options={"global-index.search-mode": mode})
            table = SimpleNamespace(table_schema=schema)
            with self.subTest(mode=mode), self.assertRaisesRegex(
                    NotImplementedError, "only supports the FAST"):
                PrimaryKeyFullTextRead(
                    table, 1, field, "paimon", definition=None)

    def test_scored_result_uses_file_local_positions(self):
        partition = GenericRow([], [])
        data_file = DataFileMeta(
            file_name="data.parquet", file_size=10, row_count=5,
            min_key=None, max_key=None, key_stats=None, value_stats=None,
            min_sequence_number=0, max_sequence_number=0,
            schema_id=0, level=1, extra_files=[])
        source = DataSplit([data_file], partition, 3)
        partition_bytes = repr(tuple(partition.values)).encode("utf-8")

        result = PrimaryKeyScoredResult(7, [source], [
            PrimaryKeySearchPosition(partition_bytes, 3, "data.parquet", 1, 0.9),
            PrimaryKeySearchPosition(partition_bytes, 3, "data.parquet", 2, 0.8),
            PrimaryKeySearchPosition(partition_bytes, 3, "data.parquet", 4, 0.7),
        ])

        self.assertEqual(7, result.snapshot_id)
        self.assertEqual([(1, 2), (4, 4)],
                         [(r.from_, r.to) for r in result.splits[0].row_ranges()])
        self.assertEqual([0.9, 0.8, 0.7], result.splits[0].scores())

    def test_pk_vector_filter_is_applied_before_top_k(self):
        split = SimpleNamespace(row_ranges_by_file={
            "first.parquet": (Range(1, 2),),
            "second.parquet": (Range(0, 1),),
        })
        sources = [
            PrimaryKeyIndexSourceFile("first.parquet", 4),
            PrimaryKeyIndexSourceFile("second.parquet", 3),
        ]
        deleted = {(id(split), "first.parquet"): {1}}

        include = _include_row_ids(split, sources, deleted)

        # first[2] maps to payload row 2; second[0,1] map to rows 4 and 5.
        self.assertEqual([2, 4, 5], list(include))

    def test_empty_pk_vector_ranges_mean_no_candidates(self):
        split = SimpleNamespace(row_ranges_by_file={"data.parquet": tuple()})
        source = PrimaryKeyIndexSourceFile("data.parquet", 3)

        include = _include_row_ids(split, [source], {})

        self.assertEqual([], list(include))
        self.assertFalse(_allowed(split, "data.parquet", 0))

    def test_pk_vector_refine_factor_overflow_is_rejected_before_search(self):
        table = SimpleNamespace(options=SimpleNamespace(
            primary_key_vector_index_type=lambda column: "ivf-pq",
            vector_index_search_mode=lambda: GlobalIndexSearchMode.FAST))
        reader = PrimaryKeyVectorRead(
            table, 2147483647, DataField(1, "embedding", AtomicType("VECTOR(1)")),
            [0.0], options={"ivf.refine_factor": "2"})

        from pypaimon.table.source.primary_key_vector_scan import (
            PrimaryKeyVectorScanPlan)
        plan = PrimaryKeyVectorScanPlan(1, [])
        with self.assertRaisesRegex(ValueError, "limit overflow"):
            reader.read_plan(plan)

    def test_pk_vector_raw_fallback_uses_vector_search_mode(self):
        from pypaimon.table.source.primary_key_vector_scan import (
            PrimaryKeyVectorScanPlan)

        field = DataField(1, "embedding", AtomicType("VECTOR(1)"))
        plan = PrimaryKeyVectorScanPlan(1, [])
        for mode, expected_calls in (
                (GlobalIndexSearchMode.FAST, 0),
                (GlobalIndexSearchMode.FULL, 1),
                (GlobalIndexSearchMode.DETAIL, 1)):
            table = SimpleNamespace(options=SimpleNamespace(
                primary_key_vector_index_type=lambda column: "ivf-pq",
                vector_index_search_mode=lambda mode=mode: mode))
            reader = PrimaryKeyVectorRead(table, 1, field, [0.0])
            with self.subTest(mode=mode), mock.patch.object(
                    reader, "_raw_candidates", return_value=[]) as raw_candidates:
                reader.read_plan(plan)
                self.assertEqual(expected_calls, raw_candidates.call_count)

    def test_full_text_payload_planner_rejects_duplicate_and_stale_levels(self):
        files = [SimpleNamespace(file_name="a.parquet", row_count=2, level=1)]
        source_meta = PrimaryKeyIndexSourceMeta(
            1, [PrimaryKeyIndexSourceFile("a.parquet", 2)])

        def payload(name, row_count=2, start=0, end=1):
            return SimpleNamespace(
                file_name=name, row_count=row_count,
                global_index_meta=SimpleNamespace(
                    row_range_start=start, row_range_end=end))

        current = payload("current")
        duplicate = payload("duplicate")
        stale = payload("stale", end=0)
        with mock.patch.object(
                PrimaryKeyIndexSourceMeta, "from_index_file",
                return_value=source_meta):
            selected, covered = _current_payloads(
                files, [current, duplicate, stale])

        self.assertEqual([], selected)
        self.assertEqual(set(), covered)

    def test_vector_payload_planner_rejects_duplicate_level(self):
        partition = GenericRow([], [])
        data_file = DataFileMeta(
            file_name="a.parquet", file_size=10, row_count=2,
            min_key=None, max_key=None, key_stats=None, value_stats=None,
            min_sequence_number=0, max_sequence_number=0,
            schema_id=0, level=1, extra_files=[], file_source=1)
        split = DataSplit([data_file], partition, 0)
        source_meta = PrimaryKeyIndexSourceMeta(
            1, [PrimaryKeyIndexSourceFile("a.parquet", 2)])

        def payload(name):
            return SimpleNamespace(
                file_name=name, row_count=2,
                global_index_meta=SimpleNamespace(
                    row_range_start=0, row_range_end=1))

        entries = [
            SimpleNamespace(partition=partition, bucket=0,
                            index_file=payload("first")),
            SimpleNamespace(partition=partition, bucket=0,
                            index_file=payload("duplicate")),
        ]
        with mock.patch.object(
                PrimaryKeyIndexSourceMeta, "from_index_file",
                return_value=source_meta):
            planned = _bucket_splits([(split, {})], entries)

        self.assertEqual(tuple(), planned[0].payloads)
        self.assertEqual(("a.parquet",), planned[0].uncovered_data_files)

    def test_full_text_source_policy_matches_java(self):
        self.assertTrue(_full_text_should_read_source(
            SimpleNamespace(file_source=1, level=1)))
        self.assertFalse(_full_text_should_read_source(
            SimpleNamespace(file_source=1, level=0)))
        self.assertFalse(_full_text_should_read_source(
            SimpleNamespace(file_source=0, level=2)))

    def test_pk_vector_scan_pins_source_and_index_to_same_snapshot(self):
        schema = TableSchema(
            fields=[DataField(1, "embedding", AtomicType("VECTOR(1)"))],
            options={"pk-vector.index.columns": "embedding",
                     "scan.tag-name": "old-tag"})
        snapshot = SimpleNamespace(id=17)
        copied_options = []
        index_snapshots = []

        class _ScanTable:
            table_schema = schema
            fields = schema.fields

            def new_read_builder(self):
                plan = SimpleNamespace(splits=lambda: [])
                scan = SimpleNamespace(plan=lambda: plan)
                return SimpleNamespace(new_scan=lambda: scan)

        class _Table:
            table_schema = schema
            fields = schema.fields

            def tag_manager(self):
                return None

            def snapshot_manager(self):
                return SimpleNamespace(get_latest_snapshot=lambda: snapshot)

            def copy(self, options):
                copied_options.append(options)
                return _ScanTable()

        class _IndexHandler:
            def __init__(self, table):
                self.table = table

            def scan(self, selected_snapshot, entry_filter):
                index_snapshots.append(selected_snapshot)
                return []

        with mock.patch(
                "pypaimon.table.source.primary_key_vector_scan."
                "TimeTravelUtil.try_travel_to_snapshot",
                return_value=snapshot), mock.patch(
                "pypaimon.table.source.primary_key_vector_scan.IndexFileHandler",
                _IndexHandler):
            plan = PrimaryKeyVectorScan(
                _Table(), schema.fields[0], index_type="vindex").scan()

        self.assertEqual(17, plan.snapshot_id)
        self.assertEqual([snapshot], index_snapshots)
        self.assertEqual("from-snapshot", copied_options[0]["scan.mode"])
        self.assertEqual("17", copied_options[0]["scan.snapshot-id"])
        self.assertIsNone(copied_options[0]["scan.tag-name"])

    def test_pk_index_source_policy_matches_java(self):
        compact = SimpleNamespace(file_source=1, level=1)
        append = SimpleNamespace(file_source=0, level=1)
        level_zero = SimpleNamespace(file_source=1, level=0)
        legacy = SimpleNamespace(file_source=None, level=1)

        self.assertTrue(_should_read_source(compact))
        self.assertFalse(_should_read_source(append))
        self.assertFalse(_should_read_source(level_zero))
        self.assertFalse(_should_read_source(legacy))


if __name__ == "__main__":
    unittest.main()
