import unittest
from time import time
from typing import Any, Optional
from urllib.parse import parse_qs, urlparse

import httpx
import jwt
from cryptography.hazmat.primitives.asymmetric import rsa
from pwo import async_test
from test_oidc import MockOIDCProvider, MockTransport

from kaya.core import HttpContext, KayaApp
from kaya.oidc import OIDCConfig, OIDCMixin, OIDCProviders
from kaya.session import InMemorySessionStore, SessionMixin

_EVENT = 'http://schemas.openid.net/event/backchannel-logout'


class BackchannelProvider(MockOIDCProvider):
    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        self.sid: Optional[str] = None

    def issue_id_token(self, *args: Any, **kwargs: Any) -> str:
        token = super().issue_id_token(*args, **kwargs)
        if self.sid is None:
            return token
        payload = jwt.decode(token, options={'verify_signature': False})
        payload['sid'] = self.sid
        return jwt.encode(
            payload,
            self.private_key,
            algorithm='RS256',
            headers={'kid': self.kid},
        )


class BackchannelLogoutTest(unittest.TestCase):

    def setUp(self) -> None:
        self.provider = BackchannelProvider()
        self.http_client = httpx.AsyncClient(transport=MockTransport(self.provider))
        self.store = InMemorySessionStore()
        self.session = SessionMixin(self.store)
        self.oidc = OIDCMixin(
            OIDCConfig(
                issuer=self.provider.issuer,
                client_id='client',
                client_secret='secret',
                redirect_uri='http://localhost:8000/auth/callback',
                backchannel_logout_path='/auth/backchannel-logout',
                http_client=self.http_client,
            ),
            session=self.session,
        )
        self.app = KayaApp(mixins=[self.session, self.oidc])
        self.transport = httpx.ASGITransport(app=self.app)

        @self.app.GET('/profile')
        @self.oidc.require_auth
        async def profile(ctx: HttpContext) -> None:
            user = self.oidc.get_user(ctx)
            assert user is not None
            await ctx.send_str(200, user.sub)

        @self.app.GET('/refresh')
        @self.oidc.require_auth
        async def refresh(ctx: HttpContext) -> None:
            token = await self.oidc.refresh_access_token(ctx.session)
            await ctx.send_str(200, token or 'no-token')

    def _client(self) -> httpx.AsyncClient:
        return httpx.AsyncClient(transport=self.transport, base_url='http://127.0.0.1:80')

    async def _login(self, client: httpx.AsyncClient, sid: str) -> None:
        self.provider.sid = sid
        response = await client.get('/auth/login', follow_redirects=False)
        self.assertEqual(302, response.status_code)
        query = parse_qs(urlparse(response.headers['Location']).query)
        self.provider.nonce = query['nonce'][0]
        response = await client.get(
            '/auth/callback',
            params={'code': 'mock-code', 'state': query['state'][0]},
            follow_redirects=False,
        )
        self.assertEqual(302, response.status_code)

    def _claims(self, **overrides: Any) -> dict[str, Any]:
        claims: dict[str, Any] = {
            'iss': self.provider.issuer,
            'aud': 'client',
            'iat': time(),
            'jti': 'logout-token-id',
            'events': {_EVENT: {}},
            'sub': self.provider.sub,
        }
        claims.update(overrides)
        return claims

    def _token(
        self,
        claims: dict[str, Any],
        *,
        typ: str = 'logout+jwt',
        key: Any = None,
        algorithm: str = 'RS256',
    ) -> str:
        return jwt.encode(
            claims,
            self.provider.private_key if key is None else key,
            algorithm=algorithm,
            headers={'kid': self.provider.kid, 'typ': typ},
        )

    async def _post_logout(self, token: str) -> httpx.Response:
        async with self._client() as client:
            return await client.post(
                '/auth/backchannel-logout',
                data={'logout_token': token},
            )

    async def _profile_status(self, client: httpx.AsyncClient) -> int:
        response = await client.get('/profile', follow_redirects=False)
        return response.status_code

    @async_test
    async def test_logout_by_sid_only_invalidates_matching_session(self) -> None:
        async with self._client() as first, self._client() as second:
            await self._login(first, 'sid-a')
            await self._login(second, 'sid-b')

            response = await self._post_logout(self._token(self._claims(sid='sid-a')))
            self.assertEqual(200, response.status_code)
            self.assertEqual('no-store', response.headers['Cache-Control'])
            self.assertNotIn('Set-Cookie', response.headers)
            self.assertEqual(302, await self._profile_status(first))
            self.assertEqual(200, await self._profile_status(second))
            self.assertEqual([], await self.store.find_sessions('oidc:default:sid', 'sid-a'))
            self.assertEqual(
                [second.cookies['session_id']],
                await self.store.find_sessions('oidc:default:sid', 'sid-b'),
            )

    @async_test
    async def test_logout_by_sub_invalidates_all_user_sessions(self) -> None:
        async with self._client() as first, self._client() as second, self._client() as other:
            await self._login(first, 'sid-a')
            await self._login(second, 'sid-b')
            self.provider.sub = 'bob'
            await self._login(other, 'sid-c')

            response = await self._post_logout(self._token(self._claims(sub='user123')))
            self.assertEqual(200, response.status_code)
            self.assertEqual(302, await self._profile_status(first))
            self.assertEqual(302, await self._profile_status(second))
            self.assertEqual(200, await self._profile_status(other))

    @async_test
    async def test_sid_logout_verifies_sub_when_both_are_present(self) -> None:
        async with self._client() as client:
            await self._login(client, 'sid-a')
            response = await self._post_logout(
                self._token(self._claims(sub='someone-else', sid='sid-a'))
            )
            self.assertEqual(200, response.status_code)
            self.assertEqual(200, await self._profile_status(client))

    @async_test
    async def test_unknown_session_is_idempotent(self) -> None:
        response = await self._post_logout(self._token(self._claims(sid='unknown')))
        self.assertEqual(200, response.status_code)

    @async_test
    async def test_invalid_logout_tokens_are_rejected(self) -> None:
        other_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
        missing_iat = self._claims(sid='sid-a')
        del missing_iat['iat']
        missing_jti = self._claims(sid='sid-a')
        del missing_jti['jti']
        missing_events = self._claims(sid='sid-a')
        del missing_events['events']
        missing_subject = self._claims()
        del missing_subject['sub']

        invalid_tokens = [
            self._token(self._claims(sid='sid-a'), typ='JWT'),
            jwt.encode(
                self._claims(sid='sid-a'),
                None,
                algorithm='none',
                headers={'typ': 'logout+jwt'},
            ),
            self._token(self._claims(sid='sid-a'), key=other_key),
            self._token(self._claims(iss='https://evil.local', sid='sid-a')),
            self._token(self._claims(aud='other-client', sid='sid-a')),
            self._token(missing_iat),
            self._token(missing_jti),
            self._token(missing_events),
            self._token(self._claims(events={'wrong-event': {}}, sid='sid-a')),
            self._token(self._claims(events={_EVENT: None}, sid='sid-a')),
            self._token(self._claims(nonce='forbidden', sid='sid-a')),
            self._token(missing_subject),
            self._token(self._claims(exp=time() - 3600, sid='sid-a')),
            self._token(self._claims(sid='sid-a'), key='x' * 32, algorithm='HS256'),
        ]

        async with self._client() as client:
            await self._login(client, 'sid-a')
            for token in invalid_tokens:
                response = await self._post_logout(token)
                self.assertEqual(400, response.status_code)
            self.assertEqual(200, await self._profile_status(client))

    @async_test
    async def test_invalid_requests_are_rejected(self) -> None:
        async with self._client() as client:
            response = await client.post('/auth/backchannel-logout', data={})
            self.assertEqual(400, response.status_code)
            response = await client.post(
                '/auth/backchannel-logout',
                content='logout_token=abc',
                headers={'Content-Type': 'text/plain'},
            )
            self.assertEqual(400, response.status_code)
            response = await client.post(
                '/auth/backchannel-logout',
                content='x' * (64 * 1024 + 1),
                headers={'Content-Type': 'application/x-www-form-urlencoded'},
            )
            self.assertEqual(400, response.status_code)

    @async_test
    async def test_local_logout_removes_indexes(self) -> None:
        async with self._client() as client:
            await self._login(client, 'sid-a')
            self.assertEqual(
                [client.cookies['session_id']],
                await self.store.find_sessions('oidc:default:sub', 'user123'),
            )
            await client.get('/auth/logout', follow_redirects=False)
            self.assertEqual([], await self.store.find_sessions('oidc:default:sub', 'user123'))
            self.assertEqual([], await self.store.find_sessions('oidc:default:sid', 'sid-a'))

    @async_test
    async def test_relogin_replaces_old_sid_index(self) -> None:
        async with self._client() as client:
            await self._login(client, 'sid-a')
            await self._login(client, 'sid-b')
            self.assertEqual([], await self.store.find_sessions('oidc:default:sid', 'sid-a'))
            self.assertEqual(
                [client.cookies['session_id']],
                await self.store.find_sessions('oidc:default:sid', 'sid-b'),
            )

    @async_test
    async def test_refresh_updates_changed_sid_index(self) -> None:
        async with self._client() as client:
            await self._login(client, 'sid-a')
            self.provider.sid = 'sid-b'
            response = await client.get('/refresh')
            self.assertEqual(200, response.status_code)
            self.assertNotEqual('no-token', response.text)
            self.assertEqual([], await self.store.find_sessions('oidc:default:sid', 'sid-a'))
            self.assertEqual(
                [client.cookies['session_id']],
                await self.store.find_sessions('oidc:default:sid', 'sid-b'),
            )

    @async_test
    async def test_existing_session_is_backfilled(self) -> None:
        async with self._client() as client:
            await self._login(client, 'sid-a')
            session_id = client.cookies['session_id']
            self.store._data[session_id].pop('oidc:default:indexed')
            self.store._session_indexes.pop(session_id, None)
            for session_ids in self.store._index_sessions.values():
                session_ids.discard(session_id)
            self.assertEqual([], await self.store.find_sessions('oidc:default:sid', 'sid-a'))

            self.assertEqual(200, await self._profile_status(client))
            self.assertEqual(
                [session_id],
                await self.store.find_sessions('oidc:default:sid', 'sid-a'),
            )

    @async_test
    async def test_index_capable_store_is_required(self) -> None:
        class LegacyStore(InMemorySessionStore):
            supports_indexes = False

        with self.assertRaises(ValueError):
            OIDCMixin(
                OIDCConfig(
                    issuer=self.provider.issuer,
                    client_id='client',
                    redirect_uri='http://localhost:8000/auth/callback',
                    backchannel_logout_path='/auth/backchannel-logout',
                ),
                session=SessionMixin(LegacyStore()),
            )


