haiku.rag/tests/test_cli.py

import subprocess
import sys
from pathlib import Path
from unittest.mock import AsyncMock, patch

import pytest
from click.exceptions import BadParameter
from typer.testing import CliRunner

from haiku.rag.cli import _cli as cli
from haiku.rag.cli import _parse_meta_options, resolve_scope
from haiku.rag.cli import cli as cli_wrapper
from haiku.rag.config import get_config, set_config
from haiku.rag.config.models import AppConfig, LanceDBConfig, StorageConfig
from haiku.rag.store.exceptions import (
    AmbiguousDatabaseError,
    MigrationRequiredError,
    UnknownDatabaseError,
)
from tests.conftest import for_path

runner = CliRunner()


def test_importing_cli_does_not_load_heavy_dependencies():
    """Importing the CLI must not pull the heavy runtime dependencies; they cost
    seconds of startup. Runs in a subprocess because another test in the same
    session may already have imported them."""
    result = subprocess.run(
        [
            sys.executable,
            "-c",
            "import haiku.rag.cli, sys; "
            "loaded = {'lancedb', 'pyarrow', 'pydantic_ai'} & sys.modules.keys(); "
            "assert not loaded, loaded",
        ],
        capture_output=True,
        text=True,
    )
    assert result.returncode == 0, result.stderr


class TestParseMetaOptions:
    def test_empty_input(self):
        assert _parse_meta_options(None) == {}
        assert _parse_meta_options([]) == {}

    def test_simple_key_value(self):
        result = _parse_meta_options(["author=alice", "topic=notes"])
        assert result == {"author": "alice", "topic": "notes"}

    def test_missing_equals_raises(self):
        with pytest.raises(BadParameter):
            _parse_meta_options(["no_equals_here"])

    def test_empty_key_raises(self):
        with pytest.raises(BadParameter):
            _parse_meta_options(["=value"])

    def test_json_number(self):
        result = _parse_meta_options(["version=3"])
        assert result == {"version": 3}
        assert isinstance(result["version"], int)

    def test_json_float(self):
        result = _parse_meta_options(["score=3.14"])
        assert result == {"score": 3.14}
        assert isinstance(result["score"], float)

    def test_json_bool(self):
        result = _parse_meta_options(["published=true", "draft=false"])
        assert result == {"published": True, "draft": False}

    def test_json_null(self):
        result = _parse_meta_options(["empty=null"])
        assert result == {"empty": None}

    def test_json_array(self):
        result = _parse_meta_options(['tags=["a","b","c"]'])
        assert result == {"tags": ["a", "b", "c"]}

    def test_json_object(self):
        result = _parse_meta_options(['nested={"x": 1}'])
        assert result == {"nested": {"x": 1}}

    def test_plain_string_not_json(self):
        result = _parse_meta_options(["name=hello world"])
        assert result == {"name": "hello world"}
        assert isinstance(result["name"], str)

    def test_value_with_equals_sign(self):
        result = _parse_meta_options(["equation=a=b+c"])
        assert result == {"equation": "a=b+c"}


class TestRebuildValidation:
    def test_embed_only_and_rechunk_mutually_exclusive(self):
        result = runner.invoke(
            cli, ["rebuild", "--embed-only", "--rechunk", "--db", "/tmp/fake.lancedb"]
        )
        assert result.exit_code == 1
        assert "mutually exclusive" in result.output


