"""
tests/test_llm_budgets.py — Tests for per-session LLM call budgets (Issue #90).

Covers:
  - Budget exhaustion blocks extraction gracefully (no crash, ok=False)
  - Budget exhaustion blocks reconciliation gracefully
  - Warning logged at 80% threshold
  - Budget resets on wake()
  - max_extraction_calls_per_batch caps claim fan-out
  - Budget disabled when max_llm_calls_per_session == 0
  - Config fields parse correctly from defaults.toml
"""
from __future__ import annotations

import asyncio
import logging
from dataclasses import dataclass, field
from typing import Any, List, Optional
from unittest.mock import AsyncMock, MagicMock, patch

import pytest


# ---------------------------------------------------------------------------
# Minimal stub types (mirrors test_plugin.py stubs)
# ---------------------------------------------------------------------------

@dataclass
class _EntityMention:
    name: str
    entity_type: str = "person"


@dataclass
class _ClaimExtraction:
    action: str = "prefers"
    subject: str = "user"
    predicate: str = "likes"
    object_value: str = "cats"
    confidence: float = 0.9
    temporal: str = "current"
    entities_mentioned: List[_EntityMention] = field(default_factory=list)
    sensitivity: str = "normal"


@dataclass
class _ExtractionResult:
    claims: List[_ClaimExtraction]
    noise_filtered: int = 0
    raw_message_count: int = 1


class _ReconciliationAction:
    def __init__(self, value: str):
        self.value = value


@dataclass
class _ReconciliationResult:
    candidate_subject: str = "user"
    candidate_predicate: str = "likes"
    candidate_value: str = "cats"
    action: Any = None

    def __post_init__(self):
        if self.action is None:
            self.action = _ReconciliationAction("ADD")


@dataclass
class _ReconciliationReport:
    results: List[_ReconciliationResult] = field(default_factory=list)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _make_config(**overrides) -> MagicMock:
    cfg = MagicMock()
    cfg.storage_db_path = ":memory:"
    cfg.storage_enable_wal = False
    cfg.storage_busy_timeout_ms = 5000
    cfg.extraction_model = "gpt-4.1-nano"
    cfg.reconciliation_model = "gpt-4.1-nano"
    cfg.retrieval_token_budget = 3000
    cfg.total_memory_token_budget = 4000
    cfg.owner_id = "default"
    cfg.cornerstone_token_budget = 800
    cfg.retrieval_default_strategy = "hybrid"
    cfg.retrieval_max_tokens = 4000
    cfg.retrieval_vector_weight = 0.7
    cfg.retrieval_bm25_weight = 0.3
    cfg.coding_noise_filter = True
    cfg.cornerstone_drift_threshold = 0.05
    cfg.max_cornerstone_count = 10
    cfg.dream_cycle_enabled = False
    cfg.max_llm_calls_per_session = overrides.pop("max_llm_calls_per_session", 1000)
    cfg.max_extraction_calls_per_batch = overrides.pop("max_extraction_calls_per_batch", 10)
    for k, v in overrides.items():
        setattr(cfg, k, v)
    return cfg


def _make_plugin(config=None, **module_overrides) -> Any:
    """Build a CortexPlugin with mocked modules, bypassing __init__."""
    from cortex.api.plugin import CortexPlugin

    p = object.__new__(CortexPlugin)
    p._init_errors = []
    p._session_id = None
    p._llm_call_count = 0
    p.config = config or _make_config()

    storage = MagicMock()
    storage.create_session = AsyncMock(return_value=MagicMock(id="db-pk-1"))
    storage.close_session = AsyncMock()
    storage.ping = AsyncMock(return_value=True)

    p.storage = module_overrides.get("storage", storage)
    p.llm = module_overrides.get("llm", MagicMock())
    p.extractor = module_overrides.get("extractor", None)
    p.entity_resolver = module_overrides.get("entity_resolver", None)
    p.reconciler = module_overrides.get("reconciler", None)
    p.cornerstone_guardian = module_overrides.get("cornerstone_guardian", None)
    p.retrieval = module_overrides.get("retrieval", None)
    p.token_budget_manager = module_overrides.get("token_budget_manager", None)
    p.feedback_system = module_overrides.get("feedback_system", None)
    p.event_emitter = module_overrides.get("event_emitter", None)
    p._webhook_dispatcher = None
    return p


