"""Tests for OAuth 2.1 provider module."""

import time
from unittest.mock import AsyncMock, MagicMock, patch

import pytest
from pydantic import AnyHttpUrl

from mcp.server.auth.provider import AccessToken, AuthorizationCode, AuthorizationParams, RefreshToken, TokenError
from mcp.shared.auth import OAuthClientInformationFull

from crustdata_mcp import oauth
from crustdata_mcp.oauth import (
    CrustdataOAuthProvider,
    _check_rate_limit,
    _cleanup_expired_sessions,
    _get_client_ip,
    _render_login_page,
    _validate_api_key,
)


def _make_client(client_id: str = "test-client") -> OAuthClientInformationFull:
    return OAuthClientInformationFull(
        client_id=client_id,
        client_secret="test-secret",
        redirect_uris=[AnyHttpUrl("https://example.com/callback")],
    )


def _make_params(**overrides) -> AuthorizationParams:
    defaults = dict(
        state="test-state",
        scopes=["read", "write"],
        code_challenge="test-challenge",
        code_challenge_method="S256",
        redirect_uri=AnyHttpUrl("https://example.com/callback"),
        redirect_uri_provided_explicitly=True,
        resource=None,
    )
    defaults.update(overrides)
    return AuthorizationParams(**defaults)


@pytest.fixture(autouse=True)
def _clear_stores():
    """Clear all in-memory stores before each test."""
    oauth._clients.clear()
    oauth._auth_codes.clear()
    oauth._code_to_api_key.clear()
    oauth._revoked_tokens.clear()
    oauth._pending_authorizations.clear()
    oauth._login_attempts.clear()
    yield
    oauth._clients.clear()
    oauth._auth_codes.clear()
    oauth._code_to_api_key.clear()
    oauth._revoked_tokens.clear()
    oauth._pending_authorizations.clear()
    oauth._login_attempts.clear()


class TestClientRegistration:
    @pytest.mark.asyncio
    async def test_register_and_get_client(self):
        provider = CrustdataOAuthProvider()
        client_info = _make_client("my-client")
        await provider.register_client(client_info)

        result = await provider.get_client("my-client")
        assert result is not None
        assert result.client_id == "my-client"

    @pytest.mark.asyncio
    async def test_get_unknown_client_returns_none(self):
        provider = CrustdataOAuthProvider()
        assert await provider.get_client("nonexistent") is None

    @pytest.mark.asyncio
    async def test_register_generates_ids_if_missing(self):
        provider = CrustdataOAuthProvider()
        client_info = OAuthClientInformationFull(
            redirect_uris=[AnyHttpUrl("https://example.com/callback")],
        )
        await provider.register_client(client_info)

        assert client_info.client_id is not None
        assert len(client_info.client_id) > 0
        assert client_info.client_secret is not None
        assert client_info.client_id_issued_at is not None


class TestAuthorize:
    @pytest.mark.asyncio
    async def test_authorize_returns_login_url(self):
        provider = CrustdataOAuthProvider()
        client = _make_client()
        params = _make_params()

        url = await provider.authorize(client, params)

        assert url.startswith("/oauth/login?session_id=")

    @pytest.mark.asyncio
    async def test_authorize_stores_pending_session(self):
        provider = CrustdataOAuthProvider()
        client = _make_client()
        params = _make_params()

        url = await provider.authorize(client, params)
        session_id = url.split("session_id=")[1]

        assert session_id in oauth._pending_authorizations
        stored_params, stored_client, created_at = oauth._pending_authorizations[session_id]
        assert stored_params.state == "test-state"
        assert created_at <= time.time()