class TestOneDatabaseCommands:
    """`lancedb.databases` names a set; most commands work on one database."""

    @staticmethod
    def _install(monkeypatch, **databases):
        import haiku.rag.config as config_module

        monkeypatch.setattr(config_module, "_config", None)
        monkeypatch.setattr("haiku.rag.cli._db_name", None)
        set_config(AppConfig(lancedb=LanceDBConfig(databases=databases)))

    def test_a_configured_set_refuses_a_one_database_command(self, monkeypatch):
        self._install(monkeypatch, alpha="/db/a.lancedb", beta="/db/b.lancedb")

        with pytest.raises(AmbiguousDatabaseError, match="alpha, beta") as raised:
            resolve_scope(None)

        # `--db-name` is global and `--db` is per-command, so the remedy says
        # where each one goes.
        assert "--db-name NAME before the command" in str(raised.value)
        assert "--db PATH after it" in str(raised.value)

    def test_the_refusal_names_the_databases_and_not_their_locations(self, monkeypatch):
        """The refusal carries names only; a location in an error message
        travels into logs and terminals."""
        self._install(
            monkeypatch,
            alpha="s3://bucket/prefix/a.lancedb",
            beta="s3://bucket/prefix/b.lancedb",
        )

        with pytest.raises(AmbiguousDatabaseError) as raised:
            resolve_scope(None)

        assert "alpha" in str(raised.value)
        assert "s3://bucket/prefix/a.lancedb" not in str(raised.value)
        assert "bucket" not in str(raised.value)

    def test_a_configured_set_of_one_needs_no_choosing(self, monkeypatch):
        """Nothing is ambiguous about a set with one database in it, and it
        keeps its name."""
        self._install(monkeypatch, alpha="/db/a.lancedb")

        assert resolve_scope(None).names == ("alpha",)

    def test_a_command_covering_the_set_is_allowed(self, monkeypatch):
        self._install(monkeypatch, alpha="/db/a.lancedb", beta="/db/b.lancedb")

        assert resolve_scope(None, covers_set=True).names == ("alpha", "beta")

    def test_naming_a_path_is_allowed(self, monkeypatch):
        self._install(monkeypatch, alpha="/db/a.lancedb")

        [ref] = resolve_scope(Path("/db/other.lancedb")).databases
        assert ref.db_path == Path("/db/other.lancedb")

    def test_naming_a_path_overrides_a_configured_set(self, monkeypatch):
        """`--db` is the operator's explicit override: it constructs the scope
        directly, where a Python caller passing a path beside `databases` is
        refused."""
        self._install(monkeypatch, alpha="/db/a.lancedb", beta="/db/b.lancedb")

        scope = resolve_scope(Path("/db/other.lancedb"))

        assert scope.names == ("other",)
        assert not scope.covers_multiple

    def test_a_path_without_a_stem_is_a_usage_error(self, monkeypatch):
        """A path that names no database is the operator's mistake, reported as
        one."""
        import typer

        self._install(monkeypatch)

        with pytest.raises(typer.BadParameter, match="no name"):
            resolve_scope(Path("/"))

    def test_no_configured_databases_is_allowed(self, monkeypatch, tmp_path):
        import haiku.rag.config as config_module

        monkeypatch.setattr(config_module, "_config", None)
        monkeypatch.setattr("haiku.rag.cli._db_name", None)
        set_config(AppConfig(storage=StorageConfig(data_dir=tmp_path)))

        [ref] = resolve_scope(None).databases
        assert ref.db_path == tmp_path / "haiku.rag.lancedb"

    def test_the_refusal_exits_with_an_error(self):
        with patch("haiku.rag.cli._cli") as mock_cli:
            mock_cli.side_effect = AmbiguousDatabaseError("names alpha, beta")

            with pytest.raises(SystemExit) as exc_info:
                cli_wrapper()
            assert exc_info.value.code == 1


class TestSelectingADatabaseByName:
    """`--db-name NAME` is the only way to reach a configured database whose
    location is a URI, since `--db` takes a path."""

    @staticmethod
    def _install(monkeypatch, **databases):
        import haiku.rag.config as config_module

        monkeypatch.setattr(config_module, "_config", None)
        set_config(AppConfig(lancedb=LanceDBConfig(databases=databases)))

    def test_a_named_database_is_passed_on_by_name(self, monkeypatch):
        """Not resolved to a path: the name is what results and citations
        carry."""
        self._install(monkeypatch, papers="s3://bucket/prefix/papers.lancedb")
        monkeypatch.setattr("haiku.rag.cli._db_name", "papers")

        assert resolve_scope(None).names == ("papers",)

    def test_naming_a_database_leaves_the_configuration_alone(self, monkeypatch):
        self._install(monkeypatch, notes="/data/notes.lancedb", other="/data/o.lancedb")
        monkeypatch.setattr("haiku.rag.cli._db_name", "notes")

        resolve_scope(None)

        config = get_config()
        assert config.lancedb.databases == {
            "notes": "/data/notes.lancedb",
            "other": "/data/o.lancedb",
        }

    def test_an_unknown_name_names_the_configured_ones(self, monkeypatch):
        self._install(monkeypatch, alpha="/data/a.lancedb", beta="/data/b.lancedb")
        monkeypatch.setattr("haiku.rag.cli._db_name", "gamma")

        with pytest.raises(UnknownDatabaseError, match="alpha, beta"):
            resolve_scope(None)

    def test_an_unknown_name_does_not_leak_locations(self, monkeypatch):
        self._install(monkeypatch, papers="s3://bucket/prefix/papers.lancedb")
        monkeypatch.setattr("haiku.rag.cli._db_name", "gamma")

        with pytest.raises(UnknownDatabaseError) as raised:
            resolve_scope(None)

        assert "bucket" not in str(raised.value)

    def test_an_unknown_name_with_nothing_configured_names_the_default(
        self, monkeypatch
    ):
        """Nothing configured is the one entry `haiku.rag`, which the message
        offers."""
        self._install(monkeypatch)
        monkeypatch.setattr("haiku.rag.cli._db_name", "papers")

        with pytest.raises(UnknownDatabaseError, match="haiku.rag"):
            resolve_scope(None)

    def test_the_callback_selects_before_a_command_runs(self, tmp_path, monkeypatch):
        """`--db-name` is resolved once the config is loaded, so every command
        and both TUIs see the selected database."""
        import haiku.rag.config as config_module

        monkeypatch.setattr(config_module, "_config", None)
        config_file = tmp_path / "haiku.rag.yaml"
        config_file.write_text("lancedb:\n  databases:\n    alpha: /data/a.lancedb\n")

        result = runner.invoke(
            cli, ["--config", str(config_file), "--db-name", "nope", "list"]
        )

        assert result.exit_code != 0
        assert isinstance(result.exception, UnknownDatabaseError)
        assert "nope" in str(result.exception)

    def test_a_selection_does_not_outlive_its_invocation(self, tmp_path, monkeypatch):
        """The selector is module state, so a second invocation without
        `--db-name` must not inherit the first one's database."""
        import haiku.rag.cli as cli_module
        import haiku.rag.config as config_module

        monkeypatch.setattr(config_module, "_config", None)
        monkeypatch.setattr(cli_module, "_db_name", None)
        config_file = tmp_path / "haiku.rag.yaml"
        selected = tmp_path / "alpha.lancedb"
        config_file.write_text(f"lancedb:\n  databases:\n    alpha: {selected}\n")

        runner.invoke(cli, ["--config", str(config_file), "--db-name", "alpha", "info"])
        assert cli_module._db_name == "alpha"

        runner.invoke(cli, ["--config", str(config_file), "settings"])

        assert cli_module._db_name is None

    def test_a_selection_does_not_outlive_its_invocation_in_process(
        self, tmp_path, monkeypatch
    ):
        """The selection is per invocation: a second one starts from a freshly
        loaded configuration."""
        import haiku.rag.cli as cli_module
        import haiku.rag.config as config_module

        monkeypatch.setattr(config_module, "_config", None)
        monkeypatch.setattr(cli_module, "_db_name", None)
        config_file = tmp_path / "haiku.rag.yaml"
        config_file.write_text(
            f"lancedb:\n  databases:\n    alpha: {tmp_path / 'alpha.lancedb'}\n"
            f"    beta: {tmp_path / 'beta.lancedb'}\n"
        )

        runner.invoke(
            cli, ["--config", str(config_file), "--db-name", "alpha", "settings"]
        )
        # Naming one leaves the configuration naming both.
        assert set(get_config().lancedb.databases) == {"alpha", "beta"}

        runner.invoke(cli, ["--config", str(config_file), "settings"])

        assert set(get_config().lancedb.databases) == {"alpha", "beta"}
        assert cli_module._db_name is None

    def test_a_selection_does_not_outlive_an_invocation_without_a_config_file(
        self, tmp_path, monkeypatch
    ):
        """No config file is still a load: the previous invocation's database
        must not be what the next one talks to."""
        import haiku.rag.cli as cli_module
        import haiku.rag.config as config_module

        monkeypatch.setattr(config_module, "_config", None)
        monkeypatch.setattr(cli_module, "_db_name", None)
        monkeypatch.delenv("HAIKU_RAG_CONFIG_PATH", raising=False)
        monkeypatch.chdir(tmp_path)
        config_file = tmp_path / "selected.yaml"
        config_file.write_text(
            "lancedb:\n  databases:\n    papers: s3://bucket/papers.lancedb\n"
        )

        runner.invoke(
            cli, ["--config", str(config_file), "--db-name", "papers", "settings"]
        )
        assert get_config().lancedb.databases == {
            "papers": "s3://bucket/papers.lancedb"
        }

        runner.invoke(cli, ["settings"])

        assert get_config().lancedb.databases == {}
        assert cli_module._db_name is None


