"""Provider shutdown joins every plugin thread within one budget (#37632, #33485, #60616)."""

import logging
import threading
import time
from types import SimpleNamespace
from unittest.mock import MagicMock

import pytest

from plugins.memory.honcho import HonchoMemoryProvider
from plugins.memory.honcho import session as session_module
from plugins.memory.honcho.client import (
    HonchoClientConfig,
    close_honcho_clients,
    join_plugin_threads,
    spawn_context_thread,
)
from plugins.memory.honcho.client_cache import _client_slots, _client_slots_lock
from plugins.memory.honcho.session import HonchoSession, HonchoSessionManager
from plugins.plugin_utils import SingletonSlot


@pytest.fixture
def async_manager(monkeypatch):
    fake = MagicMock()
    monkeypatch.setattr(session_module, "get_honcho_client", lambda *a, **k: fake)
    cfg = HonchoClientConfig(write_frequency="async", api_key="test-key", enabled=True)
    mgr = HonchoSessionManager(honcho=fake, config=cfg)
    mgr.fake_client = fake
    yield mgr
    mgr.shutdown()


class _Owner:
    pass


class TestThreadRegistry:
    def test_join_covers_only_the_given_owners(self):
        mine, theirs = _Owner(), _Owner()
        release = threading.Event()
        own = spawn_context_thread(lambda: release.wait(timeout=2), name="own", owner=mine)
        other = spawn_context_thread(lambda: release.wait(timeout=2), name="other", owner=theirs)
        own.start()
        other.start()
        try:
            started = time.monotonic()
            assert join_plugin_threads((theirs, None), timeout=0.05) == ["other"]
            release.set()
            assert join_plugin_threads((mine,), timeout=2) == []
            assert time.monotonic() - started < 2.0
        finally:
            release.set()
            own.join(timeout=1)
            other.join(timeout=1)