def _make_extractor(claims: Optional[List[_ClaimExtraction]] = None) -> MagicMock:
    ext = MagicMock()
    ext.extract = AsyncMock(return_value=_ExtractionResult(
        claims=claims or [_ClaimExtraction()],
    ))
    return ext


def _make_reconciler(n_claims: int = 1) -> MagicMock:
    rec = MagicMock()
    report = _ReconciliationReport(results=[_ReconciliationResult() for _ in range(n_claims)])
    rec.reconcile = AsyncMock(return_value=report)
    return rec


# ---------------------------------------------------------------------------
# Tests: _budget_check / _budget_consume
# ---------------------------------------------------------------------------

class TestBudgetHelpers:

    def test_budget_check_passes_within_limit(self):
        plugin = _make_plugin()
        plugin._llm_call_count = 0
        assert plugin._budget_check(1) is None

    def test_budget_check_fails_when_at_limit(self):
        plugin = _make_plugin(config=_make_config(max_llm_calls_per_session=5))
        plugin._llm_call_count = 5
        err = plugin._budget_check(1)
        assert err is not None
        assert "budget exhausted" in err.lower()

    def test_budget_check_fails_when_partial_overage(self):
        """3 calls needed but only 2 remain → should fail."""
        plugin = _make_plugin(config=_make_config(max_llm_calls_per_session=10))
        plugin._llm_call_count = 8
        err = plugin._budget_check(3)
        assert err is not None

    def test_budget_check_passes_when_exactly_enough(self):
        plugin = _make_plugin(config=_make_config(max_llm_calls_per_session=10))
        plugin._llm_call_count = 7
        assert plugin._budget_check(3) is None

    def test_budget_disabled_when_limit_zero(self):
        plugin = _make_plugin(config=_make_config(max_llm_calls_per_session=0))
        plugin._llm_call_count = 9999
        assert plugin._budget_check(1) is None  # no limit

    def test_budget_consume_increments_counter(self):
        plugin = _make_plugin()
        plugin._llm_call_count = 0
        plugin._budget_consume(5)
        assert plugin._llm_call_count == 5

    def test_budget_consume_warns_at_80_percent(self, caplog):
        plugin = _make_plugin(config=_make_config(max_llm_calls_per_session=10))
        plugin._llm_call_count = 0
        with caplog.at_level(logging.WARNING, logger="cortex.api.plugin"):
            plugin._budget_consume(8)  # 8/10 = 80%
        assert any("80%" in r.message or "budget" in r.message.lower() for r in caplog.records)

    def test_budget_consume_no_warn_below_80_percent(self, caplog):
        plugin = _make_plugin(config=_make_config(max_llm_calls_per_session=10))
        plugin._llm_call_count = 0
        with caplog.at_level(logging.WARNING, logger="cortex.api.plugin"):
            plugin._budget_consume(7)  # 7/10 = 70% — no warning yet
        budget_warns = [r for r in caplog.records if "budget" in r.message.lower() and r.levelno >= logging.WARNING]
        assert not budget_warns

    def test_budget_consume_no_warn_when_disabled(self, caplog):
        plugin = _make_plugin(config=_make_config(max_llm_calls_per_session=0))
        plugin._llm_call_count = 9999
        with caplog.at_level(logging.WARNING, logger="cortex.api.plugin"):
            plugin._budget_consume(1)
        # No warning for disabled budget
        budget_warns = [r for r in caplog.records if "budget" in r.message.lower()]
        assert not budget_warns


# ---------------------------------------------------------------------------
# Tests: remember() budget enforcement
# ---------------------------------------------------------------------------

