from contextlib import asynccontextmanager
from typing import AsyncIterator

import httpx
import pytest

from basic_memory.cli.auth import CLIAuth
from basic_memory.config import ProjectMode
from basic_memory.mcp import async_client as async_client_module
from basic_memory.mcp.async_client import (
    get_client,
    get_cloud_control_plane_client,
    set_client_factory,
)


@pytest.fixture(autouse=True)
def _reset_async_client_state(monkeypatch):
    async_client_module._client_factory = None
    async_client_module._prepared_local_asgi_databases.clear()
    async_client_module._prepared_local_asgi_database_prepare_locks.clear()
    monkeypatch.delenv("BASIC_MEMORY_FORCE_LOCAL", raising=False)
    monkeypatch.delenv("BASIC_MEMORY_FORCE_CLOUD", raising=False)
    monkeypatch.delenv("BASIC_MEMORY_EXPLICIT_ROUTING", raising=False)
    yield
    async_client_module._client_factory = None
    async_client_module._prepared_local_asgi_databases.clear()
    async_client_module._prepared_local_asgi_database_prepare_locks.clear()


@pytest.mark.asyncio
async def test_get_client_uses_injected_factory(monkeypatch):
    seen = {"used": False}

    @asynccontextmanager
    async def factory(workspace=None):
        seen["used"] = True
        async with httpx.AsyncClient(base_url="https://example.test") as client:
            yield client

    # Ensure we don't leak factory to other tests
    set_client_factory(factory)
    async with get_client() as client:
        assert str(client.base_url) == "https://example.test"
    assert seen["used"] is True


@pytest.mark.asyncio
async def test_get_client_default_uses_local_asgi_transport(config_manager):
    cfg = config_manager.load_config()
    cfg.cloud_host = "https://cloud.example.test"
    cfg.cloud_api_key = "bmc_test_key_123"
    config_manager.save_config(cfg)

    async with get_client() as client:
        assert isinstance(client._transport, httpx.ASGITransport)  # pyright: ignore[reportPrivateUsage]


@pytest.mark.asyncio
async def test_get_client_preinitializes_local_asgi_database(config_manager, monkeypatch):
    """Local ASGI routing initializes DB state before request handling."""
    from basic_memory import db
    from basic_memory.api.app import app as fastapi_app

    cfg = config_manager.load_config()
    config_manager.save_config(cfg)

    previous_engine = getattr(fastapi_app.state, "engine", None)
    previous_session_maker = getattr(fastapi_app.state, "session_maker", None)
    fastapi_app.state._state.pop("engine", None)  # pyright: ignore[reportPrivateUsage]
    fastapi_app.state._state.pop("session_maker", None)  # pyright: ignore[reportPrivateUsage]

    engine = object()
    session_maker = object()
    calls = []

    async def fake_get_or_create_db(db_path):
        calls.append(db_path)
        return engine, session_maker

    monkeypatch.setattr(db, "get_or_create_db", fake_get_or_create_db)

    try:
        async with get_client() as client:
            assert isinstance(client._transport, httpx.ASGITransport)  # pyright: ignore[reportPrivateUsage]
            assert calls == [cfg.database_path]
            assert fastapi_app.state.engine is engine
            assert fastapi_app.state.session_maker is session_maker
        assert not hasattr(fastapi_app.state, "engine")
        assert not hasattr(fastapi_app.state, "session_maker")
    finally:
        if previous_engine is None:
            fastapi_app.state._state.pop("engine", None)  # pyright: ignore[reportPrivateUsage]
        else:
            fastapi_app.state.engine = previous_engine
        if previous_session_maker is None:
            fastapi_app.state._state.pop("session_maker", None)  # pyright: ignore[reportPrivateUsage]
        else:
            fastapi_app.state.session_maker = previous_session_maker