class TestAuthorizationCodeFlow:
    @pytest.mark.asyncio
    async def test_load_valid_auth_code(self):
        provider = CrustdataOAuthProvider()
        client = _make_client()

        code = AuthorizationCode(
            code="test-code",
            scopes=["read"],
            expires_at=time.time() + 300,
            client_id="test-client",
            code_challenge="challenge",
            redirect_uri=AnyHttpUrl("https://example.com/callback"),
            redirect_uri_provided_explicitly=True,
        )
        oauth._auth_codes["test-code"] = code

        result = await provider.load_authorization_code(client, "test-code")
        assert result is not None
        assert result.code == "test-code"

    @pytest.mark.asyncio
    async def test_load_expired_auth_code_returns_none(self):
        provider = CrustdataOAuthProvider()
        client = _make_client()

        code = AuthorizationCode(
            code="expired-code",
            scopes=["read"],
            expires_at=time.time() - 10,  # expired
            client_id="test-client",
            code_challenge="challenge",
            redirect_uri=AnyHttpUrl("https://example.com/callback"),
            redirect_uri_provided_explicitly=True,
        )
        oauth._auth_codes["expired-code"] = code

        result = await provider.load_authorization_code(client, "expired-code")
        assert result is None
        assert "expired-code" not in oauth._auth_codes

    @pytest.mark.asyncio
    async def test_load_unknown_auth_code_returns_none(self):
        provider = CrustdataOAuthProvider()
        client = _make_client()

        result = await provider.load_authorization_code(client, "nonexistent")
        assert result is None

    @pytest.mark.asyncio
    async def test_load_auth_code_rejects_wrong_client(self):
        provider = CrustdataOAuthProvider()
        legitimate_client = _make_client("legitimate-client")
        malicious_client = _make_client("malicious-client")

        code = AuthorizationCode(
            code="stolen-code",
            scopes=["read"],
            expires_at=time.time() + 300,
            client_id="legitimate-client",
            code_challenge="challenge",
            redirect_uri=AnyHttpUrl("https://example.com/callback"),
            redirect_uri_provided_explicitly=True,
        )
        oauth._auth_codes["stolen-code"] = code

        # Legitimate client can load it
        result = await provider.load_authorization_code(legitimate_client, "stolen-code")
        assert result is not None

        # Malicious client cannot
        result = await provider.load_authorization_code(malicious_client, "stolen-code")
        assert result is None

    @pytest.mark.asyncio
    async def test_exchange_authorization_code_returns_api_key_as_token(self):
        provider = CrustdataOAuthProvider()
        client = _make_client()

        auth_code = AuthorizationCode(
            code="exchange-code",
            scopes=["read", "write"],
            expires_at=time.time() + 300,
            client_id="test-client",
            code_challenge="challenge",
            redirect_uri=AnyHttpUrl("https://example.com/callback"),
            redirect_uri_provided_explicitly=True,
        )
        oauth._auth_codes["exchange-code"] = auth_code
        oauth._code_to_api_key["exchange-code"] = "my-crustdata-api-key-12345"

        token = await provider.exchange_authorization_code(client, auth_code)

        assert token.access_token == "my-crustdata-api-key-12345"
        assert token.token_type == "Bearer"
        assert "exchange-code" not in oauth._auth_codes
        assert "exchange-code" not in oauth._code_to_api_key

    @pytest.mark.asyncio
    async def test_exchange_unknown_code_raises(self):
        provider = CrustdataOAuthProvider()
        client = _make_client()

        auth_code = AuthorizationCode(
            code="unknown-code",
            scopes=[],
            expires_at=time.time() + 300,
            client_id="test-client",
            code_challenge="challenge",
            redirect_uri=AnyHttpUrl("https://example.com/callback"),
            redirect_uri_provided_explicitly=True,
        )

        with pytest.raises(TokenError):
            await provider.exchange_authorization_code(client, auth_code)


class TestAccessToken:
    @pytest.mark.asyncio
    async def test_load_valid_token(self):
        provider = CrustdataOAuthProvider()
        result = await provider.load_access_token("valid-token-at-least-10-chars")

        assert result is not None
        assert result.token == "valid-token-at-least-10-chars"
        assert result.scopes == ["read", "write"]

    @pytest.mark.asyncio
    async def test_reject_empty_token(self):
        provider = CrustdataOAuthProvider()
        assert await provider.load_access_token("") is None

    @pytest.mark.asyncio
    async def test_reject_short_token(self):
        provider = CrustdataOAuthProvider()
        assert await provider.load_access_token("short") is None
        assert await provider.load_access_token("123456789") is None  # 9 chars

    @pytest.mark.asyncio
    async def test_accept_minimum_length_token(self):
        provider = CrustdataOAuthProvider()
        result = await provider.load_access_token("1234567890")  # 10 chars
        assert result is not None

    @pytest.mark.asyncio
    async def test_reject_revoked_token(self):
        provider = CrustdataOAuthProvider()
        oauth._revoked_tokens["revoked-token-12345"] = time.time()

        result = await provider.load_access_token("revoked-token-12345")
        assert result is None


