"""Tests for workspace cost gate."""

import pytest

from investibook.core import session_cost as sc
from investibook.core.session_cost import WorkspaceCostBlocked
from investibook.db.queries import _db_workspace_has_cost_bypass, _db_workspace_total_recorded_cost


@pytest.mark.asyncio
async def test_effective_prompt_bypasses_when_workspace_flag(monkeypatch: pytest.MonkeyPatch) -> None:
    async def fake_to_thread(fn, *a, **kw):
        if fn is _db_workspace_has_cost_bypass:
            return True
        raise AssertionError(f"unexpected {fn!r}")

    monkeypatch.setattr(sc.asyncio, "to_thread", fake_to_thread)
    monkeypatch.setattr(sc, "WORKSPACE_COST_LIMIT_USD", 1.0)

    blocked = await sc.effective_prompt_after_cost_gate("ws-1", "user-1")
    assert blocked is None


@pytest.mark.asyncio
async def test_effective_prompt_blocks_when_over_limit(monkeypatch: pytest.MonkeyPatch) -> None:
    async def fake_to_thread(fn, *a, **kw):
        if fn is _db_workspace_has_cost_bypass:
            return False
        if fn is _db_workspace_total_recorded_cost:
            return 10.0
        raise AssertionError(f"unexpected {fn!r}")

    monkeypatch.setattr(sc.asyncio, "to_thread", fake_to_thread)
    monkeypatch.setattr(sc, "WORKSPACE_COST_LIMIT_USD", 5.0)

    blocked = await sc.effective_prompt_after_cost_gate("ws-1", "user-1")
    assert isinstance(blocked, WorkspaceCostBlocked)
    assert blocked.usage_usd == 10.0 and blocked.limit_usd == 5.0
    msg = blocked.as_user_message()
    assert "limit" in msg.lower() and "/force-continue" in msg