@pytest.mark.parametrize("async_override", [False, True])
@pytest.mark.asyncio
async def test_get_client_uses_existing_local_asgi_database_override(
    config_manager,
    async_override,
):
    """Local ASGI routing honors FastAPI test dependency overrides."""
    from basic_memory.api.app import app as fastapi_app
    from basic_memory.deps import get_engine_factory

    cfg = config_manager.load_config()
    config_manager.save_config(cfg)

    previous_overrides = dict(fastapi_app.dependency_overrides)
    previous_engine = getattr(fastapi_app.state, "engine", None)
    previous_session_maker = getattr(fastapi_app.state, "session_maker", None)
    fastapi_app.state._state.pop("engine", None)  # pyright: ignore[reportPrivateUsage]
    fastapi_app.state._state.pop("session_maker", None)  # pyright: ignore[reportPrivateUsage]

    engine = object()
    session_maker = object()
    calls = []

    if async_override:

        async def override_engine_factory():
            calls.append("override")
            return engine, session_maker

    else:

        def override_engine_factory():
            calls.append("override")
            return engine, session_maker

    fastapi_app.dependency_overrides[get_engine_factory] = override_engine_factory

    try:
        async with get_client() as client:
            assert isinstance(client._transport, httpx.ASGITransport)  # pyright: ignore[reportPrivateUsage]
            assert calls == ["override"]
            assert fastapi_app.state.engine is engine
            assert fastapi_app.state.session_maker is session_maker
        assert not hasattr(fastapi_app.state, "engine")
        assert not hasattr(fastapi_app.state, "session_maker")
    finally:
        fastapi_app.dependency_overrides = previous_overrides
        if previous_engine is None:
            fastapi_app.state._state.pop("engine", None)  # pyright: ignore[reportPrivateUsage]
        else:
            fastapi_app.state.engine = previous_engine
        if previous_session_maker is None:
            fastapi_app.state._state.pop("session_maker", None)  # pyright: ignore[reportPrivateUsage]
        else:
            fastapi_app.state.session_maker = previous_session_maker


@pytest.mark.asyncio
async def test_get_client_resolves_local_asgi_database_override_with_fastapi_di(
    config_manager,
):
    """Local ASGI DB pre-init honors FastAPI dependency injection and cleanup."""
    from fastapi import Request

    from basic_memory.api.app import app as fastapi_app
    from basic_memory.deps import get_engine_factory

    cfg = config_manager.load_config()
    config_manager.save_config(cfg)

    previous_overrides = dict(fastapi_app.dependency_overrides)
    previous_engine = getattr(fastapi_app.state, "engine", None)
    previous_session_maker = getattr(fastapi_app.state, "session_maker", None)
    fastapi_app.state._state.pop("engine", None)  # pyright: ignore[reportPrivateUsage]
    fastapi_app.state._state.pop("session_maker", None)  # pyright: ignore[reportPrivateUsage]

    engine = object()
    session_maker = object()
    calls = []
    cleanup = []

    async def override_engine_factory(
        request: Request,
    ) -> AsyncIterator[tuple[object, object]]:
        calls.append(request.app)
        try:
            yield engine, session_maker
        finally:
            cleanup.append("closed")

    fastapi_app.dependency_overrides[get_engine_factory] = override_engine_factory

    try:
        async with get_client() as client:
            assert isinstance(client._transport, httpx.ASGITransport)  # pyright: ignore[reportPrivateUsage]
            assert calls == [fastapi_app]
            assert cleanup == []
            assert fastapi_app.state.engine is engine
            assert fastapi_app.state.session_maker is session_maker
        assert cleanup == ["closed"]
        assert not hasattr(fastapi_app.state, "engine")
        assert not hasattr(fastapi_app.state, "session_maker")
    finally:
        fastapi_app.dependency_overrides = previous_overrides
        if previous_engine is None:
            fastapi_app.state._state.pop("engine", None)  # pyright: ignore[reportPrivateUsage]
        else:
            fastapi_app.state.engine = previous_engine
        if previous_session_maker is None:
            fastapi_app.state._state.pop("session_maker", None)  # pyright: ignore[reportPrivateUsage]
        else:
            fastapi_app.state.session_maker = previous_session_maker