class TestCommandsThatWorkOnNoDatabase:
    """`settings` and `download-models` read the configuration. A name that
    selects a database is nothing to them, including a wrong one."""

    def test_settings_ignores_an_unknown_name(self, tmp_path, monkeypatch):
        import haiku.rag.config as config_module

        monkeypatch.setattr(config_module, "_config", None)
        config_file = tmp_path / "haiku.rag.yaml"
        config_file.write_text("lancedb:\n  databases:\n    alpha: /data/a.lancedb\n")

        result = runner.invoke(
            cli, ["--config", str(config_file), "--db-name", "nope", "settings"]
        )

        assert result.exit_code == 0, result.output
        assert "haiku.rag configuration" in result.output

    def test_download_models_ignores_an_unknown_name(self, tmp_path, monkeypatch):
        import haiku.rag.config as config_module

        monkeypatch.setattr(config_module, "_config", None)
        config_file = tmp_path / "haiku.rag.yaml"
        config_file.write_text("lancedb:\n  databases:\n    alpha: /data/a.lancedb\n")
        downloaded: list[object] = []

        async def nothing_to_download(config):
            downloaded.append(config)
            return
            yield  # pragma: no cover - an empty async generator needs one

        monkeypatch.setattr(
            "haiku.rag.client.downloads.download_models", nothing_to_download
        )

        result = runner.invoke(
            cli,
            ["--config", str(config_file), "--db-name", "nope", "download-models"],
        )

        assert result.exit_code == 0, result.output
        assert len(downloaded) == 1


class TestResolvingTheDatabasePath:
    def test_a_path_wins_when_nothing_is_selected(self, monkeypatch):
        monkeypatch.setattr("haiku.rag.cli._db_name", None)

        [ref] = resolve_scope(Path("/data/one.lancedb")).databases
        assert ref.db_path == Path("/data/one.lancedb")

    def test_naming_a_database_twice_is_refused(self, monkeypatch):
        monkeypatch.setattr("haiku.rag.cli._db_name", "notes")

        with pytest.raises(AmbiguousDatabaseError, match="not both"):
            resolve_scope(Path("/data/other.lancedb"))