class TestRememberBudget:

    @pytest.mark.asyncio
    async def test_budget_exhausted_before_extraction_returns_error(self):
        """When budget is fully consumed, remember() must return ok=False, no crash."""
        plugin = _make_plugin(
            config=_make_config(max_llm_calls_per_session=5),
            extractor=_make_extractor(),
        )
        plugin._llm_call_count = 5  # already at limit

        result = await plugin.remember(["user: I like dogs"])
        assert result.ok is False
        assert any("budget exhausted" in e.lower() for e in result.errors)
        assert result.claims_extracted == 0

    @pytest.mark.asyncio
    async def test_extraction_call_counted(self):
        """One remember() call should consume at least 1 LLM call (extraction)."""
        plugin = _make_plugin(
            extractor=_make_extractor(),
            entity_resolver=None,
            reconciler=_make_reconciler(1),
        )
        assert plugin._llm_call_count == 0
        await plugin.remember(["user: I like cats"])
        assert plugin._llm_call_count >= 1

    @pytest.mark.asyncio
    async def test_reconciliation_skipped_when_budget_too_low(self):
        """If extraction consumes the last budget slot, reconciliation is skipped."""
        plugin = _make_plugin(
            config=_make_config(max_llm_calls_per_session=1),  # only 1 call total
            extractor=_make_extractor([_ClaimExtraction(), _ClaimExtraction()]),
            reconciler=_make_reconciler(2),
        )
        result = await plugin.remember(["user: I like cats"])
        # Extraction consumed the 1 slot, reconciliation must have been skipped
        assert plugin.reconciler.reconcile.call_count == 0
        assert any("budget exhausted" in e.lower() for e in result.errors)

    @pytest.mark.asyncio
    async def test_successful_remember_below_budget(self):
        """Normal happy-path: budget not exceeded, all modules called."""
        reconciler = _make_reconciler(1)
        entity_resolver = MagicMock()
        entity_resolver.resolve = AsyncMock(return_value=MagicMock(entity_id="ent-1"))
        plugin = _make_plugin(
            config=_make_config(max_llm_calls_per_session=50),
            extractor=_make_extractor([_ClaimExtraction()]),
            entity_resolver=entity_resolver,
            reconciler=reconciler,
        )
        result = await plugin.remember(["user: I like cats"])
        assert result.ok is True
        assert reconciler.reconcile.call_count == 1

    @pytest.mark.asyncio
    async def test_budget_not_exceeded_across_multiple_calls(self):
        """Multiple remember() calls accumulate budget usage."""
        plugin = _make_plugin(
            config=_make_config(max_llm_calls_per_session=10),
            extractor=_make_extractor([_ClaimExtraction()]),
            reconciler=_make_reconciler(1),
        )
        # 2 calls × (1 extraction + 1 reconciliation) = 4 total
        await plugin.remember(["msg 1"])
        await plugin.remember(["msg 2"])
        assert plugin._llm_call_count == 4


# ---------------------------------------------------------------------------
# Tests: max_extraction_calls_per_batch claim capping
# ---------------------------------------------------------------------------

