import pickle
import unittest

from pwo import async_test
from test_store import FakeClock, StubMemcacheClient

from kaya.session import Session
from kaya.session.memcache import MemcacheSessionStore


class MemcacheSessionIndexTest(unittest.TestCase):

    def setUp(self) -> None:
        self.clock = FakeClock()
        self.client = StubMemcacheClient(self.clock)
        self.store = MemcacheSessionStore(self.client, clock=self.clock)  # type: ignore[arg-type]

    @async_test
    async def test_save_load_and_find(self) -> None:
        self.assertTrue(self.store.supports_indexes)
        await self.store.save('abc', Session('abc', {'user': 'alice'}, indexes={('sub', 'alice')}))
        loaded = await self.store.load('abc')
        assert loaded is not None
        self.assertEqual({'user': 'alice'}, dict(loaded))
        self.assertEqual(frozenset({('sub', 'alice')}), loaded.indexes)
        self.assertEqual(['abc'], await self.store.find_sessions('sub', 'alice'))

    @async_test
    async def test_rotation_preserves_index(self) -> None:
        await self.store.save('old', Session('old', indexes={('sub', 'alice')}))
        loaded = await self.store.load('old')
        assert loaded is not None
        loaded.set_id('new')
        await self.store.delete('old')
        await self.store.save('new', loaded)
        self.assertEqual(['new'], await self.store.find_sessions('sub', 'alice'))

    @async_test
    async def test_delete_cleans_reverse_index(self) -> None:
        await self.store.save('abc', Session('abc', indexes={('sub', 'alice')}))
        await self.store.delete('abc')
        self.assertEqual([], await self.store.find_sessions('sub', 'alice'))

    @async_test
    async def test_delete_by_index_deletes_all_sessions(self) -> None:
        await self.store.save('a', Session('a', indexes={('sub', 'alice')}))
        await self.store.save('b', Session('b', indexes={('sub', 'alice')}))
        await self.store.save('c', Session('c', indexes={('sub', 'bob')}))
        self.assertEqual(2, await self.store.delete_by_index('sub', 'alice'))
        self.assertIsNone(await self.store.load('a'))
        self.assertIsNone(await self.store.load('b'))
        self.assertIsNotNone(await self.store.load('c'))

    @async_test
    async def test_stale_member_is_pruned(self) -> None:
        await self.store.save('abc', Session('abc', indexes={('sub', 'alice')}))
        await self.client.delete(b'kaya:session:abc')
        self.assertEqual([], await self.store.find_sessions('sub', 'alice'))
        index_key = self.store._index_key('sub', 'alice')
        payload = await self.client.get(index_key)
        assert payload is not None
        members, _ = pickle.loads(payload)
        self.assertEqual(set(), members)

    @async_test
    async def test_index_ttl_tracks_session_max_age(self) -> None:
        await self.store.save(
            'abc',
            Session('abc', indexes={('sub', 'alice')}),
            max_age=60,
        )
        index_key = self.store._index_key('sub', 'alice')
        self.assertIn(index_key, self.client._expires)
        self.clock.advance(61)
        await self.store.save(
            'abc',
            Session('abc', indexes={('sub', 'alice')}),
            max_age=60,
        )
        self.assertIsNotNone(await self.client.get(index_key))

    @async_test
    async def test_legacy_payload_is_upgraded_on_save(self) -> None:
        await self.client.set(b'kaya:session:abc', pickle.dumps({'foo': 'bar'}))
        loaded = await self.store.load('abc')
        assert loaded is not None
        self.assertEqual({'foo': 'bar'}, dict(loaded))
        self.assertEqual(frozenset(), loaded.indexes)

        loaded.add_index('sub', 'alice')
        await self.store.save('abc', loaded)
        raw = await self.client.get(b'kaya:session:abc')
        assert raw is not None
        payload = pickle.loads(raw)
        self.assertEqual(1, payload['__kaya_session_v1__'])
        self.assertEqual({'foo': 'bar'}, payload['data'])
        self.assertEqual(['abc'], await self.store.find_sessions('sub', 'alice'))

    @async_test
    async def test_cas_conflict_is_retried(self) -> None:
        await self.store.save('existing', Session('existing', indexes={('sub', 'alice')}))
        original_cas = self.client.cas
        attempts = 0

        async def fail_once(key: bytes, value: bytes, cas_token: int, exptime: int = 0) -> bool:
            nonlocal attempts
            attempts += 1
            if attempts == 1:
                return False
            return await original_cas(key, value, cas_token, exptime)

        self.client.cas = fail_once  # type: ignore[method-assign]
        await self.store.save('abc', Session('abc', indexes={('sub', 'alice')}))
        self.assertEqual(
            {'existing', 'abc'},
            set(await self.store.find_sessions('sub', 'alice')),
        )
        self.assertGreaterEqual(attempts, 2)


if __name__ == '__main__':
    unittest.main()