class TestTokenRevocation:
    @pytest.mark.asyncio
    async def test_revoke_access_token_object(self):
        provider = CrustdataOAuthProvider()
        token = AccessToken(
            token="token-to-revoke-123",
            client_id="test",
            scopes=["read"],
            expires_at=None,
        )

        await provider.revoke_token(token)

        assert "token-to-revoke-123" in oauth._revoked_tokens
        assert await provider.load_access_token("token-to-revoke-123") is None

    @pytest.mark.asyncio
    async def test_revoke_string_token(self):
        provider = CrustdataOAuthProvider()

        await provider.revoke_token("string-token-to-revoke")

        assert "string-token-to-revoke" in oauth._revoked_tokens

    @pytest.mark.asyncio
    async def test_revoke_none_is_noop(self):
        provider = CrustdataOAuthProvider()
        await provider.revoke_token(None)
        assert len(oauth._revoked_tokens) == 0


class TestRefreshToken:
    @pytest.mark.asyncio
    async def test_load_refresh_token_returns_none(self):
        provider = CrustdataOAuthProvider()
        client = _make_client()
        assert await provider.load_refresh_token(client, "any-token") is None

    @pytest.mark.asyncio
    async def test_exchange_refresh_token_raises(self):
        provider = CrustdataOAuthProvider()
        client = _make_client()

        fake_token = RefreshToken(token="fake-refresh", client_id="test-client", scopes=["read"])
        with pytest.raises(TokenError) as exc_info:
            await provider.exchange_refresh_token(client, fake_token, ["read"])
        assert exc_info.value.error == "unsupported_grant_type"


class TestCleanupExpiredSessions:
    def test_cleans_expired_pending_authorizations(self):
        params = _make_params()
        client = _make_client()

        # Add an expired session (created 15 minutes ago, TTL is 10 min)
        oauth._pending_authorizations["expired"] = (params, client, time.time() - 900)
        # Add a valid session (created just now)
        oauth._pending_authorizations["valid"] = (params, client, time.time())

        _cleanup_expired_sessions()

        assert "expired" not in oauth._pending_authorizations
        assert "valid" in oauth._pending_authorizations

    def test_cleans_expired_auth_codes(self):
        expired_code = AuthorizationCode(
            code="expired",
            scopes=[],
            expires_at=time.time() - 10,
            client_id="test",
            code_challenge="c",
            redirect_uri=AnyHttpUrl("https://example.com/callback"),
            redirect_uri_provided_explicitly=True,
        )
        valid_code = AuthorizationCode(
            code="valid",
            scopes=[],
            expires_at=time.time() + 300,
            client_id="test",
            code_challenge="c",
            redirect_uri=AnyHttpUrl("https://example.com/callback"),
            redirect_uri_provided_explicitly=True,
        )
        oauth._auth_codes["expired"] = expired_code
        oauth._auth_codes["valid"] = valid_code
        oauth._code_to_api_key["expired"] = "key1"
        oauth._code_to_api_key["valid"] = "key2"

        _cleanup_expired_sessions()

        assert "expired" not in oauth._auth_codes
        assert "expired" not in oauth._code_to_api_key
        assert "valid" in oauth._auth_codes
        assert "valid" in oauth._code_to_api_key

    def test_cleans_expired_revocations(self):
        # Add old revocation (31 days ago)
        oauth._revoked_tokens["old-token"] = time.time() - (86400 * 31)
        # Add recent revocation
        oauth._revoked_tokens["new-token"] = time.time()

        _cleanup_expired_sessions()

        assert "old-token" not in oauth._revoked_tokens
        assert "new-token" in oauth._revoked_tokens