class TestConfiguredLocalDatabase:
    """One entry in `lancedb.databases` with a local path."""

    def _config_file(self, tmp_path, located: Path) -> Path:
        config_file = tmp_path / "haiku.rag.yaml"
        config_file.write_text(f"lancedb:\n  databases:\n    notes: {located}\n")
        return config_file

    def _fresh(self, monkeypatch) -> None:
        import haiku.rag.cli as cli_module
        import haiku.rag.config as config_module

        monkeypatch.setattr(config_module, "_config", None)
        monkeypatch.setattr(cli_module, "_db_name", None)

    def _reports(self, result, located: Path) -> bool:
        """Rich wraps a long path to the terminal width, so compare without it."""
        return str(located) in "".join(result.output.split())

    def test_init_creates_the_configured_local_database(self, tmp_path, monkeypatch):
        self._fresh(monkeypatch)
        located = tmp_path / "notes.lancedb"
        config_file = self._config_file(tmp_path, located)

        result = runner.invoke(cli, ["--config", str(config_file), "init"])

        assert result.exit_code == 0, result.output
        assert located.exists()

        self._fresh(monkeypatch)
        result = runner.invoke(cli, ["--config", str(config_file), "info"])
        assert result.exit_code == 0, result.output
        assert self._reports(result, located)

    def test_a_missing_configured_path_is_refused(self, tmp_path, monkeypatch):
        """A schemeless value is a local path and must exist."""
        self._fresh(monkeypatch)
        located = tmp_path / "typo.lancedb"
        config_file = self._config_file(tmp_path, located)

        result = runner.invoke(cli, ["--config", str(config_file), "info"])

        assert "does not exist" in result.output
        assert not located.exists()

    def test_db_overrides_the_configured_database(self, tmp_path, monkeypatch):
        self._fresh(monkeypatch)
        configured = tmp_path / "configured.lancedb"
        chosen = tmp_path / "chosen.lancedb"
        config_file = self._config_file(tmp_path, configured)

        result = runner.invoke(
            cli, ["--config", str(config_file), "init", "--db", str(chosen)]
        )

        assert result.exit_code == 0, result.output
        assert chosen.exists()
        assert not configured.exists()

        self._fresh(monkeypatch)
        result = runner.invoke(
            cli, ["--config", str(config_file), "info", "--db", str(chosen)]
        )
        assert result.exit_code == 0, result.output
        assert self._reports(result, chosen)


class TestCliConfigMismatchError:
    def test_a_config_mismatch_exits_with_its_remedy(self):
        """The message says which database and what to run, so it is worth more
        than a traceback."""
        from haiku.rag.store.exceptions import ConfigMismatchError

        with patch("haiku.rag.cli._cli") as mock_cli:
            mock_cli.side_effect = ConfigMismatchError(
                "database 'nemotron': vector dimension 2048 -> 2560"
            )

            with pytest.raises(SystemExit) as exc_info:
                cli_wrapper()
            assert exc_info.value.code == 1


class TestCliMigrationError:
    def test_catches_migration_required_error(self):
        with patch("haiku.rag.cli._cli") as mock_cli:
            mock_cli.side_effect = MigrationRequiredError(
                "Database requires migration. Run 'haiku-rag migrate' to upgrade."
            )

            with pytest.raises(SystemExit) as exc_info:
                cli_wrapper()
            assert exc_info.value.code == 1


class TestTagCommands:
    def test_tag_round_trip(self, temp_db_path):
        db = str(temp_db_path)

        result = runner.invoke(cli, ["init", "--db", db])
        assert result.exit_code == 0

        result = runner.invoke(cli, ["tag", "create", "release-1", "--db", db])
        assert result.exit_code == 0
        assert "release-1" in result.output

        result = runner.invoke(cli, ["tag", "list", "--db", db])
        assert result.exit_code == 0
        assert "release-1" in result.output
        assert "partial" not in result.output

        result = runner.invoke(cli, ["history", "--db", db, "-t", "documents"])
        assert result.exit_code == 0
        assert "release-1" in result.output

        result = runner.invoke(cli, ["tag", "create", "release-1", "--db", db])
        assert result.exit_code == 1
        assert "already exists" in result.output

        result = runner.invoke(cli, ["tag", "delete", "release-1", "--db", db])
        assert result.exit_code == 0

        result = runner.invoke(cli, ["tag", "list", "--db", db])
        assert result.exit_code == 0
        assert "No tags" in result.output

        result = runner.invoke(cli, ["tag", "delete", "release-1", "--db", db])
        assert result.exit_code == 1
        assert "does not exist" in result.output

    def test_tag_create_rejected_when_migrations_pending(self, temp_db_path):
        """A writable tag operation must hit the migration gate and must not
        mutate a legacy database (e.g. by creating missing tables)."""
        import asyncio

        import lancedb

        from haiku.rag.store.engine import Store

        async def _prepare_legacy_db():
            async with Store(temp_db_path, create=True) as store:
                await store.set_haiku_version("0.19.0")
            db = await lancedb.connect_async(temp_db_path.absolute())
            await db.drop_table("document_meta")
            db.close()

        asyncio.run(_prepare_legacy_db())

        result = runner.invoke(
            cli, ["tag", "create", "release-1", "--db", str(temp_db_path)]
        )
        assert result.exit_code == 1
        assert isinstance(result.exception, MigrationRequiredError)

        async def _table_names() -> list[str]:
            db = await lancedb.connect_async(temp_db_path.absolute())
            tables = (await db.list_tables()).tables
            db.close()
            return tables

        assert "document_meta" not in asyncio.run(_table_names())

    def test_tag_commands_missing_database_exit_nonzero(self, tmp_path):
        missing = str(tmp_path / "does_not_exist.lancedb")
        for args in (
            ["tag", "create", "r1", "--db", missing],
            ["tag", "delete", "r1", "--db", missing],
            ["tag", "list", "--db", missing],
        ):
            result = runner.invoke(cli, args)
            assert result.exit_code == 1, args
            assert "does not exist" in result.output, args

    def test_tag_create_invalid_name_fails_cleanly(self, temp_db_path):
        """lance restricts ref names to alphanumeric, '.', '-', '_'; the CLI
        surfaces that as a clean error instead of a traceback."""
        db = str(temp_db_path)
        result = runner.invoke(cli, ["init", "--db", db])
        assert result.exit_code == 0

        result = runner.invoke(cli, ["tag", "create", "[red]release[/red]", "--db", db])
        assert result.exit_code == 1
        assert "Error:" in result.output
        assert "Ref characters" in result.output