@pytest.mark.asyncio
async def test_get_client_keeps_local_asgi_database_during_overlapping_contexts(
    config_manager,
    monkeypatch,
):
    """Local ASGI database state stays installed until the last overlapping client exits."""
    from basic_memory import db
    from basic_memory.api.app import app as fastapi_app

    cfg = config_manager.load_config()
    config_manager.save_config(cfg)

    previous_engine = getattr(fastapi_app.state, "engine", None)
    previous_session_maker = getattr(fastapi_app.state, "session_maker", None)
    fastapi_app.state._state.pop("engine", None)  # pyright: ignore[reportPrivateUsage]
    fastapi_app.state._state.pop("session_maker", None)  # pyright: ignore[reportPrivateUsage]

    engine = object()
    session_maker = object()
    calls = []

    async def fake_get_or_create_db(db_path):
        calls.append(db_path)
        return engine, session_maker

    monkeypatch.setattr(db, "get_or_create_db", fake_get_or_create_db)

    first_context = get_client()
    second_context = get_client()
    first_entered = False
    second_entered = False
    first_exited = False

    try:
        first_client = await first_context.__aenter__()
        first_entered = True
        second_client = await second_context.__aenter__()
        second_entered = True

        assert isinstance(first_client._transport, httpx.ASGITransport)  # pyright: ignore[reportPrivateUsage]
        assert isinstance(second_client._transport, httpx.ASGITransport)  # pyright: ignore[reportPrivateUsage]
        assert calls == [cfg.database_path]

        await first_context.__aexit__(None, None, None)
        first_exited = True

        assert fastapi_app.state.engine is engine
        assert fastapi_app.state.session_maker is session_maker

        await second_context.__aexit__(None, None, None)
        second_entered = False

        assert not hasattr(fastapi_app.state, "engine")
        assert not hasattr(fastapi_app.state, "session_maker")
    finally:
        if second_entered:
            await second_context.__aexit__(None, None, None)
        if first_entered and not first_exited:
            await first_context.__aexit__(None, None, None)

        if previous_engine is None:
            fastapi_app.state._state.pop("engine", None)  # pyright: ignore[reportPrivateUsage]
        else:
            fastapi_app.state.engine = previous_engine
        if previous_session_maker is None:
            fastapi_app.state._state.pop("session_maker", None)  # pyright: ignore[reportPrivateUsage]
        else:
            fastapi_app.state.session_maker = previous_session_maker


@pytest.mark.asyncio
async def test_get_client_explicit_cloud_uses_api_key(config_manager, monkeypatch):
    cfg = config_manager.load_config()
    cfg.cloud_host = "https://cloud.example.test"
    cfg.cloud_api_key = "bmc_test_key_123"
    config_manager.save_config(cfg)

    monkeypatch.setenv("BASIC_MEMORY_FORCE_CLOUD", "true")
    monkeypatch.setenv("BASIC_MEMORY_EXPLICIT_ROUTING", "true")

    async with get_client() as client:
        assert str(client.base_url).rstrip("/") == "https://cloud.example.test/proxy"
        assert client.headers.get("Authorization") == "Bearer bmc_test_key_123"


@pytest.mark.asyncio
async def test_get_client_cloud_adds_workspace_header(config_manager):
    cfg = config_manager.load_config()
    cfg.cloud_host = "https://cloud.example.test"
    cfg.cloud_api_key = "bmc_test_key_123"
    cfg.set_project_mode("research", ProjectMode.CLOUD)
    config_manager.save_config(cfg)

    async with get_client(project_name="research", workspace="tenant-123") as client:
        assert str(client.base_url).rstrip("/") == "https://cloud.example.test/proxy"
        assert client.headers.get("X-Workspace-ID") == "tenant-123"


@pytest.mark.asyncio
async def test_get_client_cloud_uses_project_workspace_when_not_explicit(config_manager):
    cfg = config_manager.load_config()
    cfg.cloud_host = "https://cloud.example.test"
    cfg.cloud_api_key = "bmc_test_key_123"
    cfg.default_workspace = "default-tenant"
    cfg.set_project_mode("research", ProjectMode.CLOUD)
    cfg.projects["research"].workspace_id = "project-tenant"
    config_manager.save_config(cfg)

    async with get_client(project_name="research") as client:
        assert str(client.base_url).rstrip("/") == "https://cloud.example.test/proxy"
        assert client.headers.get("X-Workspace-ID") == "project-tenant"