class TestValidateApiKey:
    @pytest.mark.asyncio
    async def test_valid_key(self, respx_mock):
        respx_mock.get("https://api.crustdata.com/user/credits").mock(
            return_value=__import__("httpx").Response(200, json={"credits": 100})
        )
        assert await _validate_api_key("valid-key") is True

    @pytest.mark.asyncio
    async def test_invalid_key(self, respx_mock):
        respx_mock.get("https://api.crustdata.com/user/credits").mock(
            return_value=__import__("httpx").Response(401, json={"error": "unauthorized"})
        )
        assert await _validate_api_key("bad-key") is False

    @pytest.mark.asyncio
    async def test_network_error(self, respx_mock):
        respx_mock.get("https://api.crustdata.com/user/credits").mock(
            side_effect=__import__("httpx").ConnectError("connection refused")
        )
        assert await _validate_api_key("any-key") is None


class TestRenderLoginPage:
    def test_contains_session_id(self):
        html = _render_login_page("my-session-123")
        assert 'value="my-session-123"' in html

    def test_error_hidden_by_default(self):
        html = _render_login_page("sid")
        assert "display: none" in html

    def test_error_shown_when_provided(self):
        html = _render_login_page("sid", error="Bad key")
        assert "display: block" in html
        assert "Bad key" in html

    def test_password_input_type(self):
        html = _render_login_page("sid")
        assert 'type="password"' in html

    def test_show_hide_toggle(self):
        html = _render_login_page("sid")
        assert "Show" in html


class TestLoginRateLimiting:
    def test_under_limit_allowed(self):
        now = time.time()
        # 4 attempts in the last minute (limit is 5)
        oauth._login_attempts["sid"] = [now - 30, now - 20, now - 10, now - 5]
        recent = [t for t in oauth._login_attempts["sid"] if now - t < oauth.LOGIN_RATE_LIMIT_WINDOW]
        assert len(recent) < oauth.LOGIN_RATE_LIMIT_MAX

    def test_at_limit_blocked(self):
        now = time.time()
        # 5 attempts in the last minute
        oauth._login_attempts["sid"] = [now - 40, now - 30, now - 20, now - 10, now - 5]
        recent = [t for t in oauth._login_attempts["sid"] if now - t < oauth.LOGIN_RATE_LIMIT_WINDOW]
        assert len(recent) >= oauth.LOGIN_RATE_LIMIT_MAX

    def test_old_attempts_dont_count(self):
        now = time.time()
        # 5 attempts but all older than the window
        oauth._login_attempts["sid"] = [now - 120, now - 110, now - 100, now - 90, now - 80]
        recent = [t for t in oauth._login_attempts["sid"] if now - t < oauth.LOGIN_RATE_LIMIT_WINDOW]
        assert len(recent) < oauth.LOGIN_RATE_LIMIT_MAX

    def test_cleanup_removes_stale_attempts(self):
        now = time.time()
        oauth._login_attempts["stale"] = [now - 120]  # older than window
        oauth._login_attempts["active"] = [now - 5]  # recent

        _cleanup_expired_sessions()

        assert "stale" not in oauth._login_attempts
        assert "active" in oauth._login_attempts


def _mock_request(query_params: dict | None = None, cookies: dict | None = None, headers: dict | None = None, client_host: str = "127.0.0.1"):
    """Create a mock Starlette Request with the given query params, cookies, and headers."""
    req = MagicMock()
    req.query_params = query_params or {}
    req.cookies = cookies or {}
    req.headers = headers or {}
    req.client = MagicMock()
    req.client.host = client_host
    return req


class TestAuth0LoginUrl:
    def test_builds_url_with_base_url(self):
        provider = CrustdataOAuthProvider(base_url="https://mcp.crustdata.com")
        url = provider._auth0_login_url("test-session-123")

        assert url.startswith("https://api.crustdata.com/auth0/login?redirect_url=")
        assert "session_id%3Dtest-session-123" in url or "session_id=test-session-123" in url

    def test_callback_url_is_properly_encoded(self):
        provider = CrustdataOAuthProvider(base_url="https://mcp.crustdata.com")
        url = provider._auth0_login_url("abc")

        # The redirect_url value should be URL-encoded (no bare ? or & from the inner URL)
        redirect_param = url.split("redirect_url=")[1]
        assert "?" not in redirect_param  # inner ? should be encoded as %3F


