# 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 pypaimon.common.options.core_options import CoreOptions
from pypaimon.schema.column_directive_utils import (
    apply_add_column_directive,
    apply_directives,
    parse_add_column_comment,
    remove_dropped_directive_options,
)
from pypaimon.schema.data_types import (
    ArrayType, AtomicType, DataField, MapType, VectorType,
)


class TestParseAddColumnComment(unittest.TestCase):

    def test_none_and_empty(self):
        self.assertIsNone(parse_add_column_comment(None))
        self.assertIsNone(parse_add_column_comment(""))
        self.assertIsNone(parse_add_column_comment("normal comment"))

    def test_blob_field(self):
        d = parse_add_column_comment("__BLOB_FIELD; picture")
        self.assertEqual(d.option_key, CoreOptions.BLOB_FIELD.key())
        self.assertEqual(d.real_comment, "picture")
        self.assertFalse(d.is_vector)

    def test_blob_field_bare(self):
        d = parse_add_column_comment("__BLOB_FIELD")
        self.assertIsNone(d.real_comment)

    def test_blob_descriptor_field(self):
        d = parse_add_column_comment("__BLOB_DESCRIPTOR_FIELD; desc")
        self.assertEqual(d.option_key, CoreOptions.BLOB_DESCRIPTOR_FIELD.key())

    def test_blob_view_field(self):
        d = parse_add_column_comment("__BLOB_VIEW_FIELD; view")
        self.assertEqual(d.option_key, CoreOptions.BLOB_VIEW_FIELD.key())

    def test_vector_field(self):
        d = parse_add_column_comment("__VECTOR_FIELD;128; embedding")
        self.assertEqual(d.option_key, CoreOptions.VECTOR_FIELD.key())
        self.assertTrue(d.is_vector)
        self.assertEqual(d.vector_dim, 128)
        self.assertEqual(d.real_comment, "embedding")

    def test_vector_field_no_comment(self):
        d = parse_add_column_comment("__VECTOR_FIELD;64")
        self.assertEqual(d.vector_dim, 64)
        self.assertIsNone(d.real_comment)

    def test_unknown_blob_directive_rejected(self):
        with self.assertRaises(ValueError):
            parse_add_column_comment("__BLOB_UNKNOWN")

    def test_vector_without_dim_rejected(self):
        with self.assertRaises(ValueError):
            parse_add_column_comment("__VECTOR_FIELD")

    def test_vector_non_integer_dim_rejected(self):
        with self.assertRaises(ValueError):
            parse_add_column_comment("__VECTOR_FIELD;abc")


class TestApplyAddColumnDirective(unittest.TestCase):

    def test_non_directive_returns_none(self):
        opts = {}
        result = apply_add_column_directive(
            "normal", "col", AtomicType("BYTES"), opts
        )
        self.assertIsNone(result)
        self.assertEqual(opts, {})

    def test_blob_field(self):
        opts = {}
        result = apply_add_column_directive(
            "__BLOB_FIELD; pic", "pic", AtomicType("BYTES"), opts
        )
        self.assertIsNotNone(result)
        self.assertEqual(result.type.type, "BLOB")
        self.assertEqual(result.comment, "pic")
        self.assertEqual(opts[CoreOptions.BLOB_FIELD.key()], "pic")

    def test_blob_field_array_source_type(self):
        opts = {}
        result = apply_add_column_directive(
            "__BLOB_FIELD; pictures",
            "pictures",
            ArrayType(True, AtomicType("BYTES", False)),
            opts,
        )

        self.assertIsNotNone(result)
        self.assertIsInstance(result.type, ArrayType)
        self.assertEqual(result.type.element.type, "BLOB")
        self.assertFalse(result.type.element.nullable)
        self.assertEqual(result.comment, "pictures")
        self.assertEqual(opts[CoreOptions.BLOB_FIELD.key()], "pictures")

    def test_inline_blob_directive_rejects_array_source_type(self):
        with self.assertRaisesRegex(ValueError, "ARRAY<BLOB> is only supported"):
            apply_add_column_directive(
                "__BLOB_DESCRIPTOR_FIELD",
                "pictures",
                ArrayType(True, AtomicType("BYTES")),
                {},
            )

    def test_blob_field_map_source_type(self):
        opts = {}
        result = apply_add_column_directive(
            "__BLOB_FIELD; pictures",
            "pictures",
            MapType(
                True,
                AtomicType("STRING", False),
                AtomicType("BYTES", False),
            ),
            opts,
        )

        self.assertIsInstance(result.type, MapType)
        self.assertEqual(result.type.key, AtomicType("STRING", False))
        self.assertEqual(result.type.value, AtomicType("BLOB", False))
        self.assertEqual(result.comment, "pictures")
        self.assertEqual(opts[CoreOptions.BLOB_FIELD.key()], "pictures")

    def test_inline_blob_directive_rejects_map_source_type(self):
        with self.assertRaisesRegex(ValueError, "MAP<X, BLOB> is only supported"):
            apply_add_column_directive(
                "__BLOB_DESCRIPTOR_FIELD",
                "pictures",
                MapType(
                    True,
                    AtomicType("INT", False),
                    AtomicType("BYTES"),
                ),
                {},
            )

    def test_blob_field_map_allows_unsupported_key_type(self):
        result = apply_add_column_directive(
            "__BLOB_FIELD",
            "pictures",
            MapType(
                True,
                AtomicType("BOOLEAN", False),
                AtomicType("BYTES"),
            ),
            {},
        )

        self.assertIsInstance(result.type, MapType)
        self.assertEqual(result.type.key, AtomicType("BOOLEAN", False))
        self.assertEqual(result.type.value, AtomicType("BLOB"))

    def test_vector_field(self):
        opts = {}
        result = apply_add_column_directive(
            "__VECTOR_FIELD;128; emb",
            "emb",
            ArrayType(True, AtomicType("FLOAT")),
            opts
        )
        self.assertIsNotNone(result)
        self.assertIsInstance(result.type, VectorType)
        self.assertEqual(result.type.length, 128)
        self.assertEqual(result.comment, "emb")
        self.assertEqual(opts[CoreOptions.VECTOR_FIELD.key()], "emb")

    def test_blob_rejects_non_binary(self):
        with self.assertRaises(ValueError):
            apply_add_column_directive(
                "__BLOB_FIELD", "col", AtomicType("INT"), {}
            )

    def test_vector_rejects_non_array(self):
        with self.assertRaises(ValueError):
            apply_add_column_directive(
                "__VECTOR_FIELD;128", "col", AtomicType("INT"), {}
            )

    def test_appends_to_existing(self):
        opts = {CoreOptions.BLOB_FIELD.key(): "a"}
        apply_add_column_directive(
            "__BLOB_FIELD", "b", AtomicType("BYTES"), opts
        )
        self.assertEqual(opts[CoreOptions.BLOB_FIELD.key()], "a,b")

    def test_migrates_legacy_fallback_key(self):
        opts = {"blob.stored-descriptor-fields": "legacy_col"}
        apply_add_column_directive(
            "__BLOB_DESCRIPTOR_FIELD", "new_col", AtomicType("BYTES"), opts
        )
        self.assertEqual(
            opts[CoreOptions.BLOB_DESCRIPTOR_FIELD.key()], "legacy_col,new_col"
        )
        self.assertNotIn("blob.stored-descriptor-fields", opts)