@pytest.mark.asyncio
async def test_get_client_cloud_uses_default_workspace_when_project_has_none(config_manager):
    cfg = config_manager.load_config()
    cfg.cloud_host = "https://cloud.example.test"
    cfg.cloud_api_key = "bmc_test_key_123"
    cfg.default_workspace = "default-tenant"
    cfg.set_project_mode("research", ProjectMode.CLOUD)
    config_manager.save_config(cfg)

    async with get_client(project_name="research") as client:
        assert str(client.base_url).rstrip("/") == "https://cloud.example.test/proxy"
        assert client.headers.get("X-Workspace-ID") == "default-tenant"


@pytest.mark.asyncio
async def test_get_client_explicit_cloud_raises_without_credentials(config_manager, monkeypatch):
    cfg = config_manager.load_config()
    cfg.cloud_host = "https://cloud.example.test"
    cfg.cloud_api_key = None
    cfg.cloud_client_id = "cid"
    cfg.cloud_domain = "https://auth.example.test"
    config_manager.save_config(cfg)

    monkeypatch.setenv("BASIC_MEMORY_FORCE_CLOUD", "true")
    monkeypatch.setenv("BASIC_MEMORY_EXPLICIT_ROUTING", "true")

    with pytest.raises(RuntimeError, match="Cloud routing requested but no credentials found"):
        async with get_client():
            pass


@pytest.mark.asyncio
async def test_get_client_per_project_cloud_uses_api_key(config_manager):
    """Cloud-mode project routes through cloud with API key auth."""
    cfg = config_manager.load_config()
    cfg.cloud_host = "https://cloud.example.test"
    cfg.cloud_api_key = "bmc_test_key_123"
    cfg.set_project_mode("research", ProjectMode.CLOUD)
    config_manager.save_config(cfg)

    async with get_client(project_name="research") as client:
        assert str(client.base_url).rstrip("/") == "https://cloud.example.test/proxy"
        assert client.headers.get("Authorization") == "Bearer bmc_test_key_123"


@pytest.mark.asyncio
async def test_get_client_per_project_cloud_raises_without_credentials(config_manager):
    """Cloud-mode project raises with actionable auth guidance when no credentials exist."""
    cfg = config_manager.load_config()
    cfg.cloud_api_key = None
    cfg.set_project_mode("research", ProjectMode.CLOUD)
    config_manager.save_config(cfg)

    with pytest.raises(RuntimeError, match="Project 'research' is set to cloud mode"):
        async with get_client(project_name="research"):
            pass


@pytest.mark.asyncio
async def test_get_client_local_project_uses_asgi_transport(config_manager):
    """Local-mode project uses ASGI transport even if API key exists."""
    cfg = config_manager.load_config()
    cfg.cloud_api_key = "bmc_test_key_123"
    cfg.set_project_mode("main", ProjectMode.LOCAL)
    config_manager.save_config(cfg)

    async with get_client(project_name="main") as client:
        assert isinstance(client._transport, httpx.ASGITransport)  # pyright: ignore[reportPrivateUsage]


@pytest.mark.asyncio
async def test_get_client_no_project_name_defaults_local(config_manager):
    """No project_name defaults to local ASGI routing."""
    cfg = config_manager.load_config()
    cfg.cloud_api_key = "bmc_test_key_123"
    cfg.set_project_mode("research", ProjectMode.CLOUD)
    config_manager.save_config(cfg)

    async with get_client() as client:
        assert isinstance(client._transport, httpx.ASGITransport)  # pyright: ignore[reportPrivateUsage]


@pytest.mark.asyncio
async def test_get_client_factory_overrides_per_project_routing(config_manager):
    """Injected factory takes priority over per-project routing."""
    cfg = config_manager.load_config()
    cfg.cloud_api_key = "bmc_test_key_123"
    cfg.set_project_mode("research", ProjectMode.CLOUD)
    config_manager.save_config(cfg)

    @asynccontextmanager
    async def factory(workspace=None):
        async with httpx.AsyncClient(base_url="https://factory.test") as client:
            yield client

    set_client_factory(factory)

    async with get_client(project_name="research") as client:
        assert str(client.base_url) == "https://factory.test"


@pytest.mark.asyncio
async def test_get_client_force_local_without_explicit_does_not_override_project_mode(
    config_manager, monkeypatch
):
    """FORCE_LOCAL alone should not bypass per-project cloud routing."""
    cfg = config_manager.load_config()
    cfg.cloud_host = "https://cloud.example.test"
    cfg.cloud_api_key = "bmc_test_key_123"
    cfg.set_project_mode("research", ProjectMode.CLOUD)
    config_manager.save_config(cfg)

    monkeypatch.setenv("BASIC_MEMORY_FORCE_LOCAL", "true")

    async with get_client(project_name="research") as client:
        assert str(client.base_url).rstrip("/") == "https://cloud.example.test/proxy"