class TestTagRestore:
    def test_restore_requires_confirmation_and_decline_changes_nothing(
        self, temp_db_path
    ):
        db = str(temp_db_path)
        assert runner.invoke(cli, ["init", "--db", db]).exit_code == 0
        assert runner.invoke(cli, ["tag", "create", "r1", "--db", db]).exit_code == 0

        result = runner.invoke(cli, ["tag", "restore", "r1", "--db", db], input="n\n")
        assert result.exit_code == 1
        assert "live database state" in result.output
        assert "Stop all ingestion" in result.output
        assert "not transactionally atomic" in result.output
        assert "safety tag" in result.output

        result = runner.invoke(cli, ["tag", "list", "--db", db])
        assert "before-restore" not in result.output

    def test_restore_non_interactive_without_yes_fails(self, temp_db_path):
        db = str(temp_db_path)
        assert runner.invoke(cli, ["init", "--db", db]).exit_code == 0
        assert runner.invoke(cli, ["tag", "create", "r1", "--db", db]).exit_code == 0

        result = runner.invoke(cli, ["tag", "restore", "r1", "--db", db])
        assert result.exit_code == 1

        result = runner.invoke(cli, ["tag", "list", "--db", db])
        assert "before-restore" not in result.output

    def test_restore_with_yes(self, temp_db_path):
        db = str(temp_db_path)
        assert runner.invoke(cli, ["init", "--db", db]).exit_code == 0
        assert runner.invoke(cli, ["tag", "create", "r1", "--db", db]).exit_code == 0

        result = runner.invoke(cli, ["tag", "restore", "r1", "--yes", "--db", db])
        assert result.exit_code == 0
        assert "Restored database to tag 'r1'" in result.output
        assert "before-restore-" in result.output
        assert "now live" in result.output
        assert "migrate" in result.output

        result = runner.invoke(cli, ["tag", "list", "--db", db])
        assert "before-restore-" in result.output

    def test_restore_missing_tag_errors(self, temp_db_path):
        db = str(temp_db_path)
        assert runner.invoke(cli, ["init", "--db", db]).exit_code == 0

        result = runner.invoke(cli, ["tag", "restore", "nope", "--yes", "--db", db])
        assert result.exit_code == 1
        assert "does not exist" in result.output

    def test_restore_partial_tag_errors(self, temp_db_path):
        import asyncio

        from haiku.rag.store.engine import Store

        async def _partial_tag():
            async with Store(temp_db_path, create=True) as store:
                version = await store.chunks_table.version()
                await store.chunks_table.tags.create("stale", version)

        asyncio.run(_partial_tag())

        result = runner.invoke(
            cli, ["tag", "restore", "stale", "--yes", "--db", str(temp_db_path)]
        )
        assert result.exit_code == 1
        assert "partial" in result.output
        assert "documents" in result.output

    def test_restore_missing_database_exits_nonzero(self, tmp_path):
        missing = tmp_path / "does_not_exist.lancedb"
        result = runner.invoke(
            cli, ["tag", "restore", "r1", "--yes", "--db", str(missing)]
        )
        assert result.exit_code == 1
        assert "does not exist" in result.output

        # Without --yes the missing database is reported before the
        # confirmation prompt, not after the user confirms.
        result = runner.invoke(cli, ["tag", "restore", "r1", "--db", str(missing)])
        assert result.exit_code == 1
        assert "does not exist" in result.output
        assert "Continue?" not in result.output

    def test_tag_help_includes_restore(self):
        result = runner.invoke(cli, ["tag", "--help"])
        assert result.exit_code == 0
        assert "restore" in result.output

        result = runner.invoke(cli, ["--help"])
        assert "--before" not in result.output
        assert "--at" not in result.output