class TestRenderLoginPageAuth0:
    def test_no_auth0_button_without_url(self):
        html = _render_login_page("sid")
        assert "Sign in with Crustdata account" not in html

    def test_auth0_button_shown_with_url(self):
        html = _render_login_page("sid", auth0_url="https://example.com/auth0/login?redirect_url=callback")
        assert "Sign in with Crustdata account" in html
        assert "https://example.com/auth0/login" in html

    def test_auth0_button_with_error(self):
        html = _render_login_page("sid", error="Bad key", auth0_url="https://example.com/auth0")
        assert "Sign in with Crustdata account" in html
        assert "Bad key" in html

    def test_divider_shown_with_auth0(self):
        html = _render_login_page("sid", auth0_url="https://example.com/auth0")
        assert ">or<" in html.replace(" ", "")


class TestCompleteAuthorization:
    def test_generates_auth_code_and_redirects(self):
        provider = CrustdataOAuthProvider()
        params = _make_params()
        client = _make_client()
        oauth._pending_authorizations["sid-123"] = (params, client, time.time())

        response = provider._complete_authorization("sid-123", "my-api-key-12345678")

        assert response.status_code == 302
        assert "sid-123" not in oauth._pending_authorizations
        # Auth code should be stored
        assert len(oauth._auth_codes) == 1
        code = list(oauth._auth_codes.keys())[0]
        assert oauth._code_to_api_key[code] == "my-api-key-12345678"


class TestAuth0Callback:
    @pytest.mark.asyncio
    async def test_invalid_session_returns_400(self):
        provider = CrustdataOAuthProvider(base_url="https://mcp.crustdata.com")
        request = _mock_request(query_params={"session_id": "nonexistent"})

        response = await provider.handle_auth0_callback(request)

        assert response.status_code == 400

    @pytest.mark.asyncio
    async def test_missing_token_cookie_shows_error(self):
        provider = CrustdataOAuthProvider(base_url="https://mcp.crustdata.com")
        params = _make_params()
        client = _make_client()
        oauth._pending_authorizations["sid-ok"] = (params, client, time.time())

        request = _mock_request(
            query_params={"session_id": "sid-ok"},
            cookies={},  # no token cookie
        )
        response = await provider.handle_auth0_callback(request)

        assert response.status_code == 200
        assert "no token received" in response.body.decode().lower()
        # Session should still be valid for retry
        assert "sid-ok" in oauth._pending_authorizations

    @pytest.mark.asyncio
    async def test_valid_token_cookie_completes_flow(self, respx_mock):
        respx_mock.get("https://api.crustdata.com/user/credits").mock(
            return_value=__import__("httpx").Response(200, json={"credits": 50})
        )

        provider = CrustdataOAuthProvider(base_url="https://mcp.crustdata.com")
        params = _make_params()
        client = _make_client()
        oauth._pending_authorizations["sid-auth0"] = (params, client, time.time())

        request = _mock_request(
            query_params={"session_id": "sid-auth0"},
            cookies={"token": "valid-api-key-from-auth0"},
        )
        response = await provider.handle_auth0_callback(request)

        assert response.status_code == 302
        assert "sid-auth0" not in oauth._pending_authorizations
        # Verify auth code was created with the right API key
        code = list(oauth._code_to_api_key.keys())[0]
        assert oauth._code_to_api_key[code] == "valid-api-key-from-auth0"

    @pytest.mark.asyncio
    async def test_invalid_token_cookie_shows_error(self, respx_mock):
        respx_mock.get("https://api.crustdata.com/user/credits").mock(
            return_value=__import__("httpx").Response(401, json={"error": "unauthorized"})
        )

        provider = CrustdataOAuthProvider(base_url="https://mcp.crustdata.com")
        params = _make_params()
        client = _make_client()
        oauth._pending_authorizations["sid-bad"] = (params, client, time.time())

        request = _mock_request(
            query_params={"session_id": "sid-bad"},
            cookies={"token": "expired-or-invalid-token"},
        )
        response = await provider.handle_auth0_callback(request)

        assert response.status_code == 200
        assert "invalid" in response.body.decode().lower()
        # Session should still be valid for retry
        assert "sid-bad" in oauth._pending_authorizations

    @pytest.mark.asyncio
    async def test_network_error_shows_unavailable(self, respx_mock):
        respx_mock.get("https://api.crustdata.com/user/credits").mock(
            side_effect=__import__("httpx").ConnectError("connection refused")
        )

        provider = CrustdataOAuthProvider(base_url="https://mcp.crustdata.com")
        params = _make_params()
        client = _make_client()
        oauth._pending_authorizations["sid-net"] = (params, client, time.time())

        request = _mock_request(
            query_params={"session_id": "sid-net"},
            cookies={"token": "some-token-value-here"},
        )
        response = await provider.handle_auth0_callback(request)

        assert response.status_code == 200
        assert "unavailable" in response.body.decode().lower()

    @pytest.mark.asyncio
    async def test_rate_limited(self):
        provider = CrustdataOAuthProvider(base_url="https://mcp.crustdata.com")
        params = _make_params()
        client = _make_client()
        oauth._pending_authorizations["sid-rl"] = (params, client, time.time())

        # Fill up the rate limit bucket for this IP
        now = time.time()
        oauth._login_attempts["10.0.0.1"] = [now - 4, now - 3, now - 2, now - 1, now]

        request = _mock_request(
            query_params={"session_id": "sid-rl"},
            cookies={"token": "some-token"},
            headers={"x-forwarded-for": "10.0.0.1"},
        )
        response = await provider.handle_auth0_callback(request)

        assert response.status_code == 429
        assert "too many attempts" in response.body.decode().lower()