@pytest.mark.asyncio
async def test_get_client_explicit_local_overrides_cloud_project(config_manager, monkeypatch):
    """EXPLICIT_ROUTING + FORCE_LOCAL should override a cloud project to local ASGI."""
    cfg = config_manager.load_config()
    cfg.cloud_host = "https://cloud.example.test"
    cfg.cloud_api_key = "bmc_test_key_123"
    cfg.set_project_mode("research", ProjectMode.CLOUD)
    config_manager.save_config(cfg)

    monkeypatch.setenv("BASIC_MEMORY_FORCE_LOCAL", "true")
    monkeypatch.setenv("BASIC_MEMORY_EXPLICIT_ROUTING", "true")

    async with get_client(project_name="research") as client:
        assert isinstance(client._transport, httpx.ASGITransport)  # pyright: ignore[reportPrivateUsage]


@pytest.mark.asyncio
async def test_get_client_per_project_cloud_oauth_fallback(config_manager):
    """Cloud-mode project uses OAuth token when no API key is configured."""
    cfg = config_manager.load_config()
    cfg.cloud_host = "https://cloud.example.test"
    cfg.cloud_api_key = None
    cfg.cloud_client_id = "cid"
    cfg.cloud_domain = "https://auth.example.test"
    cfg.set_project_mode("research", ProjectMode.CLOUD)
    config_manager.save_config(cfg)

    # Write OAuth token file so CLIAuth.get_valid_token() returns it
    auth = CLIAuth(client_id=cfg.cloud_client_id, authkit_domain=cfg.cloud_domain)
    auth.token_file.parent.mkdir(parents=True, exist_ok=True)
    auth.token_file.write_text(
        '{"access_token":"oauth-token-456","refresh_token":null,"expires_at":9999999999,"token_type":"Bearer"}',
        encoding="utf-8",
    )

    async with get_client(project_name="research") as client:
        assert str(client.base_url).rstrip("/") == "https://cloud.example.test/proxy"
        assert client.headers.get("Authorization") == "Bearer oauth-token-456"


@pytest.mark.asyncio
async def test_get_client_explicit_cloud_overrides_local_project(config_manager, monkeypatch):
    """EXPLICIT_ROUTING + FORCE_CLOUD should override a local project to cloud."""
    cfg = config_manager.load_config()
    cfg.cloud_host = "https://cloud.example.test"
    cfg.cloud_api_key = "bmc_test_key_123"
    config_manager.save_config(cfg)

    monkeypatch.setenv("BASIC_MEMORY_FORCE_CLOUD", "true")
    monkeypatch.setenv("BASIC_MEMORY_EXPLICIT_ROUTING", "true")

    async with get_client(project_name="main") as client:
        assert str(client.base_url).rstrip("/") == "https://cloud.example.test/proxy"
        assert client.headers.get("Authorization") == "Bearer bmc_test_key_123"


@pytest.mark.asyncio
async def test_get_cloud_control_plane_client_uses_api_key_when_available(config_manager):
    cfg = config_manager.load_config()
    cfg.cloud_host = "https://cloud.example.test"
    cfg.cloud_api_key = "bmc_test_key_123"
    cfg.cloud_client_id = "cid"
    cfg.cloud_domain = "https://auth.example.test"
    config_manager.save_config(cfg)

    async with get_cloud_control_plane_client() as client:
        assert str(client.base_url).rstrip("/") == "https://cloud.example.test"
        assert client.headers.get("Authorization") == "Bearer bmc_test_key_123"


@pytest.mark.asyncio
async def test_get_cloud_control_plane_client_adds_workspace_header(config_manager):
    cfg = config_manager.load_config()
    cfg.cloud_host = "https://cloud.example.test"
    cfg.cloud_api_key = "bmc_test_key_123"
    config_manager.save_config(cfg)

    async with get_cloud_control_plane_client(workspace="tenant-123") as client:
        assert str(client.base_url).rstrip("/") == "https://cloud.example.test"
        assert client.headers.get("Authorization") == "Bearer bmc_test_key_123"
        assert client.headers.get("X-Workspace-ID") == "tenant-123"