class TestAskAnalyzeImageOption:
    def test_ask_forwards_image_paths(self):
        from unittest.mock import AsyncMock

        from haiku.rag.app import HaikuRAGApp

        with patch.object(HaikuRAGApp, "ask", new_callable=AsyncMock) as mock_ask:
            result = runner.invoke(
                cli,
                ["ask", "q", "--image", "/tmp/a.png", "--image", "/tmp/b.jpg"],
            )
        assert result.exit_code == 0
        from pathlib import Path

        assert mock_ask.call_args.kwargs["images"] == [
            Path("/tmp/a.png"),
            Path("/tmp/b.jpg"),
        ]

    def test_analyze_forwards_image_paths(self):
        from unittest.mock import AsyncMock

        from haiku.rag.app import HaikuRAGApp

        with patch.object(HaikuRAGApp, "analyze", new_callable=AsyncMock) as mock:
            result = runner.invoke(cli, ["analyze", "q", "--image", "/tmp/a.png"])
        assert result.exit_code == 0
        from pathlib import Path

        assert mock.call_args.kwargs["images"] == [Path("/tmp/a.png")]

    @pytest.mark.asyncio
    async def test_app_ask_reads_image_bytes(self, temp_db_path, tmp_path):
        from io import BytesIO
        from unittest.mock import AsyncMock

        from PIL import Image as PILImage

        from haiku.rag.app import HaikuRAGApp
        from haiku.rag.client import HaikuRAG

        buffer = BytesIO()
        PILImage.new("RGB", (4, 4)).save(buffer, format="PNG")
        img_path = tmp_path / "img.png"
        img_path.write_bytes(buffer.getvalue())

        async with HaikuRAG(temp_db_path, create=True):
            pass

        with patch.object(
            HaikuRAG, "ask", new_callable=AsyncMock, return_value=("answer", [])
        ) as mock_ask:
            app = HaikuRAGApp(scope=for_path(temp_db_path))
            await app.ask("q", images=[img_path])

        assert mock_ask.call_args.kwargs["images"] == [buffer.getvalue()]


class TestChatCoversTheSet:
    """Chat is a read verb: it answers with the same capabilities `ask` uses,
    and covers the configured set."""

    @staticmethod
    def _config_file(tmp_path):
        config_file = tmp_path / "haiku.rag.yaml"
        config_file.write_text(
            f"lancedb:\n  databases:\n    arxiv: {tmp_path / 'a.lancedb'}\n"
            f"    wiki: {tmp_path / 'w.lancedb'}\n"
        )
        return config_file

    def test_a_configured_set_is_covered_rather_than_refused(
        self, tmp_path, monkeypatch
    ):
        import haiku.rag.cli as cli_module
        import haiku.rag.config as config_module

        monkeypatch.setattr(config_module, "_config", None)
        monkeypatch.setattr(cli_module, "_db_name", None)

        with patch("haiku.rag.chat.run_chat") as run_chat:
            result = runner.invoke(
                cli, ["--config", str(self._config_file(tmp_path)), "chat"]
            )

        assert result.exit_code == 0, result.output
        # The scope covers both, which is what makes chat read the set.
        assert run_chat.call_args.kwargs["scope"].names == ("arxiv", "wiki")

    def test_naming_one_database_opens_that_one(self, tmp_path, monkeypatch):
        import haiku.rag.cli as cli_module
        import haiku.rag.config as config_module

        monkeypatch.setattr(config_module, "_config", None)
        monkeypatch.setattr(cli_module, "_db_name", None)

        with patch("haiku.rag.chat.run_chat") as run_chat:
            result = runner.invoke(
                cli,
                [
                    "--config",
                    str(self._config_file(tmp_path)),
                    "--db-name",
                    "wiki",
                    "chat",
                ],
            )

        assert result.exit_code == 0, result.output
        # Passed on by name, so the database it opens keeps its identity.
        assert run_chat.call_args.kwargs["scope"].names == ("wiki",)

    def test_a_single_database_setup_is_unchanged(self, tmp_path, monkeypatch):
        """Without a configured set, chat opens the path it always did."""
        import haiku.rag.cli as cli_module
        import haiku.rag.config as config_module

        monkeypatch.setattr(config_module, "_config", None)
        monkeypatch.setattr(cli_module, "_db_name", None)

        with patch("haiku.rag.chat.run_chat") as run_chat:
            result = runner.invoke(cli, ["chat", "--db", str(tmp_path / "one.lancedb")])

        assert result.exit_code == 0, result.output
        [ref] = run_chat.call_args.kwargs["scope"].databases
        assert ref.db_path == tmp_path / "one.lancedb"


class TestRenderingTheDatabase:
    """Across databases a result has to say which one it came from. One database
    needs no such label, so single-database output is unchanged."""

    @staticmethod
    def _app(tmp_path, **databases):
        """An app over the configured set. The scope is what decides the label,
        so it has to be the one the configuration resolves to."""
        from haiku.rag.app import HaikuRAGApp
        from haiku.rag.client.scope import DatabaseScope

        config = AppConfig(lancedb=LanceDBConfig(databases=databases))
        return HaikuRAGApp(scope=DatabaseScope.resolve(config), config=config)

    @staticmethod
    def _rendered(app, result) -> str:
        from rich.console import Console

        app.console = Console(record=True, width=200)
        app._rich_print_search_result(result)
        return app.console.export_text()

    @staticmethod
    def _result():
        from haiku.rag.store.models import SearchResult

        return SearchResult(
            content="a body",
            score=0.9,
            source="alpha",
            chunk_id="c1",
            document_id="d1",
            document_uri="test://alpha/one",
        )

    def test_a_set_labels_each_result(self, tmp_path):
        app = self._app(
            tmp_path,
            alpha=str(tmp_path / "alpha.lancedb"),
            beta=str(tmp_path / "beta.lancedb"),
        )

        assert "database: alpha" in self._rendered(app, self._result())

    def test_one_database_is_not_labelled(self, tmp_path):
        app = self._app(tmp_path, alpha=str(tmp_path / "alpha.lancedb"))

        assert "database:" not in self._rendered(app, self._result())


