"""Resolving which databases an operation covers."""

import pytest
from pydantic import ValidationError

from haiku.rag.client import HaikuRAG
from haiku.rag.client.scope import DatabaseScope
from haiku.rag.config.models import AppConfig, LanceDBConfig
from haiku.rag.store.exceptions import (
    AmbiguousDatabaseError,
    SourceUnavailableError,
    UnknownDatabaseError,
)
from haiku.rag.utils import locate_database
from tests.multi_db.helpers import (
    _config,
    _seed,
)


class TestConfig:
    def test_databases_is_the_one_placement(self):
        config = LanceDBConfig(databases={"one": "s3://b/one.lancedb"})
        assert config.databases == {"one": "s3://b/one.lancedb"}

    def test_uri_is_refused_naming_the_replacement(self):
        with pytest.raises(ValidationError, match="lancedb.databases"):
            LanceDBConfig.model_validate({"uri": "s3://b/one.lancedb"})


class TestNamingIsRequired:
    def test_a_blank_name_is_rejected(self):
        """An unnamed database is unreachable: every source check reads the
        empty name as no name at all."""
        with pytest.raises(ValidationError, match="entry with no name"):
            LanceDBConfig(databases={"": "/tmp/a.lancedb"})
        with pytest.raises(ValidationError, match="entry with no name"):
            LanceDBConfig(databases={"   ": "/tmp/a.lancedb"})

    def test_a_blank_location_is_rejected(self):
        """A blank location resolves to the working directory."""
        with pytest.raises(
            ValidationError, match=r"databases\[alpha\] has no location"
        ):
            LanceDBConfig(databases={"alpha": ""})


class TestNamingADatabaseDirectly:
    @pytest.mark.asyncio
    async def test_a_db_path_beside_the_configured_set_is_refused(
        self, tmp_path, temp_db_path
    ):
        """The configuration places databases; a path beside it is a second
        placement, and the refusal names both."""
        config = _config(tmp_path, ["alpha", "beta"])

        with pytest.raises(AmbiguousDatabaseError, match="alpha") as raised:
            async with HaikuRAG(temp_db_path, config=config, create=True):
                pass
        assert str(temp_db_path) in str(raised.value)
        assert not temp_db_path.exists()

    @pytest.mark.asyncio
    async def test_one_configured_database_is_opened_by_name(self, tmp_path):
        """A set of one is not federated, and the client resolves it."""
        config = _config(tmp_path, ["alpha"])
        await _seed(config, "alpha", ["alpha document about cats"])

        async with HaikuRAG(config=config) as rag:
            assert not rag.covers_multiple
            assert rag.source == "alpha"
            results = await rag.search("cats", search_type="fts", limit=10)

        assert [r.source for r in results] == ["alpha"]


class TestOneConfiguredLocation:
    """One entry in `lancedb.databases` places one named database, at a URI or
    at a local path."""

    def _config(self, location) -> AppConfig:
        return AppConfig(lancedb=LanceDBConfig(databases={"notes": str(location)}))

    @pytest.mark.asyncio
    async def test_a_local_location_opens_the_configured_database(self, tmp_path):
        located = tmp_path / "notes.lancedb"
        config = self._config(located)

        async with HaikuRAG(config=config, create=True) as rag:
            assert rag.store.db_path == located
            assert rag.source == "notes"
        assert located.exists()

    @pytest.mark.asyncio
    async def test_a_path_beside_the_configured_database_is_refused(self, tmp_path):
        config = self._config(tmp_path / "configured.lancedb")
        chosen = tmp_path / "chosen.lancedb"

        with pytest.raises(AmbiguousDatabaseError, match="notes"):
            async with HaikuRAG(chosen, config=config, create=True):
                pass
        assert not chosen.exists()
        assert not (tmp_path / "configured.lancedb").exists()

    @pytest.mark.asyncio
    async def test_a_local_location_that_does_not_exist_is_refused(self, tmp_path):
        """A schemeless location is a local path and must exist. The error names
        the configured database, never its location."""
        config = self._config(tmp_path / "typo.lancedb")

        with pytest.raises(SourceUnavailableError, match="notes") as caught:
            async with HaikuRAG(config=config):
                pass
        assert "typo.lancedb" not in str(caught.value)
        assert not (tmp_path / "typo.lancedb").exists()

    def test_a_uri_with_a_scheme_stays_a_uri(self, tmp_path):
        """Object storage has no local path to check, and a location that does
        not exist yet is normal there."""
        from haiku.rag.store.engine import ConnectionMode

        config = self._config("s3://bucket/one.lancedb")

        [ref] = DatabaseScope.resolve(config).databases

        assert ref.location == "s3://bucket/one.lancedb"
        assert ConnectionMode.of(ref.location) == ConnectionMode.OBJECT_STORAGE