class TestGetClientIp:
    def test_uses_rightmost_xff(self):
        req = _mock_request(headers={"x-forwarded-for": "spoofed.ip, 1.2.3.4"})
        assert _get_client_ip(req) == "1.2.3.4"

    def test_single_xff_entry(self):
        req = _mock_request(headers={"x-forwarded-for": "5.6.7.8"})
        assert _get_client_ip(req) == "5.6.7.8"

    def test_falls_back_to_client_host(self):
        req = _mock_request(client_host="192.168.1.1")
        assert _get_client_ip(req) == "192.168.1.1"

    def test_strips_whitespace(self):
        req = _mock_request(headers={"x-forwarded-for": "1.1.1.1,  2.2.2.2 "})
        assert _get_client_ip(req) == "2.2.2.2"

    def test_ignores_spoofed_leftmost(self):
        # Attacker sends X-Forwarded-For: fake, ALB appends real IP
        req = _mock_request(headers={"x-forwarded-for": "attacker-controlled, real-client-ip"})
        assert _get_client_ip(req) == "real-client-ip"


class TestCheckRateLimit:
    def test_allows_under_limit(self):
        req = _mock_request(headers={"x-forwarded-for": "10.10.10.10"})
        assert _check_rate_limit(req) is False

    def test_blocks_at_limit(self):
        now = time.time()
        oauth._login_attempts["10.10.10.10"] = [now - 4, now - 3, now - 2, now - 1, now]
        req = _mock_request(headers={"x-forwarded-for": "10.10.10.10"})
        assert _check_rate_limit(req) is True

    def test_different_ips_independent(self):
        now = time.time()
        oauth._login_attempts["10.10.10.10"] = [now - 4, now - 3, now - 2, now - 1, now]
        req = _mock_request(headers={"x-forwarded-for": "10.10.10.99"})
        assert _check_rate_limit(req) is False


class TestClientCleanup:
    def test_cleans_stale_clients(self):
        old_client = _make_client("old-client")
        old_client.client_id_issued_at = int(time.time()) - 90000  # 25 hours ago
        oauth._clients["old-client"] = old_client

        new_client = _make_client("new-client")
        new_client.client_id_issued_at = int(time.time())
        oauth._clients["new-client"] = new_client

        _cleanup_expired_sessions()

        assert "old-client" not in oauth._clients
        assert "new-client" in oauth._clients


class TestHtmlEscaping:
    def test_session_id_escaped(self):
        html = _render_login_page('"><script>alert(1)</script>')
        assert "alert(1)</script>" not in html
        assert "&lt;script&gt;alert(1)&lt;/script&gt;" in html

    def test_error_escaped(self):
        html = _render_login_page("sid", error='<img onerror="alert(1)">')
        assert 'onerror="alert(1)"' not in html
        assert "&lt;img" in html

    def test_auth0_url_escaped(self):
        html = _render_login_page("sid", auth0_url='https://evil.com" onclick="alert(1)')
        assert 'onclick="alert(1)' not in html