class MultiProviderBackchannelLogoutTest(unittest.TestCase):

    def setUp(self) -> None:
        self.google = BackchannelProvider('google', sub='shared-sub', email='g@example.com')
        self.github = BackchannelProvider('github', sub='shared-sub', email='h@example.com')
        self.store = InMemorySessionStore()
        self.session = SessionMixin(self.store)
        self.oidc = OIDCProviders(
            [
                self._config(self.google, 'google'),
                self._config(self.github, 'github'),
            ],
            session=self.session,
            strategy='multi-identity',
        )
        self.app = KayaApp(mixins=[self.session, self.oidc])
        self.transport = httpx.ASGITransport(app=self.app)

        @self.app.GET('/google-profile')
        @self.oidc.require_auth('google')
        async def google_profile(ctx: HttpContext) -> None:
            await ctx.send_str(200, 'google')

        @self.app.GET('/github-profile')
        @self.oidc.require_auth('github')
        async def github_profile(ctx: HttpContext) -> None:
            await ctx.send_str(200, 'github')

    @staticmethod
    def _config(provider: BackchannelProvider, name: str) -> OIDCConfig:
        return OIDCConfig(
            issuer=provider.issuer,
            client_id='client',
            client_secret='secret',
            redirect_uri=f'http://localhost:8000/auth/{name}/callback',
            provider=name,
            backchannel_logout_path=f'/auth/{name}/backchannel-logout',
            http_client=httpx.AsyncClient(transport=MockTransport(provider)),
        )

    async def _login(
        self,
        client: httpx.AsyncClient,
        provider: BackchannelProvider,
        name: str,
        sid: str,
    ) -> None:
        provider.sid = sid
        response = await client.get(f'/auth/{name}/login', follow_redirects=False)
        query = parse_qs(urlparse(response.headers['Location']).query)
        provider.nonce = query['nonce'][0]
        await client.get(
            f'/auth/{name}/callback',
            params={'code': 'mock-code', 'state': query['state'][0]},
            follow_redirects=False,
        )

    @async_test
    async def test_single_identity_eviction_removes_old_provider_indexes(self) -> None:
        session = SessionMixin(self.store)
        oidc = OIDCProviders(
            [
                self._config(self.google, 'google'),
                self._config(self.github, 'github'),
            ],
            session=session,
            strategy='single-identity',
        )
        app = KayaApp(mixins=[session, oidc])
        transport = httpx.ASGITransport(app=app)
        async with httpx.AsyncClient(
            transport=transport,
            base_url='http://127.0.0.1:80',
        ) as client:
            await self._login(client, self.google, 'google', 'google-sid')
            self.assertEqual(
                [client.cookies['session_id']],
                await self.store.find_sessions('oidc:google:sid', 'google-sid'),
            )
            await self._login(client, self.github, 'github', 'github-sid')
            self.assertEqual(
                [],
                await self.store.find_sessions('oidc:google:sid', 'google-sid'),
            )
            self.assertEqual(
                [client.cookies['session_id']],
                await self.store.find_sessions('oidc:github:sid', 'github-sid'),
            )

    @async_test
    async def test_provider_logout_is_isolated(self) -> None:
        async with httpx.AsyncClient(
            transport=self.transport,
            base_url='http://127.0.0.1:80',
        ) as client:
            await self._login(client, self.google, 'google', 'google-sid')
            await self._login(client, self.github, 'github', 'github-sid')
            token = jwt.encode(
                {
                    'iss': self.google.issuer,
                    'aud': 'client',
                    'iat': time(),
                    'jti': 'google-logout',
                    'events': {_EVENT: {}},
                    'sub': 'shared-sub',
                    'sid': 'google-sid',
                },
                self.google.private_key,
                algorithm='RS256',
                headers={'kid': self.google.kid, 'typ': 'logout+jwt'},
            )
            response = await client.post(
                '/auth/google/backchannel-logout',
                data={'logout_token': token},
            )
            self.assertEqual(200, response.status_code)

            response = await client.get('/google-profile', follow_redirects=False)
            self.assertEqual(302, response.status_code)
            response = await client.get('/github-profile', follow_redirects=False)
            self.assertEqual(200, response.status_code)


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