class TestProviderShutdown:
    def _provider(self, manager, cfg=None):
        provider = HonchoMemoryProvider()
        provider._manager = manager
        provider._config = cfg or manager._config
        provider._session_key = "test-session"
        provider._session_initialized = True
        return provider

    def _shutdown_blocks_until(self, provider, release):
        """shutdown() off-thread must still be running until ``release`` is set, then finish."""
        done = threading.Event()
        threading.Thread(target=lambda: (provider.shutdown(), done.set()), daemon=True).start()
        try:
            assert not done.wait(timeout=0.05)
            release.set()
            assert done.wait(timeout=2)
        finally:
            release.set()

    def test_shutdown_waits_for_an_in_flight_context_prefetch(self, async_manager):
        provider = self._provider(async_manager)
        started, release = threading.Event(), threading.Event()

        def slow_prefetch(session_key, user_message=None):
            started.set()
            release.wait(timeout=2)
            return {"representation": "ready"}

        async_manager.get_prefetch_context = slow_prefetch
        async_manager.prefetch_context("test-session", "query")
        assert started.wait(timeout=1)
        self._shutdown_blocks_until(provider, release)
        assert not any(t.name == "honcho-context-prefetch" and t.is_alive() for t in threading.enumerate())

    def test_shutdown_joins_a_stalled_init_thread_within_the_budget(self, monkeypatch, caplog):
        monkeypatch.setattr(HonchoMemoryProvider, "_SHUTDOWN_JOIN_FLOOR", 0.2)
        cfg = HonchoClientConfig(api_key="test-key", enabled=True, timeout=0.1)
        release = threading.Event()
        entered = threading.Event()

        class StalledManager:
            def __init__(self, *args, **kwargs):
                pass

            def get_or_create(self, session_key):
                entered.set()
                release.wait(timeout=5)
                return SimpleNamespace(messages=[])

            def migrate_memory_files(self, *a, **k):
                pass

        monkeypatch.setattr("plugins.memory.honcho.client.HonchoClientConfig.from_global_config", lambda: cfg)
        monkeypatch.setattr("plugins.memory.honcho.client.get_honcho_client", lambda cfg: object())
        monkeypatch.setattr("plugins.memory.honcho.session.HonchoSessionManager", StalledManager)
        provider = HonchoMemoryProvider()
        provider.initialize("session-1", platform="cli")
        try:
            assert entered.wait(timeout=1)
            started = time.monotonic()
            with caplog.at_level(logging.WARNING, logger="plugins.memory.honcho"):
                provider.shutdown()
            assert 0.15 <= time.monotonic() - started < 2.0
            assert "honcho-session-init" in caplog.text
            assert "timed out after 0.2s" in caplog.text
        finally:
            release.set()
            provider._init_thread.join(timeout=2)

    def test_shutdown_joins_the_prewarm_dialectic_thread(self, async_manager):
        provider = self._provider(async_manager)
        release = threading.Event()
        async_manager.dialectic_query = lambda *a, **k: release.wait(timeout=2) and "answer"
        provider._spawn_dialectic("who?", thread_name="honcho-prewarm-dialectic", fired_at=0, log_label="prewarm")
        self._shutdown_blocks_until(provider, release)
        assert not provider._prefetch_thread.is_alive()

    def test_recall_sync_worker_is_registered_under_the_provider(self, async_manager):
        from plugins.memory.honcho.recall_sync import prefetch_sync

        provider = self._provider(async_manager, cfg=HonchoClientConfig(api_key="test-key", enabled=True, timeout=0.05))
        entered, release = threading.Event(), threading.Event()

        def blocked(session_key, query, **kwargs):
            entered.set()
            release.wait(timeout=2)
            return {}

        async_manager.get_prefetch_context = blocked
        provider.on_turn_start(1, "what did we decide about the schema?")
        try:
            assert prefetch_sync(provider, "what did we decide about the schema?") == ""
            assert entered.wait(timeout=1)
            assert join_plugin_threads((provider,), timeout=0.01) == ["honcho-recall-sync"]
            release.set()
            assert join_plugin_threads((provider,), timeout=2) == []
        finally:
            release.set()
            provider._recall_sync_thread.join(timeout=2)

    def test_shutdown_skips_a_flush_whose_lock_an_upload_holds_past_the_budget(self, async_manager, monkeypatch, caplog):
        provider = self._provider(async_manager)
        monkeypatch.setattr(provider, "_shutdown_join_budget", lambda: 0.2)
        session = HonchoSession("test-session", "user", "assistant", "sid")
        session.add_message("user", "pending")
        async_manager._cache["test-session"] = session
        entered, release = threading.Event(), threading.Event()

        def upload(messages):
            entered.set()
            release.wait(timeout=5)

        async_manager._sessions_cache["sid"] = SimpleNamespace(add_messages=upload)
        async_manager.save(session)
        try:
            assert entered.wait(timeout=1)
            started = time.monotonic()
            with caplog.at_level(logging.WARNING, logger="plugins.memory.honcho"):
                provider.shutdown()
            assert time.monotonic() - started < 2.0
            assert "1 message(s) in 1 session(s) still unsynced" in caplog.text
            assert session.messages[0].get("_synced") is None
        finally:
            release.set()

    def test_shutdown_does_not_close_the_shared_client(self, async_manager):
        provider = self._provider(async_manager)
        provider.shutdown()
        async_manager.fake_client._http.close.assert_not_called()


@pytest.mark.parametrize("timeout, budget", [(60.0, 60.0), (2.0, 5.0), (None, 30.0)])
def test_shutdown_join_budget_is_the_floor_or_the_longer_timeout(timeout, budget):
    """Unset timeout resolves to the client's 30s default, so the join waits as long as a blocked call can."""
    provider = HonchoMemoryProvider()
    provider._config = HonchoClientConfig(api_key="k", timeout=timeout)
    assert provider._shutdown_join_budget() == budget


class TestManagerAfterShutdown:
    def test_prefetch_is_skipped_after_shutdown(self, async_manager):
        calls = []
        async_manager.get_prefetch_context = lambda *a, **k: calls.append(a) or {}
        async_manager.shutdown()
        async_manager.prefetch_context("test-session", "query")
        time.sleep(0.05)
        assert calls == []


def test_close_honcho_clients_closes_every_pool_and_drops_the_slots():
    client = MagicMock()
    slot = SingletonSlot()
    slot.get(lambda: client)
    key = ("test", "close-me")
    with _client_slots_lock:
        _client_slots[key] = slot
    try:
        close_honcho_clients()
    finally:
        with _client_slots_lock:
            _client_slots.pop(key, None)
    client._http.close.assert_called_once_with()
    assert key not in _client_slots