class TestSessionsOwnTheRef:
    """A session is built from the resolved reference and hands storage only
    its location; the configuration it keeps is the one the caller named."""

    @pytest.mark.asyncio
    async def test_a_session_opens_the_location_with_the_undivided_config(
        self, tmp_path
    ):
        from haiku.rag.client.session import SingleDatabaseSession

        config = _config(tmp_path, ["alpha", "beta"])
        await _seed(config, "alpha", ["alpha document about cats"])
        [ref] = DatabaseScope.resolve(config, database_name="alpha").databases

        session = await SingleDatabaseSession(ref, config, read_only=True).open()
        try:
            assert session.source == "alpha"
            assert session.location == ref.location
            assert session.db_path == ref.location
            assert session.store.location == ref.location
            assert session.store._config is config
        finally:
            await session.aclose()

    @pytest.mark.asyncio
    async def test_a_client_keeps_the_configuration_it_was_given(self, tmp_path):
        config = _config(tmp_path, ["alpha", "beta"])
        await _seed(config, "alpha", ["alpha document about cats"])

        async with HaikuRAG(config=config, sources=["alpha"]) as rag:
            assert rag._config is config
            assert set(rag._config.lancedb.databases) == {"alpha", "beta"}
            assert rag.store.location == tmp_path / "alpha.lancedb"


class TestLocate:
    def test_a_scheme_is_a_uri(self):
        assert locate_database("s3://bucket/one.lancedb") == "s3://bucket/one.lancedb"

    def test_anything_else_is_a_local_path(self):
        from pathlib import Path

        assert locate_database("/data/one.lancedb") == Path("/data/one.lancedb")


class TestSelection:
    @pytest.mark.asyncio
    async def test_unknown_source_at_construction_is_rejected(self, tmp_path):
        config = _config(tmp_path, ["alpha", "beta"])

        with pytest.raises(UnknownDatabaseError, match="nope"):
            async with HaikuRAG(config=config, sources=["nope"]):
                pass

    @pytest.mark.asyncio
    async def test_unknown_source_across_several_databases_is_rejected(self, tmp_path):
        config = _config(tmp_path, ["alpha", "beta"])
        await _seed(config, "alpha", ["alpha document about cats"])
        await _seed(config, "beta", ["beta document about cats"])

        async with HaikuRAG(config=config) as rag:
            with pytest.raises(UnknownDatabaseError, match="nope"):
                await rag.search("cats", search_type="fts", sources=["nope"])

    @pytest.mark.asyncio
    async def test_no_matches_anywhere_returns_nothing(self, tmp_path):
        config = _config(tmp_path, ["alpha", "beta"])
        await _seed(config, "alpha", ["alpha document about cats"])
        await _seed(config, "beta", ["beta document about cats"])

        async with HaikuRAG(config=config) as rag:
            assert await rag.search("aardvarks", search_type="fts") == []