class TestApplyDirectives(unittest.TestCase):

    def test_no_directives(self):
        fields = [DataField(0, "k", AtomicType("INT"))]
        opts = {}
        changed = apply_directives(fields, opts)
        self.assertFalse(changed)
        self.assertEqual(fields[0].type.type, "INT")

    def test_mixed_fields(self):
        fields = [
            DataField(0, "k", AtomicType("INT")),
            DataField(1, "pic", AtomicType("BYTES"), "__BLOB_FIELD; picture"),
            DataField(
                2, "emb", ArrayType(True, AtomicType("FLOAT")),
                "__VECTOR_FIELD;64; my emb"
            ),
        ]
        opts = {}
        changed = apply_directives(fields, opts)
        self.assertTrue(changed)
        self.assertEqual(fields[1].type.type, "BLOB")
        self.assertEqual(fields[1].description, "picture")
        self.assertIsInstance(fields[2].type, VectorType)
        self.assertEqual(fields[2].type.length, 64)
        self.assertEqual(fields[2].description, "my emb")
        self.assertEqual(opts[CoreOptions.BLOB_FIELD.key()], "pic")
        self.assertEqual(opts[CoreOptions.VECTOR_FIELD.key()], "emb")


class TestRemoveDroppedDirectiveOptions(unittest.TestCase):

    def test_drop_blob(self):
        opts = {
            CoreOptions.BLOB_FIELD.key(): "a,b",
            CoreOptions.BLOB_DESCRIPTOR_FIELD.key(): "b",
            CoreOptions.VECTOR_FIELD.key(): "v",
        }
        remove_dropped_directive_options("b", "BLOB", opts)
        self.assertEqual(opts[CoreOptions.BLOB_FIELD.key()], "a")
        self.assertNotIn(CoreOptions.BLOB_DESCRIPTOR_FIELD.key(), opts)
        self.assertEqual(opts[CoreOptions.VECTOR_FIELD.key()], "v")

    def test_drop_vector(self):
        opts = {
            CoreOptions.VECTOR_FIELD.key(): "emb,emb2",
            "field.emb.vector-dim": "128",
            CoreOptions.BLOB_FIELD.key(): "a",
        }
        remove_dropped_directive_options("emb", "VECTOR", opts)
        self.assertEqual(opts[CoreOptions.VECTOR_FIELD.key()], "emb2")
        self.assertNotIn("field.emb.vector-dim", opts)
        self.assertEqual(opts[CoreOptions.BLOB_FIELD.key()], "a")

    def test_drop_blob_cleans_fallback_keys(self):
        opts = {
            CoreOptions.BLOB_DESCRIPTOR_FIELD.key(): "b,c",
            "blob.stored-descriptor-fields": "b,legacy",
        }
        remove_dropped_directive_options("b", "BLOB", opts)
        self.assertEqual(opts[CoreOptions.BLOB_DESCRIPTOR_FIELD.key()], "c")
        self.assertEqual(opts["blob.stored-descriptor-fields"], "legacy")

    def test_drop_non_directive_is_noop(self):
        opts = {CoreOptions.BLOB_FIELD.key(): "a"}
        remove_dropped_directive_options("x", "INT", opts)
        self.assertEqual(opts[CoreOptions.BLOB_FIELD.key()], "a")


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