class TestExtractionBatchCap:

    @pytest.mark.asyncio
    async def test_claims_capped_at_batch_limit(self):
        """If extractor returns 15 claims but limit=10, only 10 passed to reconciler."""
        claims = [_ClaimExtraction(object_value=str(i)) for i in range(15)]
        reconciler = _make_reconciler(10)
        plugin = _make_plugin(
            config=_make_config(max_extraction_calls_per_batch=10, max_llm_calls_per_session=200),
            extractor=_make_extractor(claims),
            reconciler=reconciler,
        )
        result = await plugin.remember(["big conversation"])
        # Reconciler was called with at most 10 candidates
        call_args = reconciler.reconcile.call_args
        assert len(call_args.kwargs.get("candidates", call_args.args[0] if call_args.args else [])) <= 10

    @pytest.mark.asyncio
    async def test_claims_not_capped_below_limit(self):
        """If extractor returns 5 claims and limit=10, all 5 pass through."""
        claims = [_ClaimExtraction(object_value=str(i)) for i in range(5)]
        reconciler = _make_reconciler(5)
        plugin = _make_plugin(
            config=_make_config(max_extraction_calls_per_batch=10, max_llm_calls_per_session=200),
            extractor=_make_extractor(claims),
            reconciler=reconciler,
        )
        await plugin.remember(["small conversation"])
        call_args = reconciler.reconcile.call_args
        candidates = call_args.kwargs.get("candidates", call_args.args[0] if call_args.args else [])
        assert len(candidates) == 5

    @pytest.mark.asyncio
    async def test_batch_cap_zero_means_no_cap(self):
        """max_extraction_calls_per_batch=0 disables the cap."""
        claims = [_ClaimExtraction(object_value=str(i)) for i in range(20)]
        reconciler = _make_reconciler(20)
        plugin = _make_plugin(
            config=_make_config(max_extraction_calls_per_batch=0, max_llm_calls_per_session=500),
            extractor=_make_extractor(claims),
            reconciler=reconciler,
        )
        await plugin.remember(["huge conversation"])
        call_args = reconciler.reconcile.call_args
        candidates = call_args.kwargs.get("candidates", call_args.args[0] if call_args.args else [])
        assert len(candidates) == 20


# ---------------------------------------------------------------------------
# Tests: wake() resets budget
# ---------------------------------------------------------------------------

class TestWakeBudgetReset:

    @pytest.mark.asyncio
    async def test_wake_resets_llm_call_count(self):
        """wake() must zero out _llm_call_count."""
        plugin = _make_plugin()
        plugin._llm_call_count = 42

        cornerstone_guardian = MagicMock()
        cornerstone_guardian.load_all = AsyncMock(return_value=[])
        plugin.cornerstone_guardian = cornerstone_guardian

        await plugin.wake(session_id="new-session")
        assert plugin._llm_call_count == 0

    @pytest.mark.asyncio
    async def test_budget_allows_calls_after_wake(self):
        """After waking, a previously-exhausted budget is restored."""
        reconciler = _make_reconciler(1)
        plugin = _make_plugin(
            config=_make_config(max_llm_calls_per_session=5),
            extractor=_make_extractor([_ClaimExtraction()]),
            reconciler=reconciler,
        )
        plugin._llm_call_count = 5  # exhaust budget

        # Without wake — should fail
        result = await plugin.remember(["msg"])
        assert result.ok is False

        # After wake — should succeed
        plugin.cornerstone_guardian = MagicMock()
        plugin.cornerstone_guardian.load_all = AsyncMock(return_value=[])
        entity_resolver = MagicMock()
        entity_resolver.resolve = AsyncMock(return_value=MagicMock(entity_id="ent-1"))
        plugin.entity_resolver = entity_resolver
        await plugin.wake()

        result2 = await plugin.remember(["msg after wake"])
        assert result2.ok is True


# ---------------------------------------------------------------------------
# Tests: config parsing
# ---------------------------------------------------------------------------

class TestBudgetConfigParsing:

    def test_defaults_toml_has_budget_keys(self):
        """defaults.toml must include max_llm_calls_per_session + max_extraction_calls_per_batch."""
        from cortex.config import CortexConfig
        cfg = CortexConfig.defaults()
        assert cfg.max_llm_calls_per_session == 1000
        assert cfg.max_extraction_calls_per_batch == 25

    def test_overlay_overrides_budget(self):
        """In-process overlay can override budget values."""
        from cortex.config import CortexConfig
        cfg = CortexConfig.from_toml(
            overlay={"cortex": {"budgets": {
                "max_llm_calls_per_session": 100,
                "max_extraction_calls_per_batch": 5,
            }}}
        )
        assert cfg.max_llm_calls_per_session == 100
        assert cfg.max_extraction_calls_per_batch == 5

    def test_default_config_dataclass_values(self):
        """CortexConfig() with no args uses the right defaults."""
        from cortex.config import CortexConfig
        cfg = CortexConfig()
        assert cfg.max_llm_calls_per_session == 1000
        assert cfg.max_extraction_calls_per_batch == 25
