#
# 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 __future__ import annotations

from unittest.mock import MagicMock, patch

import pytest

from airflow.sdk import Asset, BaseHook, plugins_manager
from airflow.sdk.lineage import (
    AssetLineageInfo,
    HookLineage,
    HookLineageCollector,
    HookLineageReader,
    NoOpCollector,
)

from tests_common.test_utils.config import conf_vars
from tests_common.test_utils.mock_plugins import mock_plugin_manager


class TestHookLineageCollector:
    @pytest.fixture  # default scope is function
    def collector(self):
        return HookLineageCollector()

    def test_default_max_collected_assets(self):
        """Default MAX_COLLECTED_ASSETS is 100 (configurable via [lineage] conf)."""
        from airflow.sdk.lineage import MAX_COLLECTED_ASSETS

        assert MAX_COLLECTED_ASSETS == 100

    def test_default_max_collected_extra(self):
        """Default MAX_COLLECTED_EXTRA is 200 (configurable via [lineage] conf)."""
        from airflow.sdk.lineage import MAX_COLLECTED_EXTRA

        assert MAX_COLLECTED_EXTRA == 200

    def test_generate_hash_handles_non_serializable(self, collector):
        class Obj:
            def __str__(self):
                return "<obj>"

        h1 = collector._generate_hash({"a": Obj()})
        h2 = collector._generate_hash({"a": "<obj>"})
        assert isinstance(h1, str)
        assert h1 == h2

    def test_generate_hash_is_deterministic(self, collector):
        h1 = collector._generate_hash({"foo": "bar"})
        h2 = collector._generate_hash({"foo": "bar"})
        assert h1 == h2

    def test_generate_hash_changes_with_value(self, collector):
        h1 = collector._generate_hash({"foo": "bar"})
        h2 = collector._generate_hash({"foo": "different"})
        assert h1 != h2

    def test_generate_asset_entry_id_deterministic(self, collector):
        asset = Asset(uri="s3://bucket/file", extra={"x": 1})
        ctx = BaseHook()
        key1 = collector._generate_asset_entry_id(asset, ctx)
        key2 = collector._generate_asset_entry_id(asset, ctx)
        assert key1 == key2

    def test_generate_asset_entry_id_differs_by_context(self, collector):
        asset = Asset(uri="s3://bucket/file")
        ctx1 = BaseHook()
        ctx2 = BaseHook()
        key1 = collector._generate_asset_entry_id(asset, ctx1)
        key2 = collector._generate_asset_entry_id(asset, ctx2)
        assert key1 != key2

    def test_generate_asset_entry_id_differs_by_extra(self, collector):
        ctx = BaseHook()
        asset1 = Asset(uri="s3://bucket/file", extra={"foo": "bar"})
        asset2 = Asset(uri="s3://bucket/file", extra={"foo": "other"})
        key1 = collector._generate_asset_entry_id(asset1, ctx)
        key2 = collector._generate_asset_entry_id(asset2, ctx)
        assert key1 != key2

    def test_generate_extra_entry_id_deterministic(self, collector):
        ctx = BaseHook()
        key1 = collector._generate_extra_entry_id("k", "v", ctx)
        key2 = collector._generate_extra_entry_id("k", "v", ctx)
        assert key1 == key2

    def test_generate_extra_entry_id_differs_by_context(self, collector):
        ctx1 = BaseHook()
        ctx2 = BaseHook()
        key1 = collector._generate_extra_entry_id("k", "v", ctx1)
        key2 = collector._generate_extra_entry_id("k", "v", ctx2)
        assert key1 != key2

    def test_generate_extra_entry_id_differs_by_key_value(self, collector):
        ctx = BaseHook()
        key1 = collector._generate_extra_entry_id("k", "v1", ctx)
        key2 = collector._generate_extra_entry_id("k", "v2", ctx)
        key3 = collector._generate_extra_entry_id("k2", "v1", ctx)
        assert key1 != key2
        assert key1 != key3

    def test_are_assets_collected(self, collector):
        assert collector is not None
        assert collector.collected_assets == HookLineage()
        input_hook = BaseHook()
        output_hook = BaseHook()
        collector.add_input_asset(input_hook, uri="s3://in_bucket/file", name="asset-1", group="test")
        collector.add_output_asset(
            output_hook,
            uri="postgres://example.com:5432/database/default/table",
        )
        assert collector.collected_assets == HookLineage(
            [
                AssetLineageInfo(
                    asset=Asset(uri="s3://in_bucket/file", name="asset-1", group="test"),
                    count=1,
                    context=input_hook,
                )
            ],
            [
                AssetLineageInfo(
                    asset=Asset(
                        uri="postgres://example.com:5432/database/default/table",
                        name="postgres://example.com:5432/database/default/table",
                        group="asset",
                    ),
                    count=1,
                    context=output_hook,
                )
            ],
        )

    @patch("airflow.sdk.lineage.Asset")
    def test_add_input_asset(self, mock_asset, collector):
        asset = MagicMock(spec=Asset, extra={})
        mock_asset.return_value = asset

        mock_hook = MagicMock(spec=BaseHook)
        collector.add_input_asset(mock_hook, uri="test_uri")

        assert next(iter(collector._inputs.values())) == (asset, mock_hook)
        mock_asset.assert_called_once_with(uri="test_uri")

    def test_grouping_assets(self, collector):
        hook_1 = MagicMock(spec=BaseHook)
        hook_2 = MagicMock(spec=BaseHook)

        uri = "test://uri/"

        collector.add_input_asset(context=hook_1, uri=uri)
        collector.add_input_asset(context=hook_2, uri=uri)
        collector.add_input_asset(context=hook_1, uri=uri, asset_extra={"key": "value"})

        collected_inputs = collector.collected_assets.inputs

        assert len(collected_inputs) == 3
        assert collected_inputs[0].asset.uri == "test://uri/"
        assert collected_inputs[0].asset == collected_inputs[1].asset
        assert collected_inputs[0].count == 1
        assert collected_inputs[0].context == collected_inputs[2].context == hook_1
        assert collected_inputs[1].count == 1
        assert collected_inputs[1].context == hook_2
        assert collected_inputs[2].count == 1
        assert collected_inputs[2].asset.extra == {"key": "value"}

    def test_create_asset(self, collector):
        def create_asset(arg1, arg2="default", extra=None):
            return Asset(
                uri=f"myscheme://{arg1}/{arg2}", name=f"asset-{arg1}", group="test", extra=extra or {}
            )

        collector._asset_factories = {"myscheme": create_asset}
        assert collector.create_asset(
            scheme="myscheme",
            uri=None,
            name=None,
            group=None,
            asset_kwargs={"arg1": "value_1"},
            asset_extra=None,
        ) == Asset(uri="myscheme://value_1/default", name="asset-value_1", group="test")
        assert collector.create_asset(
            scheme="myscheme",
            uri=None,
            name=None,
            group=None,
            asset_kwargs={"arg1": "value_1", "arg2": "value_2"},
            asset_extra={"key": "value"},
        ) == Asset(
            uri="myscheme://value_1/value_2", name="asset-value_1", group="test", extra={"key": "value"}
        )

    @patch("airflow.sdk.lineage.ProvidersManagerTaskRuntime")
    def test_create_asset_no_factory(self, mock_providers_manager, collector):
        test_scheme = "myscheme"
        mock_providers_manager.return_value.asset_factories = {}

        test_kwargs = {"arg1": "value_1"}

        assert (
            collector.create_asset(
                scheme=test_scheme,
                uri=None,
                name=None,
                group=None,
                asset_kwargs=test_kwargs,
                asset_extra=None,
            )
            is None
        )

    @patch("airflow.sdk.lineage.ProvidersManagerTaskRuntime")
    def test_create_asset_factory_exception(self, mock_providers_manager, collector):
        def create_asset(extra=None, **kwargs):
            raise RuntimeError("Factory error")

        test_scheme = "myscheme"
        mock_providers_manager.return_value.asset_factories = {test_scheme: create_asset}

        test_kwargs = {"arg1": "value_1"}

        assert (
            collector.create_asset(scheme=test_scheme, uri=None, asset_kwargs=test_kwargs, asset_extra=None)
            is None
        )

    def test_create_asset_missing_parameters_returns_none(self, collector):
        assert collector.create_asset() is None

    def test_collected_assets(self, collector):
        context_input = MagicMock(spec=BaseHook)
        context_output = MagicMock(spec=BaseHook)

        collector.add_input_asset(context_input, uri="test://input")
        collector.add_output_asset(context_output, uri="test://output")

        hook_lineage = collector.collected_assets
        assert len(hook_lineage.inputs) == 1
        assert hook_lineage.inputs[0].asset.uri == "test://input/"
        assert hook_lineage.inputs[0].context == context_input

        assert len(hook_lineage.outputs) == 1
        assert hook_lineage.outputs[0].asset.uri == "test://output/"

    def test_has_collected(self, collector):
        assert not collector.has_collected

        collector._inputs = {"unique_key": (MagicMock(spec=Asset), MagicMock(spec=BaseHook))}
        assert collector.has_collected

    def test_add_asset_count_tracking(self, collector):
        """Test that duplicate assets are counted correctly."""
        ctx = MagicMock(spec=BaseHook)

        # Add same input multiple times
        collector.add_input_asset(ctx, uri="s3://bucket/input")
        collector.add_input_asset(ctx, uri="s3://bucket/input")
        collector.add_input_asset(ctx, uri="s3://bucket/input")

        # Add same output multiple times
        collector.add_output_asset(ctx, uri="s3://bucket/output")
        collector.add_output_asset(ctx, uri="s3://bucket/output")

        assert collector.has_collected

        lineage = collector.collected_assets
        assert len(lineage.inputs) == 1
        assert lineage.inputs[0].asset.uri == "s3://bucket/input"
        assert lineage.inputs[0].count == 3

        assert len(lineage.outputs) == 1
        assert lineage.outputs[0].asset.uri == "s3://bucket/output"
        assert lineage.outputs[0].count == 2

    def test_add_asset_different_uris(self, collector):
        """Test that different URIs are tracked separately."""
        ctx = MagicMock(spec=BaseHook)

        collector.add_input_asset(ctx, uri="s3://bucket/file1")
        collector.add_input_asset(ctx, uri="s3://bucket/file2")
        collector.add_input_asset(ctx, uri="postgres://example.com:5432/database/default/table")

        collector.add_output_asset(ctx, uri="s3://output/file1")
        collector.add_output_asset(ctx, uri="s3://output/file2")

        assert collector.has_collected

        lineage = collector.collected_assets
        assert len(lineage.inputs) == 3
        assert lineage.inputs[0].asset.uri == "s3://bucket/file1"
        assert lineage.inputs[1].asset.uri == "s3://bucket/file2"
        assert lineage.inputs[2].asset.uri == "postgres://example.com:5432/database/default/table"

        assert len(lineage.outputs) == 2
        assert lineage.outputs[0].asset.uri == "s3://output/file1"
        assert lineage.outputs[1].asset.uri == "s3://output/file2"

    def test_add_asset_different_contexts(self, collector):
        """Test that different contexts are tracked separately."""
        ctx1 = MagicMock(spec=BaseHook)
        ctx2 = MagicMock(spec=BaseHook)

        collector.add_input_asset(ctx1, uri="s3://bucket/file")
        collector.add_input_asset(ctx2, uri="s3://bucket/file")

        collector.add_output_asset(ctx1, uri="s3://output/file")
        collector.add_output_asset(ctx2, uri="s3://output/file")

        assert collector.has_collected

        lineage = collector.collected_assets
        assert len(lineage.inputs) == 2
        assert lineage.inputs[0].context == ctx1
        assert lineage.inputs[1].context == ctx2

        assert len(lineage.outputs) == 2
        assert lineage.outputs[0].context == ctx1
        assert lineage.outputs[1].context == ctx2

    def test_add_asset_with_extra_metadata(self, collector):
        """Test adding assets with extra metadata."""
        ctx = MagicMock(spec=BaseHook)

        collector.add_input_asset(
            ctx,
            uri="postgres://example.com:5432/database/default/table",
            asset_extra={"schema": "public", "table": "users"},
        )
        collector.add_output_asset(
            ctx,
            uri="postgres://example.com:5432/database/default/table",
            asset_extra={"schema": "public", "table": "results"},
        )

        assert collector.has_collected

        lineage = collector.collected_assets
        assert len(lineage.inputs) == 1
        assert lineage.inputs[0].asset.extra == {"schema": "public", "table": "users"}

        assert len(lineage.outputs) == 1
        assert lineage.outputs[0].asset.extra == {"schema": "public", "table": "results"}

    def test_add_asset_different_extra_values(self, collector):
        """Test that assets with different extra values are tracked separately."""
        ctx = MagicMock(spec=BaseHook)

        # Same URI but different extra metadata
        collector.add_input_asset(ctx, uri="s3://bucket/file", asset_extra={"version": "1"})
        collector.add_input_asset(ctx, uri="s3://bucket/file", asset_extra={"version": "2"})

        collector.add_output_asset(ctx, uri="s3://output/file", asset_extra={"format": "parquet"})
        collector.add_output_asset(ctx, uri="s3://output/file", asset_extra={"format": "csv"})

        assert collector.has_collected

        lineage = collector.collected_assets
        assert len(lineage.inputs) == 2
        assert lineage.inputs[0].asset.extra == {"version": "1"}
        assert lineage.inputs[1].asset.extra == {"version": "2"}

        assert len(lineage.outputs) == 2
        assert lineage.outputs[0].asset.extra == {"format": "parquet"}
        assert lineage.outputs[1].asset.extra == {"format": "csv"}

    def test_hooks_limit_input_output_assets(self, collector):
        assert not collector.has_collected

        for i in range(1000):
            collector.add_input_asset(MagicMock(spec=BaseHook), uri=f"test://input/{i}")
            collector.add_output_asset(MagicMock(spec=BaseHook), uri=f"test://output/{i}")

        assert collector.has_collected
        assert len(collector._inputs) == 100
        assert len(collector._outputs) == 100

    def test_configurable_max_collected_assets(self):
        """MAX_COLLECTED_ASSETS is read from [lineage] conf and limits both inputs and outputs."""
        import importlib

        import airflow.sdk.lineage as lineage_mod

        new_max = 4

        with conf_vars({("lineage", "max_assets_per_collector"): str(new_max)}):
            importlib.reload(lineage_mod)
            assert new_max == lineage_mod.MAX_COLLECTED_ASSETS

            collector = lineage_mod.HookLineageCollector()
            for i in range(new_max * 2):
                collector.add_input_asset(MagicMock(spec=BaseHook), uri=f"test://input/{i}")
                collector.add_output_asset(MagicMock(spec=BaseHook), uri=f"test://output/{i}")
            assert len(collector._inputs) == new_max
            assert len(collector._outputs) == new_max

        importlib.reload(lineage_mod)

    @pytest.mark.parametrize("uri", ["", None])
    def test_invalid_uri_none(self, collector, uri):
        """Test handling of None or empty URI - should not raise."""
        ctx = MagicMock(spec=BaseHook)

        # Should not raise exceptions
        collector.add_input_asset(ctx, uri=uri)
        collector.add_output_asset(ctx, uri=uri)

        # Collector should handle gracefully and not collect invalid URIs
        assert not collector.has_collected

        lineage = collector.collected_assets
        assert len(lineage.inputs) == 0
        assert len(lineage.outputs) == 0

    def test_malformed_uri(self, collector):
        """Test handling of malformed URIs - should not raise."""
        ctx = MagicMock(spec=BaseHook)

        # Various malformed URIs should not cause crashes
        collector.add_input_asset(ctx, uri="not-a-valid-uri")
        collector.add_input_asset(ctx, uri="://missing-scheme")
        collector.add_input_asset(ctx, uri="scheme:")
        collector.add_output_asset(ctx, uri="//no-scheme")

        assert collector.has_collected

        lineage = collector.collected_assets
        assert len(lineage.inputs) == 3
        assert lineage.inputs[0].asset.uri == "not-a-valid-uri"
        assert lineage.inputs[1].asset.uri == "://missing-scheme"
        assert lineage.inputs[2].asset.uri == "scheme:/"

        assert len(lineage.outputs) == 1
        assert lineage.outputs[0].asset.uri == "//no-scheme"

    def test_very_long_uri(self, collector):
        """Test handling of very long URIs - 1000 chars OK, 2000 chars raises ValueError."""
        ctx = MagicMock(spec=BaseHook)

        # Create very long URI (1000 chars - should work)
        long_path = "a" * 1000
        long_uri = f"s3://bucket/{long_path}"

        # Create too long URI (2000 chars - should raise)
        too_long_uri = f"s3://bucket/{long_path * 2}"

        collector.add_input_asset(ctx, uri=long_uri)

        # Too long URI should raise ValueError
        with pytest.raises(ValueError, match="Asset name cannot exceed"):
            collector.add_output_asset(ctx, uri=too_long_uri)

        assert collector.has_collected

        lineage = collector.collected_assets
        assert len(lineage.inputs) == 1
        assert lineage.inputs[0].asset.uri == long_uri

        assert len(lineage.outputs) == 0

    def test_uri_with_special_characters(self, collector):
        """Test URIs with special characters - should not raise."""
        ctx = MagicMock(spec=BaseHook)

        # URIs with various special characters
        special_uris = {
            "s3://bucket/path with spaces/file": "s3://bucket/path with spaces/file",
            "s3://bucket/path%20encoded/file": "s3://bucket/path%20encoded/file",
            "file:///path/with/üñíçødé/chars": "file:///path/with/üñíçødé/chars",
            "scheme://host/path?query=value&other=123": "scheme://host/path?other=123&query=value",
            "scheme://host/path#fragment": "scheme://host/path",
            "postgres://user:p@ss!word@host:5432/db/sche$ma/table": "postgres://host:5432/db/sche$ma/table",
        }

        for uri in special_uris:
            collector.add_input_asset(ctx, uri=uri)

        assert collector.has_collected

        lineage = collector.collected_assets
        assert len(lineage.inputs) == 6

        for i, expected_uri in enumerate(special_uris.values()):
            assert lineage.inputs[i].asset.uri == expected_uri

    def test_empty_asset_extra(self, collector):
        """Test that empty asset_extra is handled correctly."""
        ctx = MagicMock(spec=BaseHook)

        collector.add_input_asset(ctx, uri="s3://bucket/file", asset_extra={})
        collector.add_output_asset(ctx, uri="s3://bucket/file", asset_extra={})

        assert collector.has_collected

        lineage = collector.collected_assets
        assert len(lineage.inputs) == 1
        assert lineage.inputs[0].asset.extra == {}
        assert len(lineage.outputs) == 1
        assert lineage.outputs[0].asset.extra == {}

    def test_asset_with_all_optional_parameters(self, collector):
        """Test asset creation with all optional parameters provided."""
        ctx = MagicMock(spec=BaseHook)

        collector.add_input_asset(
            ctx,
            uri="s3://bucket/file",
            name="custom-name",
            group="custom-group",
            asset_extra={"key": "value"},
        )

        assert collector.has_collected

        lineage = collector.collected_assets
        assert len(lineage.inputs) == 1
        assert lineage.inputs[0].asset.uri == "s3://bucket/file"
        assert lineage.inputs[0].asset.name == "custom-name"
        assert lineage.inputs[0].asset.group == "custom-group"
        assert lineage.inputs[0].asset.extra == {"key": "value"}

    def test_asset_extra_with_non_serializable(self, collector):
        """Test that asset_extra with non-JSON-serializable values is handled."""
        ctx = MagicMock(spec=BaseHook)

        class CustomObject:
            pass

        # Should not raise
        collector.add_input_asset(ctx, uri="s3://bucket/file", asset_extra={"obj": CustomObject()})

        # May or may not be collected depending on implementation
        lineage = collector.collected_assets
        # Just verify it doesn't crash and structure is intact
        assert isinstance(lineage.inputs, list)
        assert isinstance(lineage.outputs, list)

    def test_empty_name_and_group(self, collector):
        """Test that empty strings for name and group are handled."""
        ctx = MagicMock(spec=BaseHook)

        # Empty strings for optional parameters
        collector.add_input_asset(ctx, uri="s3://bucket/file", name="", group="")

        assert collector.has_collected

        lineage = collector.collected_assets
        assert len(lineage.inputs) == 1
        assert lineage.inputs[0].asset.uri == "s3://bucket/file"
        assert lineage.inputs[0].asset.name == "s3://bucket/file"
        assert lineage.inputs[0].asset.group == "asset"

    def test_add_extra(self, collector):
        ctx = MagicMock(spec=BaseHook)
        collector.add_extra(ctx, "k", "v")

        data = collector.collected_assets.extra
        assert len(data) == 1
        assert data[0].key == "k"
        assert data[0].value == "v"
        assert data[0].context == ctx
        assert data[0].count == 1

        # adding again with same values only increments count
        collector.add_extra(ctx, "k", "v")
        assert collector.collected_assets.extra[0].count == 2
        data = collector.collected_assets.extra
        assert len(data) == 1

    def test_add_extra_missing_key_or_value(self, collector):
        ctx = MagicMock(spec=BaseHook)

        collector.add_extra(ctx, "", "v")
        collector.add_extra(ctx, "k", None)

        # nothing added
        assert len(collector.collected_assets.extra) == 0

    def test_extra_limit(self, collector):
        ctx = MagicMock(spec=BaseHook)

        for i in range(501):
            collector.add_extra(ctx, f"k{i}", f"v{i}")

        assert len(collector.collected_assets.extra) == 200

    def test_configurable_max_collected_extra(self):
        """MAX_COLLECTED_EXTRA is read from [lineage] conf and limits extra collection."""
        import importlib

        import airflow.sdk.lineage as lineage_mod

        new_max = 4

        with conf_vars({("lineage", "max_extras_per_collector"): str(new_max)}):
            importlib.reload(lineage_mod)
            assert new_max == lineage_mod.MAX_COLLECTED_EXTRA

            collector = lineage_mod.HookLineageCollector()
            ctx = MagicMock(spec=BaseHook)
            for i in range(new_max * 2):
                collector.add_extra(ctx, f"key_{i}", f"value_{i}")
            assert len(collector.collected_assets.extra) == new_max

        importlib.reload(lineage_mod)

    def test_add_extra_different_values(self, collector):
        """Test that different values are tracked separately."""
        ctx = MagicMock(spec=BaseHook)

        collector.add_extra(ctx, "key1", {"data": "value1"})
        collector.add_extra(ctx, "key2", {"data": "value2"})
        collector.add_extra(ctx, "key1", {"data": "value3"})

        assert collector.has_collected

        lineage = collector.collected_assets
        assert len(lineage.extra) == 3
        assert lineage.extra[0].key == "key1"
        assert lineage.extra[0].value == {"data": "value1"}
        assert lineage.extra[0].count == 1
        assert lineage.extra[1].key == "key2"
        assert lineage.extra[1].value == {"data": "value2"}
        assert lineage.extra[1].count == 1
        assert lineage.extra[2].key == "key1"
        assert lineage.extra[2].value == {"data": "value3"}
        assert lineage.extra[2].count == 1

    def test_add_extra_different_contexts(self, collector):
        """Test that different contexts are tracked separately."""
        ctx1 = MagicMock(spec=BaseHook)
        ctx2 = MagicMock(spec=BaseHook)

        collector.add_extra(ctx1, "test_key", {"data": "value"})
        collector.add_extra(ctx2, "test_key", {"data": "value"})

        assert collector.has_collected

        lineage = collector.collected_assets
        assert len(lineage.extra) == 2
        assert lineage.extra[0].context == ctx1
        assert lineage.extra[1].context == ctx2

    def test_add_extra_complex_values(self, collector):
        """Test that add_extra handles complex JSON-serializable values."""
        ctx = MagicMock(spec=BaseHook)

        collector.add_extra(ctx, "dict", {"nested": {"data": "value"}})
        collector.add_extra(ctx, "list", [1, 2, 3, "test"])
        collector.add_extra(ctx, "number", 42)
        collector.add_extra(ctx, "string", "simple string")

        assert collector.has_collected

        lineage = collector.collected_assets
        assert len(lineage.extra) == 4
        assert lineage.extra[0].value == {"nested": {"data": "value"}}
        assert lineage.extra[1].value == [1, 2, 3, "test"]
        assert lineage.extra[2].value == 42
        assert lineage.extra[3].value == "simple string"

    def test_special_characters_in_extra_key(self, collector):
        """Test that extra keys with special characters work."""
        ctx = MagicMock(spec=BaseHook)

        collector.add_extra(ctx, "key-with-dashes", {"data": "value"})
        collector.add_extra(ctx, "key.with.dots", {"data": "value"})
        collector.add_extra(ctx, "key_with_underscores", {"data": "value"})
        collector.add_extra(ctx, "key/with/slashes", {"data": "value"})
        collector.add_extra(ctx, "key:with:colons", {"data": "value"})

        assert collector.has_collected

        lineage = collector.collected_assets
        assert len(lineage.extra) == 5
        assert lineage.extra[0].key == "key-with-dashes"
        assert lineage.extra[1].key == "key.with.dots"
        assert lineage.extra[2].key == "key_with_underscores"
        assert lineage.extra[3].key == "key/with/slashes"
        assert lineage.extra[4].key == "key:with:colons"

    def test_unicode_in_extra_key_and_value(self, collector):
        """Test that unicode characters in extra work correctly."""
        ctx = MagicMock(spec=BaseHook)

        collector.add_extra(ctx, "clé_française", {"données": "valeur"})
        collector.add_extra(ctx, "中文键", {"中文": "值"})
        collector.add_extra(ctx, "مفتاح", {"بيانات": "قيمة"})
        collector.add_extra(ctx, "emoji_🚀", {"status": "✅"})

        assert collector.has_collected

        lineage = collector.collected_assets
        assert len(lineage.extra) == 4
        assert lineage.extra[0].key == "clé_française"
        assert lineage.extra[0].value == {"données": "valeur"}
        assert lineage.extra[1].key == "中文键"
        assert lineage.extra[1].value == {"中文": "值"}
        assert lineage.extra[2].key == "مفتاح"
        assert lineage.extra[2].value == {"بيانات": "قيمة"}
        assert lineage.extra[3].key == "emoji_🚀"
        assert lineage.extra[3].value == {"status": "✅"}

    def test_very_large_extra_value(self, collector):
        """Test that large extra values are handled."""
        ctx = MagicMock(spec=BaseHook)

        # Create a large value
        large_value = {"data": "x" * 1000, "list": list(range(1000))}

        collector.add_extra(ctx, "large_key", large_value)

        assert collector.has_collected

        lineage = collector.collected_assets
        assert len(lineage.extra) == 1
        assert lineage.extra[0].key == "large_key"
        assert lineage.extra[0].value == large_value

    def test_deeply_nested_extra_value(self, collector):
        """Test that deeply nested data structures in extra are handled."""
        ctx = MagicMock(spec=BaseHook)

        # Create deeply nested structure
        nested_value = {"level1": {"level2": {"level3": {"level4": {"level5": {"data": "deep"}}}}}}

        collector.add_extra(ctx, "nested", nested_value)

        assert collector.has_collected

        lineage = collector.collected_assets
        assert len(lineage.extra) == 1
        assert lineage.extra[0].value == nested_value

    def test_extra_value_with_various_types(self, collector):
        """Test that extra can handle various data types."""
        ctx = MagicMock(spec=BaseHook)

        collector.add_extra(ctx, "string", "text")
        collector.add_extra(ctx, "integer", 42)
        collector.add_extra(ctx, "float", 3.14)
        collector.add_extra(ctx, "boolean", True)
        collector.add_extra(ctx, "list", [1, 2, 3])
        collector.add_extra(ctx, "dict", {"key": "value"})
        collector.add_extra(ctx, "null", None)

        assert collector.has_collected

        # None value should not be collected (based on validation)
        lineage = collector.collected_assets
        assert len(lineage.extra) == 6  # None is filtered out

        assert lineage.extra[0].value == "text"
        assert lineage.extra[1].value == 42
        assert lineage.extra[2].value == 3.14
        assert lineage.extra[3].value is True
        assert lineage.extra[4].value == [1, 2, 3]
        assert lineage.extra[5].value == {"key": "value"}

    def test_non_json_serializable_value_in_extra(self, collector):
        """Test that non-JSON-serializable values are handled gracefully."""
        ctx = MagicMock(spec=BaseHook)

        # Create a non-serializable object
        class CustomObject:
            def __str__(self):
                return "custom_object"

        # Should not raise - collector should handle via str conversion or skip
        collector.add_extra(ctx, "custom_key", CustomObject())

        # May or may not be collected depending on implementation
        lineage = collector.collected_assets
        # Just verify it doesn't crash
        assert isinstance(lineage.extra, list)

    def test_extremely_long_extra_key(self, collector):
        """Test that extremely long extra keys are handled."""
        ctx = MagicMock(spec=BaseHook)

        long_key = "k" * 1000
        collector.add_extra(ctx, long_key, "value")

        assert collector.has_collected

        lineage = collector.collected_assets
        assert len(lineage.extra) == 1
        assert lineage.extra[0].key == long_key
        assert lineage.extra[0].value == "value"

    def test_collected_assets_called_multiple_times(self, collector):
        """Test that collected_assets property can be called multiple times."""
        ctx = MagicMock(spec=BaseHook)

        collector.add_input_asset(ctx, uri="s3://bucket/file")

        assert collector.has_collected

        # Call multiple times - should return same data
        lineage1 = collector.collected_assets
        lineage2 = collector.collected_assets
        lineage3 = collector.collected_assets

        assert lineage1.inputs == lineage2.inputs == lineage3.inputs
        assert len(lineage1.inputs) == 1

    def test_has_collected_only_extra(self, collector):
        assert collector.has_collected is False

        collector.add_extra(MagicMock(spec=BaseHook), "event", "trigger")

        assert collector.has_collected is True
        assert len(collector.collected_assets.inputs) == 0
        assert len(collector.collected_assets.outputs) == 0
        assert len(collector.collected_assets.extra) == 1

    def test_none_context(self, collector):
        """Test handling of None context - should not raise."""
        # Should not raise exceptions
        collector.add_input_asset(None, uri="s3://bucket/input")
        collector.add_output_asset(None, uri="s3://bucket/output")
        collector.add_extra(None, "key", "value")

        assert collector.has_collected

        lineage = collector.collected_assets
        assert len(lineage.inputs) == 1
        assert lineage.inputs[0].context is None

        assert len(lineage.outputs) == 1
        assert lineage.outputs[0].context is None

        assert len(lineage.extra) == 1
        assert lineage.extra[0].context is None

    def test_rapid_repeated_calls(self, collector):
        """Test that rapid repeated calls don't cause issues."""
        ctx = MagicMock(spec=BaseHook)

        # Simulate rapid repeated calls
        for _ in range(50):
            collector.add_input_asset(ctx, uri="s3://bucket/file")
            collector.add_output_asset(ctx, uri="s3://bucket/output")
            collector.add_extra(ctx, "key", "value")

        assert collector.has_collected

        lineage = collector.collected_assets
        # Should have counted properly
        assert len(lineage.inputs) == 1
        assert lineage.inputs[0].count == 50
        assert len(lineage.outputs) == 1
        assert lineage.outputs[0].count == 50
        assert len(lineage.extra) == 1
        assert lineage.extra[0].count == 50

    def test_mixed_valid_invalid_operations(self, collector):
        """Test mixing valid and invalid operations."""
        ctx = MagicMock(spec=BaseHook)

        # Mix valid and invalid calls
        collector.add_input_asset(ctx, uri="s3://bucket/valid")
        collector.add_input_asset(ctx, uri=None)  # Invalid - should not be collected
        collector.add_input_asset(ctx, uri="")  # Invalid - should not be collected
        collector.add_input_asset(ctx, uri="s3://bucket/another-valid")

        collector.add_extra(ctx, "valid_key", "valid_value")
        collector.add_extra(ctx, "", "invalid_key")  # Invalid key - should not be collected
        collector.add_extra(ctx, "another_key", "another_value")

        assert collector.has_collected

        # Should collect only valid items
        lineage = collector.collected_assets
        assert len(lineage.inputs) == 2
        assert lineage.inputs[0].asset.uri == "s3://bucket/valid"
        assert lineage.inputs[1].asset.uri == "s3://bucket/another-valid"

        assert len(lineage.extra) == 2
        assert lineage.extra[0].key == "valid_key"
        assert lineage.extra[0].value == "valid_value"
        assert lineage.extra[1].key == "another_key"
        assert lineage.extra[1].value == "another_value"

    def test_noop_collector(self):
        noop = NoOpCollector()
        ctx = MagicMock(spec=BaseHook)
        noop.add_input_asset(ctx, uri="x")
        noop.add_output_asset(ctx, uri="y")
        noop.add_extra(ctx, "k", "v")

        lineage = noop.collected_assets
        assert lineage.inputs == []
        assert lineage.outputs == []
        assert lineage.extra == []


class FakePlugin(plugins_manager.AirflowPlugin):
    name = "FakePluginHavingHookLineageCollector"
    hook_lineage_readers = [HookLineageReader]


@pytest.mark.parametrize(
    ("has_readers", "expected_class_name"),
    [
        (True, "HookLineageCollector"),
        (False, "NoOpCollector"),
    ],
)
def test_get_hook_lineage_collector(has_readers, expected_class_name):
    # After importlib.reload in other tests, the top-level imported references
    # (get_hook_lineage_collector, HookLineageCollector, NoOpCollector) may point
    # to stale class/function objects. Always fetch from the live module so that
    # the function, the returned instance, and the class used in isinstance() all
    # belong to the same reload generation.
    import airflow.sdk.lineage as lineage_mod

    lineage_mod.get_hook_lineage_collector.cache_clear()
    expected_class = getattr(lineage_mod, expected_class_name)
    plugins = [FakePlugin()] if has_readers else []
    with mock_plugin_manager(plugins=plugins):
        assert isinstance(lineage_mod.get_hook_lineage_collector(), expected_class)
        assert lineage_mod.get_hook_lineage_collector() is lineage_mod.get_hook_lineage_collector()