@pytest.fixture
def app_stub(monkeypatch, tmp_path):
    """Stand in for HaikuRAGApp so a command's wiring can be checked without a
    database or a model. These tests pin argument parsing and dispatch, not what
    the application layer renders."""
    # AsyncMock so every command's `asyncio.run(app.x(...))` gets a coroutine.
    stub = AsyncMock()
    monkeypatch.setattr(
        "haiku.rag.cli.create_app",
        lambda db=None, *, covers_set=False: stub,
    )
    return stub


DB_ARGS = ["--db", "/tmp/test.lancedb"]


@pytest.mark.parametrize(
    "argv, method, expected",
    [
        (["list"], "list_documents", {"filter": None}),
        (
            ["list", "--filter", "uri LIKE 'x%'"],
            "list_documents",
            {"filter": "uri LIKE 'x%'"},
        ),
        (
            ["add", "some text", "--title", "T"],
            "add_document_from_text",
            {"text": "some text", "title": "T", "metadata": None},
        ),
        (
            ["add-src", "/tmp/doc.md"],
            "add_document_from_source",
            {"source": "/tmp/doc.md", "title": None, "metadata": None},
        ),
        (["get", "doc-1"], "get_document", {"doc_id": "doc-1"}),
        (["delete", "doc-1"], "delete_document", {"doc_id": "doc-1"}),
        (
            ["visualize", "chunk-1"],
            "visualize_chunk",
            {"chunk_id": "chunk-1", "expand": True},
        ),
        (
            ["visualize", "chunk-1", "--no-expand"],
            "visualize_chunk",
            {"chunk_id": "chunk-1", "expand": False},
        ),
        (["vacuum"], "vacuum", {}),
        (["create-index"], "create_index", {}),
        (["init"], "init", {}),
        (["info"], "info", {}),
        # limit/search_type default to None: the app layer resolves the config
        # default, so the CLI must not invent one.
        (["history"], "history", {"table": None, "limit": None}),
        (
            ["history", "--table", "chunks", "--limit", "5"],
            "history",
            {"table": "chunks", "limit": 5},
        ),
    ],
)
def test_command_dispatches_to_the_app(app_stub, argv, method, expected):
    result = runner.invoke(cli, argv + DB_ARGS)

    assert result.exit_code == 0, result.output
    getattr(app_stub, method).assert_called_once_with(**expected)


@pytest.mark.parametrize(
    "argv, expected",
    [
        (
            ["search", "q"],
            {
                "query": "q",
                "limit": None,
                "filter": None,
                "search_type": None,
                "image": None,
            },
        ),
        (
            ["search", "q", "--limit", "3", "--search-type", "vector"],
            {
                "query": "q",
                "limit": 3,
                "filter": None,
                "search_type": "vector",
                "image": None,
            },
        ),
    ],
)
def test_search_dispatch(app_stub, argv, expected):
    result = runner.invoke(cli, argv + DB_ARGS)

    assert result.exit_code == 0, result.output
    app_stub.search.assert_called_once_with(**expected)


@pytest.mark.parametrize("command, method", [("ask", "ask"), ("analyze", "analyze")])
def test_question_commands_dispatch(app_stub, command, method):
    result = runner.invoke(cli, [command, "why?"] + DB_ARGS)

    assert result.exit_code == 0, result.output
    getattr(app_stub, method).assert_called_once_with(
        question="why?", filter=None, images=None, full_citations=False
    )


@pytest.mark.parametrize("command, method", [("ask", "ask"), ("analyze", "analyze")])
def test_full_citations_flag_dispatches(app_stub, command, method):
    result = runner.invoke(cli, [command, "why?", "--full-citations"] + DB_ARGS)

    assert result.exit_code == 0, result.output
    assert getattr(app_stub, method).call_args.kwargs["full_citations"] is True


@pytest.mark.parametrize(
    "flag, mode_name",
    [
        (None, "FULL"),
        ("--embed-only", "EMBED_ONLY"),
        ("--rechunk", "RECHUNK"),
        ("--title-only", "TITLE_ONLY"),
        ("--descriptions", "DESCRIPTIONS"),
        ("--set-embedder", "SET_EMBEDDER"),
    ],
)
def test_rebuild_flag_selects_the_mode(app_stub, flag, mode_name):
    """Each flag picks one rebuild mode, and no flag means a full rebuild."""
    result = runner.invoke(cli, ["rebuild"] + ([flag] if flag else []) + DB_ARGS)

    assert result.exit_code == 0, result.output
    (_, kwargs) = app_stub.rebuild.call_args
    assert kwargs["mode"].name == mode_name


def test_migrate_reports_applied_migrations(app_stub):
    app_stub.migrate.return_value = ["v0_40_0: add document_items"]

    result = runner.invoke(cli, ["migrate"] + DB_ARGS)

    assert result.exit_code == 0, result.output
    assert "Applied 1 migration(s)" in result.output
    assert "add document_items" in result.output


def test_migrate_reports_an_up_to_date_database(app_stub):
    app_stub.migrate.return_value = []

    result = runner.invoke(cli, ["migrate"] + DB_ARGS)

    assert result.exit_code == 0, result.output
    assert "No migrations pending" in result.output


