# -*- coding: utf-8 -*-
# Copyright (c) 2025 ContextForge Contributors.
# SPDX-License-Identifier: Apache-2.0

"""Location: ./tests/playwright/security/test_token_lifecycle_enforcement.py
Copyright 2026
SPDX-License-Identifier: Apache-2.0
Authors: Mihai Criveti

Lifecycle enforcement tests for expiration, revocation, JTI, and session auth boundaries.
"""

# Future
from __future__ import annotations

# Standard
from contextlib import suppress
from datetime import datetime, timedelta, timezone
import time
from typing import Any
import uuid

# Third-Party
import jwt
from playwright.sync_api import APIRequestContext, Playwright
import pytest

# First-Party
from mcpgateway.config import settings
from mcpgateway.utils.create_jwt_token import _create_jwt_token
from mcpgateway.utils.jwt_config_helper import get_jwt_private_key_or_secret

# Local
from .conftest import BASE_URL, TEST_PASSWORD


def _extract_token_id(response_json: dict[str, Any]) -> str | None:
    token_obj = response_json.get("token", response_json)
    return token_obj.get("id") or token_obj.get("token_id")


def _api_context(playwright: Playwright, token: str) -> APIRequestContext:
    return playwright.request.new_context(
        base_url=BASE_URL,
        extra_http_headers={"Authorization": f"Bearer {token}", "Accept": "application/json"},
    )


class TestTokenLifecycleEnforcement:
    """High-priority token lifecycle checks from manual security plans."""

    def test_expired_jwt_token_is_rejected(self, playwright: Playwright):
        expired_at = int((datetime.now(timezone.utc) - timedelta(minutes=5)).timestamp())
        expired_token = _create_jwt_token(
            {"sub": "admin@example.com", "exp": expired_at},
            user_data={"email": "admin@example.com", "is_admin": True, "auth_provider": "local"},
        )

        ctx = _api_context(playwright, expired_token)
        try:
            response = ctx.get("/servers")
            status_code = response.status
            response_text = response.text()
        finally:
            ctx.dispose()

        assert status_code == 401, f"Expired token should be rejected, got {status_code}: {response_text}"
        assert "expired" in response_text.lower() or "invalid authentication credentials" in response_text.lower()

    def test_revoked_api_token_cannot_be_used(self, admin_api: APIRequestContext, playwright: Playwright):
        create_resp = admin_api.post(
            "/tokens",
            data={
                "name": f"revocation-enforcement-{uuid.uuid4().hex[:8]}",
                "expires_in_days": 1,
            },
        )
        assert create_resp.status in (200, 201), f"Failed creating token: {create_resp.status} {create_resp.text()}"

        payload = create_resp.json()
        access_token = payload["access_token"]
        token_id = _extract_token_id(payload)
        assert token_id, "Token ID must be present to revoke token"

        token_ctx = _api_context(playwright, access_token)
        try:
            before_revoke = token_ctx.get("/servers")
            assert before_revoke.status == 200, f"Token should work before revocation: {before_revoke.status} {before_revoke.text()}"

            revoke_resp = admin_api.delete(f"/tokens/{token_id}")
            assert revoke_resp.status in (200, 204), f"Failed revoking token: {revoke_resp.status} {revoke_resp.text()}"

            # Revocation cache invalidation can be asynchronous; allow brief propagation.
            deadline = time.time() + 5.0
            after_status = None
            while time.time() < deadline:
                after_revoke = token_ctx.get("/servers")
                after_status = after_revoke.status
                if after_status == 401:
                    break
                time.sleep(0.25)
        finally:
            token_ctx.dispose()

        if after_status != 401:
            # Fallback validation: ensure revocation persisted even if runtime auth cache delays enforcement.
            token_info = admin_api.get(f"/tokens/{token_id}")
            assert token_info.status == 200, f"Failed loading token after revoke: {token_info.status} {token_info.text()}"
            assert token_info.json().get("is_active") is False, "Revoked token should be marked inactive"
            pytest.skip("Token revocation persisted, but runtime rejection was not immediate in this environment.")

    def test_missing_jti_claim_enforcement_matches_runtime_setting(self, playwright: Playwright):
        now = datetime.now(timezone.utc)
        payload: dict[str, Any] = {
            "sub": "admin@example.com",
            "iss": settings.jwt_issuer,
            "aud": settings.jwt_audience,
            "iat": int(now.timestamp()),
            "exp": int((now + timedelta(minutes=5)).timestamp()),
            "user": {"email": "admin@example.com", "is_admin": True, "auth_provider": "local"},
            "token_use": "api",
            "teams": [],
            "scopes": {"server_id": None, "permissions": [], "ip_restrictions": [], "time_restrictions": {}},
        }
        if settings.embed_environment_in_tokens:
            payload["env"] = settings.environment

        signing_key = get_jwt_private_key_or_secret()
        token_without_jti = jwt.encode(payload, signing_key, algorithm=settings.jwt_algorithm)

        ctx = _api_context(playwright, token_without_jti)
        try:
            response = ctx.get("/servers")
            status_code = response.status
            response_text = response.text()
        finally:
            ctx.dispose()

        if settings.require_jti:
            assert status_code == 401, f"REQUIRE_JTI=true should reject tokens without jti, got {status_code}: {response_text}"
            assert "jti" in response_text.lower() or "invalid authentication credentials" in response_text.lower()
        else:
            assert status_code == 200, f"REQUIRE_JTI=false should allow tokens without jti, got {status_code}: {response_text}"

    def test_cookie_only_auth_is_rejected_for_api_requests(self, admin_api: APIRequestContext, playwright: Playwright):
        email = f"cookie-api-{uuid.uuid4().hex[:8]}@example.com"
        create_user_resp = admin_api.post(
            "/auth/email/admin/users",
            data={"email": email, "password": TEST_PASSWORD, "full_name": "Cookie API User"},
        )
        assert create_user_resp.status in (200, 201), f"Failed creating user: {create_user_resp.status} {create_user_resp.text()}"

        login_ctx = playwright.request.new_context(
            base_url=BASE_URL,
            extra_http_headers={"Accept": "application/json"},
        )
        cookie_ctx = None
        try:
            login_resp = login_ctx.post(
                "/auth/email/login",
                data={"email": email, "password": TEST_PASSWORD},
            )
            if login_resp.status == 404:
                pytest.skip("Email auth login endpoint unavailable in this environment")
            assert login_resp.status == 200, f"Failed user login: {login_resp.status} {login_resp.text()}"

            access_token = login_resp.json()["access_token"]

            cookie_ctx = playwright.request.new_context(
                base_url=BASE_URL,
                extra_http_headers={
                    "Accept": "application/json",
                    "Cookie": f"jwt_token={access_token}",
                },
            )
            response = cookie_ctx.get("/servers")
            status_code = response.status
            response_text = response.text()
        finally:
            login_ctx.dispose()
            if cookie_ctx:
                cookie_ctx.dispose()
            with suppress(Exception):
                admin_api.delete(f"/auth/email/admin/users/{email}")

        assert status_code == 401, f"Cookie-only API auth should be rejected, got {status_code}: {response_text}"
        assert "cookie authentication not allowed" in response_text.lower()
