import pickle
import unittest

import fakeredis.aioredis
from pwo import async_test

from kaya.session import Session
from kaya.session.redis import RedisSessionStore


class RedisSessionIndexTest(unittest.TestCase):

    def setUp(self) -> None:
        self.client = fakeredis.aioredis.FakeRedis()
        self.store = RedisSessionStore(self.client)

    @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')}))
        index_key = self.store._index_key('sub', 'alice')
        self.assertEqual({b'abc'}, await self.client.smembers(index_key))
        await self.store.delete('abc')
        self.assertEqual(set(), await self.client.smembers(index_key))

    @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('kaya:session:abc')
        self.assertEqual([], await self.store.find_sessions('sub', 'alice'))
        index_key = self.store._index_key('sub', 'alice')
        self.assertEqual(set(), await self.client.smembers(index_key))

    @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')
        ttl = await self.client.ttl(index_key)
        self.assertGreater(ttl, 60)
        self.assertLessEqual(ttl, 120)

    @async_test
    async def test_legacy_payload_is_upgraded_on_save(self) -> None:
        await self.client.set('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('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'))


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