def test_migrate_exits_nonzero_on_failure(app_stub):
    app_stub.migrate.side_effect = RuntimeError("schema is from the future")

    result = runner.invoke(cli, ["migrate"] + DB_ARGS)

    assert result.exit_code == 1
    assert "Migration failed: schema is from the future" in result.output


def test_mcp_stdio_selects_the_transport(app_stub):
    result = runner.invoke(cli, ["mcp", "--stdio"] + DB_ARGS)

    assert result.exit_code == 0, result.output
    app_stub.run_mcp.assert_called_once()
    kwargs = app_stub.run_mcp.call_args.kwargs
    assert kwargs["transport"] == "stdio"


def test_mcp_without_stdio_leaves_the_transport_unset(app_stub):
    result = runner.invoke(cli, ["mcp"] + DB_ARGS)

    assert result.exit_code == 0, result.output
    assert app_stub.run_mcp.call_args.kwargs["transport"] is None


def test_mcp_covers_the_configured_set(monkeypatch):
    seen = {}

    def create_app(db=None, *, covers_set=False):
        seen["covers_set"] = covers_set
        return AsyncMock()

    monkeypatch.setattr("haiku.rag.cli.create_app", create_app)

    result = runner.invoke(cli, ["mcp", "--stdio"])

    assert result.exit_code == 0, result.output
    assert seen["covers_set"] is True


def test_version_flag_prints_the_version():
    result = runner.invoke(cli, ["--version"])

    assert result.exit_code == 0, result.output
    assert "haiku.rag version" in result.output


def test_outdated_install_warns(app_stub, monkeypatch):
    """The startup check warns but does not block the command."""

    async def outdated():
        return False, "0.1.0", "9.9.9"

    monkeypatch.setattr("haiku.rag.cli.is_up_to_date", outdated)

    result = runner.invoke(cli, ["list"] + DB_ARGS)

    assert result.exit_code == 0, result.output
    assert "haiku.rag is outdated" in result.output
    assert "Current: 0.1.0, Latest: 9.9.9" in result.output
    app_stub.list_documents.assert_called_once()


def test_up_to_date_install_says_nothing(app_stub, monkeypatch):
    async def current():
        return True, "9.9.9", "9.9.9"

    monkeypatch.setattr("haiku.rag.cli.is_up_to_date", current)

    result = runner.invoke(cli, ["list"] + DB_ARGS)

    assert result.exit_code == 0, result.output
    assert "outdated" not in result.output


def test_a_failing_version_check_does_not_block_the_cli(app_stub, monkeypatch):
    """PyPI being unreachable must not stop a command from running."""

    async def boom():
        raise RuntimeError("no network")

    monkeypatch.setattr("haiku.rag.cli.is_up_to_date", boom)

    result = runner.invoke(cli, ["list"] + DB_ARGS)

    assert result.exit_code == 0, result.output
    app_stub.list_documents.assert_called_once()


def test_settings_command_shows_the_configuration(monkeypatch):
    shown = []

    class StubApp:
        def __init__(self, **kwargs):
            shown.append(kwargs)

        def show_settings(self):
            shown.append("shown")

    monkeypatch.setattr("haiku.rag.app.HaikuRAGApp", StubApp)

    result = runner.invoke(cli, ["settings"])

    assert result.exit_code == 0, result.output
    assert "shown" in shown
    assert shown[0]["read_only"] is True


def test_download_models_reports_failure(monkeypatch):
    class StubApp:
        def __init__(self, **kwargs):
            pass

        async def download_models(self):
            raise RuntimeError("hub unreachable")

    monkeypatch.setattr("haiku.rag.app.HaikuRAGApp", StubApp)

    result = runner.invoke(cli, ["download-models"])

    assert result.exit_code == 1
    assert "Error downloading models: hub unreachable" in result.output


def test_download_models_succeeds(monkeypatch):
    calls = []

    class StubApp:
        def __init__(self, **kwargs):
            pass

        async def download_models(self):
            calls.append("downloaded")

    monkeypatch.setattr("haiku.rag.app.HaikuRAGApp", StubApp)

    result = runner.invoke(cli, ["download-models"])

    assert result.exit_code == 0, result.output
    assert calls == ["downloaded"]


def test_chat_reports_a_missing_tui_extra(monkeypatch):
    """run_chat imports the Textual app itself, so a missing tui extra surfaces
    from the call, not from importing haiku.rag.chat. The CLI must report it and
    exit nonzero rather than traceback."""
    import sys

    monkeypatch.setitem(sys.modules, "haiku.rag.chat.app", None)

    result = runner.invoke(cli, ["chat"])

    assert result.exit_code == 1
    assert "textual is not installed" in result.output
    assert "haiku.rag-slim[tui]" in result.output


def test_inspect_reports_a_missing_tui_extra(monkeypatch):
    """run_inspector raises at import instead, so the guard sits on the import."""
    import sys

    monkeypatch.delitem(sys.modules, "haiku.rag.inspector", raising=False)
    monkeypatch.setitem(sys.modules, "haiku.rag.inspector.app", None)

    result = runner.invoke(cli, ["inspect"])

    assert result.exit_code == 1
    assert "textual is not installed" in result.output
    assert "haiku.rag-slim[tui]" in result.output