class TestPlacingADatabase:
    """What a client says about the databases it covers, so nothing outside has
    to read its private state to find out."""

    @pytest.mark.asyncio
    async def test_a_set_names_every_database_it_covers(self, tmp_path):
        config = _config(tmp_path, ["alpha", "beta"])
        await _seed(config, "alpha", ["alpha one"])
        await _seed(config, "beta", ["beta one"])

        async with HaikuRAG(config=config, read_only=True) as rag:
            assert rag.covers_multiple
            assert rag.source_names == ("alpha", "beta")
            assert rag.source is None

    @pytest.mark.asyncio
    async def test_one_named_database_names_itself(self, tmp_path):
        config = _config(tmp_path, ["alpha", "beta"])
        await _seed(config, "alpha", ["alpha one"])

        async with HaikuRAG(config=config, read_only=True, sources=["alpha"]) as rag:
            assert not rag.covers_multiple
            assert rag.source_names == ("alpha",)
            assert rag.source == "alpha"

    @pytest.mark.asyncio
    async def test_a_named_database_keeps_its_name_on_re_entry(self, tmp_path):
        """Entering derives a single-database configuration from what was
        configured. Deriving it from the last derivation loses the name."""
        config = _config(tmp_path, ["alpha", "beta"])
        await _seed(config, "alpha", ["alpha document about cats"])

        rag = HaikuRAG(config=config, read_only=True, sources=["alpha"])
        async with rag:
            assert rag.source == "alpha"
        async with rag:
            assert rag.source == "alpha"
            assert rag.source_names == ("alpha",)
            results = await rag.search("cats", search_type="fts")

        assert {r.source for r in results} == {"alpha"}

    @pytest.mark.asyncio
    async def test_a_database_at_a_path_is_named_by_its_stem(self, temp_db_path):
        async with HaikuRAG(temp_db_path, create=True) as rag:
            assert rag.source_names == (temp_db_path.stem,)
            assert rag.source == temp_db_path.stem

    @pytest.mark.asyncio
    async def test_the_default_database_is_selectable_by_name(self, tmp_path):
        """Nothing configured is the one entry `haiku.rag`, an ordinary
        configured database that `sources` can name."""
        from haiku.rag.config.models import StorageConfig

        config = AppConfig(storage=StorageConfig(data_dir=tmp_path))

        async with HaikuRAG(config=config, sources=["haiku.rag"], create=True) as rag:
            assert rag.source == "haiku.rag"
            assert rag.store.db_path == tmp_path / "haiku.rag.lancedb"

    def test_coverage_is_known_before_the_client_opens(self, tmp_path):
        """Coverage is a fact of the resolved scope, readable before entering,
        and `source_names` and `covers_multiple` agree on it."""
        config = _config(tmp_path, ["alpha", "beta"])

        covering = HaikuRAG(config=config)
        assert covering.source_names == ("alpha", "beta")
        assert covering.covers_multiple

        narrowed = HaikuRAG(config=config, sources=["beta"])
        assert narrowed.source_names == ("beta",)
        assert not narrowed.covers_multiple

        at_path = HaikuRAG(tmp_path / "other.lancedb")
        assert at_path.source_names == ("other",)
        assert not at_path.covers_multiple

    @pytest.mark.asyncio
    async def test_the_reader_for_a_database_is_the_client_holding_it(self, tmp_path):
        config = _config(tmp_path, ["alpha", "beta"])
        await _seed(config, "alpha", ["alpha one"])
        await _seed(config, "beta", ["beta one"])

        async with HaikuRAG(config=config, read_only=True) as rag:
            reader = await rag.reader_for("beta")

            assert reader is not None
            assert reader.source == "beta"
            # Asked twice, the same wrapper comes back.
            assert await rag.reader_for("beta") is reader

    @pytest.mark.asyncio
    async def test_a_client_reading_one_database_is_its_own_reader(self, temp_db_path):
        async with HaikuRAG(temp_db_path, create=True) as rag:
            assert await rag.reader_for(None) is rag

    def test_one_type_answers_for_a_name_nothing_covers(self, tmp_path):
        """Selecting by name is a lookup wherever it happens, so it fails the
        same way at construction, per query, and when placing evidence."""
        config = _config(tmp_path, ["alpha", "beta"])

        with pytest.raises(UnknownDatabaseError):
            DatabaseScope.resolve(config, database_name="typo")
        with pytest.raises(UnknownDatabaseError):
            DatabaseScope.resolve(config).select(["typo"])

        # A KeyError, so a caller treating selection as a lookup still catches it.
        assert issubclass(UnknownDatabaseError, KeyError)
        # ...but the message reads as a sentence, not as a missing key.
        assert str(UnknownDatabaseError("unknown database 'typo'")) == (
            "unknown database 'typo'"
        )

    def test_a_path_and_sources_cannot_both_choose(self, tmp_path):
        """A path names one database and `sources` names others; together they
        are refused, whatever the selection."""
        config = _config(tmp_path, ["alpha", "beta"])

        for sources in ([], ["alpha"], ["nope"]):
            with pytest.raises(AmbiguousDatabaseError, match="pass one of them"):
                HaikuRAG(tmp_path / "alpha.lancedb", config=config, sources=sources)

    @pytest.mark.asyncio
    async def test_one_database_refuses_a_name_it_does_not_cover(self, tmp_path):
        """A citation naming another database must not get this database's
        reader."""
        config = _config(tmp_path, ["alpha", "beta"])
        await _seed(config, "alpha", ["alpha one"])
        await _seed(config, "beta", ["beta one"])

        async with HaikuRAG(config=config, sources=["alpha"]) as alpha:
            assert await alpha.reader_for("alpha") is alpha
            with pytest.raises(UnknownDatabaseError, match="beta"):
                await alpha.reader_for("beta")

    @pytest.mark.asyncio
    async def test_a_database_at_a_path_answers_to_its_stem_alone(self, temp_db_path):
        async with HaikuRAG(temp_db_path, create=True) as rag:
            assert await rag.reader_for(temp_db_path.stem) is rag
            with pytest.raises(UnknownDatabaseError, match=temp_db_path.stem):
                await rag.reader_for("anything")

    @pytest.mark.asyncio
    async def test_a_set_cannot_place_evidence_that_names_no_database(self, tmp_path):
        """A sourceless citation names no database a set could place."""
        config = _config(tmp_path, ["alpha", "beta"])
        await _seed(config, "alpha", ["alpha one"])

        async with HaikuRAG(config=config, read_only=True) as rag:
            assert await rag.reader_for(None) is None