@pytest.mark.asyncio
async def test_get_cloud_control_plane_client_uses_oauth_token(config_manager):
    cfg = config_manager.load_config()
    cfg.cloud_host = "https://cloud.example.test"
    cfg.cloud_api_key = None
    cfg.cloud_client_id = "cid"
    cfg.cloud_domain = "https://auth.example.test"
    config_manager.save_config(cfg)

    auth = CLIAuth(client_id=cfg.cloud_client_id, authkit_domain=cfg.cloud_domain)
    auth.token_file.parent.mkdir(parents=True, exist_ok=True)
    auth.token_file.write_text(
        '{"access_token":"oauth-control-123","refresh_token":null,"expires_at":9999999999,"token_type":"Bearer"}',
        encoding="utf-8",
    )

    async with get_cloud_control_plane_client() as client:
        assert str(client.base_url).rstrip("/") == "https://cloud.example.test"
        assert client.headers.get("Authorization") == "Bearer oauth-control-123"


@pytest.mark.asyncio
async def test_get_cloud_control_plane_client_with_workspace(config_manager):
    """Control plane client passes X-Workspace-ID header when workspace is provided."""
    cfg = config_manager.load_config()
    cfg.cloud_host = "https://cloud.example.test"
    cfg.cloud_api_key = "bmc_test_key_123"
    config_manager.save_config(cfg)

    async with get_cloud_control_plane_client(workspace="tenant-abc") as client:
        assert client.headers.get("X-Workspace-ID") == "tenant-abc"

    # Without workspace, header should not be present
    async with get_cloud_control_plane_client() as client:
        assert "X-Workspace-ID" not in client.headers


@pytest.mark.asyncio
async def test_get_client_auto_resolves_workspace_from_project_config(config_manager):
    """get_client resolves workspace from project entry when not explicitly passed."""
    cfg = config_manager.load_config()
    cfg.cloud_host = "https://cloud.example.test"
    cfg.cloud_api_key = "bmc_test_key_123"
    cfg.set_project_mode("research", ProjectMode.CLOUD)
    cfg.projects["research"].workspace_id = "tenant-from-config"
    config_manager.save_config(cfg)

    async with get_client(project_name="research") as client:
        assert client.headers.get("X-Workspace-ID") == "tenant-from-config"


@pytest.mark.asyncio
async def test_get_client_auto_resolves_workspace_from_default(config_manager):
    """get_client falls back to default_workspace when project has no workspace_id."""
    cfg = config_manager.load_config()
    cfg.cloud_host = "https://cloud.example.test"
    cfg.cloud_api_key = "bmc_test_key_123"
    cfg.set_project_mode("research", ProjectMode.CLOUD)
    cfg.default_workspace = "default-tenant-456"
    config_manager.save_config(cfg)

    async with get_client(project_name="research") as client:
        assert client.headers.get("X-Workspace-ID") == "default-tenant-456"


@pytest.mark.asyncio
async def test_get_client_explicit_workspace_overrides_config(config_manager):
    """Explicit workspace param takes priority over project config."""
    cfg = config_manager.load_config()
    cfg.cloud_host = "https://cloud.example.test"
    cfg.cloud_api_key = "bmc_test_key_123"
    cfg.set_project_mode("research", ProjectMode.CLOUD)
    cfg.projects["research"].workspace_id = "tenant-from-config"
    config_manager.save_config(cfg)

    async with get_client(project_name="research", workspace="explicit-tenant") as client:
        assert client.headers.get("X-Workspace-ID") == "explicit-tenant"


@pytest.mark.asyncio
async def test_get_cloud_control_plane_client_raises_without_credentials(config_manager):
    cfg = config_manager.load_config()
    cfg.cloud_host = "https://cloud.example.test"
    cfg.cloud_api_key = None
    cfg.cloud_client_id = "cid"
    cfg.cloud_domain = "https://auth.example.test"
    config_manager.save_config(cfg)

    with pytest.raises(RuntimeError, match="Cloud routing requested but no credentials found"):
        async with get_cloud_control_plane_client():
            pass