class TestNamingOneOfTheSetOnTheCommandLine:
    """`--db-name NAME` reaches the application layer as a name, and every
    client it opens has to honour it — one that ignores it covers the set and
    quietly answers from the wrong database."""

    @pytest.mark.asyncio
    async def test_a_named_database_is_the_one_read(self, tmp_path, capsys):
        from haiku.rag.app import HaikuRAGApp

        config = _config(tmp_path, ["alpha", "beta"])
        await _seed(config, "alpha", ["alpha document about cats"])
        await _seed(config, "beta", ["beta document about cats"])

        scope = DatabaseScope.resolve(config).select(["beta"])
        app = HaikuRAGApp(scope=scope, config=config, read_only=True)
        await app.list_documents()

        # Rich wraps long lines, so match the unwrapped part of the URI.
        printed = capsys.readouterr().out
        assert "test://beta/" in printed
        assert "test://alpha/" not in printed

    @pytest.mark.asyncio
    async def test_naming_none_of_them_covers_the_set(self, tmp_path, capsys):
        from haiku.rag.app import HaikuRAGApp

        config = _config(tmp_path, ["alpha", "beta"])
        await _seed(config, "alpha", ["alpha document about cats"])
        await _seed(config, "beta", ["beta document about cats"])

        app = HaikuRAGApp(
            scope=DatabaseScope.resolve(config), config=config, read_only=True
        )
        await app.list_documents()

        printed = capsys.readouterr().out
        assert "test://alpha/" in printed
        assert "test://beta/" in printed