"""
tests/test_auth.py — Tests for Cortex JWT + API key authentication (#542).

Covers:
    - Valid JWT → correct owner_id extracted
    - Expired JWT → AuthError
    - Tampered JWT (wrong signature) → AuthError
    - Valid API key → success
    - Invalid API key → AuthError
    - Open access mode (no secrets configured) → success with 'default' owner
    - Authorization: Bearer header parsing
    - X-Cortex-API-Key header parsing
    - X-API-Key legacy header parsing
    - Public endpoints skip auth
    - Constant-time comparison used for both JWT sig and API key
    - FastAPI integration: 401 when auth fails and require_auth=True
    - request.state.owner_id is set correctly by verify_api_key dependency
"""
from __future__ import annotations

import base64
import hashlib
import hmac
import json
import time
from typing import Any, Dict, Optional
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from cortex.api.auth import AuthError, AuthProvider, AuthResult, _b64url_decode, _b64url_encode


# ---------------------------------------------------------------------------
# JWT test helpers
# ---------------------------------------------------------------------------

def _b64url_encode_str(s: str) -> str:
    """Base64url-encode a UTF-8 string."""
    return _b64url_encode(s.encode())


def make_test_jwt(payload: dict, secret: str, algorithm: str = "HS256") -> str:
    """
    Create a minimal HS256 JWT for testing.

    Args:
        payload: Claims dict (must include at least 'sub' and 'exp').
        secret:  HMAC secret key.

    Returns:
        Dot-separated JWT string.
    """
    header = _b64url_encode_str(json.dumps({"alg": algorithm, "typ": "JWT"}))
    body = _b64url_encode_str(json.dumps(payload))
    signing_input = f"{header}.{body}".encode()
    sig_bytes = hmac.new(secret.encode(), signing_input, hashlib.sha256).digest()
    sig = _b64url_encode(sig_bytes)
    return f"{header}.{body}.{sig}"


def _future_exp(seconds: int = 3600) -> int:
    return int(time.time()) + seconds


def _past_exp(seconds: int = 3600) -> int:
    return int(time.time()) - seconds


# ---------------------------------------------------------------------------
# AuthResult dataclass
# ---------------------------------------------------------------------------

class TestAuthResult:
    def test_fields(self):
        r = AuthResult(owner_id="user-123", method="jwt", plan="pro")
        assert r.owner_id == "user-123"
        assert r.method == "jwt"
        assert r.plan == "pro"

    def test_plan_defaults_to_none(self):
        r = AuthResult(owner_id="u", method="open")
        assert r.plan is None


# ---------------------------------------------------------------------------
# JWT validation
# ---------------------------------------------------------------------------

class TestJWTValidation:
    SECRET = "test-jwt-secret-supersecret"

    def _provider(self) -> AuthProvider:
        return AuthProvider(jwt_secret=self.SECRET)

    def test_valid_jwt_returns_correct_owner_id(self):
        jwt = make_test_jwt({"sub": "user-uuid-abc", "exp": _future_exp()}, self.SECRET)
        result = self._provider().validate_jwt(jwt)
        assert result.owner_id == "user-uuid-abc"
        assert result.method == "jwt"
        assert result.plan is None

    def test_valid_jwt_with_plan_claim(self):
        jwt = make_test_jwt(
            {"sub": "user-uuid-pro", "exp": _future_exp(), "plan": "pro"},
            self.SECRET,
        )
        result = self._provider().validate_jwt(jwt)
        assert result.owner_id == "user-uuid-pro"
        assert result.plan == "pro"

    def test_valid_jwt_with_user_plan_claim(self):
        """Supabase custom claims may use 'user_plan' instead of 'plan'."""
        jwt = make_test_jwt(
            {"sub": "user-uuid-scale", "exp": _future_exp(), "user_plan": "scale"},
            self.SECRET,
        )
        result = self._provider().validate_jwt(jwt)
        assert result.plan == "scale"

    def test_expired_jwt_raises_auth_error(self):
        jwt = make_test_jwt({"sub": "user-xyz", "exp": _past_exp()}, self.SECRET)
        with pytest.raises(AuthError, match="expired"):
            self._provider().validate_jwt(jwt)

    def test_tampered_jwt_wrong_signature_raises_auth_error(self):
        jwt = make_test_jwt({"sub": "user-xyz", "exp": _future_exp()}, self.SECRET)
        # Tamper: replace last char of signature
        parts = jwt.split(".")
        parts[2] = parts[2][:-1] + ("A" if parts[2][-1] != "A" else "B")
        tampered = ".".join(parts)
        with pytest.raises(AuthError, match="signature"):
            self._provider().validate_jwt(tampered)

    def test_tampered_jwt_wrong_secret_raises_auth_error(self):
        jwt = make_test_jwt({"sub": "user-xyz", "exp": _future_exp()}, self.SECRET)
        provider = AuthProvider(jwt_secret="wrong-secret")
        with pytest.raises(AuthError, match="signature"):
            provider.validate_jwt(jwt)

    def test_invalid_jwt_format_missing_parts(self):
        with pytest.raises(AuthError, match="format"):
            self._provider().validate_jwt("only.two")

    def test_jwt_missing_sub_claim(self):
        jwt = make_test_jwt({"exp": _future_exp()}, self.SECRET)
        with pytest.raises(AuthError, match="sub"):
            self._provider().validate_jwt(jwt)

    def test_jwt_missing_exp_claim(self):
        jwt = make_test_jwt({"sub": "user-xyz"}, self.SECRET)
        with pytest.raises(AuthError, match="exp"):
            self._provider().validate_jwt(jwt)

    def test_jwt_secret_not_configured_raises_auth_error(self):
        provider = AuthProvider(jwt_secret=None)
        with pytest.raises(AuthError, match="not configured"):
            provider.validate_jwt("any.token.value")

    def test_constant_time_comparison_used(self):
        """
        Verify that hmac.compare_digest is used (constant-time).
        The implementation should call hmac.compare_digest, not ==.
        We check this by ensuring a near-match still fails (no short-circuit).
        """
        jwt = make_test_jwt({"sub": "user-xyz", "exp": _future_exp()}, self.SECRET)
        # Construct a token signed with a secret that differs only in the last byte
        # If == were used instead of compare_digest, this would be timing-vulnerable.
        bad_secret = self.SECRET[:-1] + chr(ord(self.SECRET[-1]) ^ 1)
        provider = AuthProvider(jwt_secret=bad_secret)
        with pytest.raises(AuthError):
            provider.validate_jwt(jwt)


# ---------------------------------------------------------------------------
# API key validation
# ---------------------------------------------------------------------------

class TestAPIKeyValidation:
    KEY = "my-api-key-12345"

    def _provider(self) -> AuthProvider:
        return AuthProvider(api_key=self.KEY)

    def test_valid_api_key_returns_auth_result(self):
        result = self._provider().validate_api_key(self.KEY)
        assert result.method == "api_key"
        assert result.owner_id == "default"

    def test_invalid_api_key_raises_auth_error(self):
        with pytest.raises(AuthError, match="Invalid"):
            self._provider().validate_api_key("wrong-key")

    def test_api_key_not_configured_raises_auth_error(self):
        provider = AuthProvider(api_key=None)
        with pytest.raises(AuthError, match="not configured"):
            provider.validate_api_key(self.KEY)

    def test_constant_time_comparison_for_api_key(self):
        """Verify near-match API key still fails — constant-time compare."""
        bad_key = self.KEY[:-1] + chr(ord(self.KEY[-1]) ^ 1)
        with pytest.raises(AuthError):
            self._provider().validate_api_key(bad_key)


# ---------------------------------------------------------------------------
# authenticate() — header dispatch
# ---------------------------------------------------------------------------

class TestAuthenticate:
    JWT_SECRET = "jwt-secret-xyz"
    API_KEY = "api-key-abc"

    def _provider(
        self,
        jwt_secret: Optional[str] = None,
        api_key: Optional[str] = None,
    ) -> AuthProvider:
        return AuthProvider(jwt_secret=jwt_secret, api_key=api_key)

    # ---- JWT path ----------------------------------------------------------

    def test_authorization_bearer_uses_jwt(self):
        jwt = make_test_jwt({"sub": "jwt-user", "exp": _future_exp()}, self.JWT_SECRET)
        provider = self._provider(jwt_secret=self.JWT_SECRET)
        result = provider.authenticate({"Authorization": f"Bearer {jwt}"})
        assert result.method == "jwt"
        assert result.owner_id == "jwt-user"

    def test_bearer_token_with_jwt_secret_validates(self):
        jwt = make_test_jwt({"sub": "user-123", "exp": _future_exp()}, self.JWT_SECRET)
        provider = self._provider(jwt_secret=self.JWT_SECRET, api_key=self.API_KEY)
        result = provider.authenticate({"Authorization": f"Bearer {jwt}"})
        assert result.method == "jwt"

    def test_bearer_with_expired_jwt_raises_auth_error(self):
        jwt = make_test_jwt({"sub": "user-xyz", "exp": _past_exp()}, self.JWT_SECRET)
        provider = self._provider(jwt_secret=self.JWT_SECRET)
        with pytest.raises(AuthError):
            provider.authenticate({"Authorization": f"Bearer {jwt}"})

    def test_bearer_without_jwt_secret_falls_through_to_open(self):
        """If no JWT secret configured, Bearer token is ignored → open access."""
        fake_token = "header.payload.sig"
        provider = self._provider(jwt_secret=None, api_key=None)
        result = provider.authenticate({"Authorization": f"Bearer {fake_token}"})
        assert result.method == "open"

    def test_lowercase_authorization_header_works(self):
        """Case-insensitive header lookup."""
        jwt = make_test_jwt({"sub": "user-lower", "exp": _future_exp()}, self.JWT_SECRET)
        provider = self._provider(jwt_secret=self.JWT_SECRET)
        # Starlette Headers lower-cases internally; simulate plain dict with lowercase
        result = provider.authenticate({"authorization": f"Bearer {jwt}"})
        assert result.method == "jwt"
        assert result.owner_id == "user-lower"

    # ---- API key paths -----------------------------------------------------

    def test_x_cortex_api_key_header_used(self):
        provider = self._provider(api_key=self.API_KEY)
        result = provider.authenticate({"X-Cortex-API-Key": self.API_KEY})
        assert result.method == "api_key"

    def test_x_cortex_api_key_lowercase_header_works(self):
        provider = self._provider(api_key=self.API_KEY)
        result = provider.authenticate({"x-cortex-api-key": self.API_KEY})
        assert result.method == "api_key"

    def test_legacy_x_api_key_header_works(self):
        """Backward compat: old self-host setups use X-API-Key."""
        provider = self._provider(api_key=self.API_KEY)
        result = provider.authenticate({"X-API-Key": self.API_KEY})
        assert result.method == "api_key"

    def test_invalid_x_cortex_api_key_raises_auth_error(self):
        provider = self._provider(api_key=self.API_KEY)
        with pytest.raises(AuthError):
            provider.authenticate({"X-Cortex-API-Key": "wrong-key"})

    # ---- JWT beats API key -------------------------------------------------

    def test_jwt_takes_priority_over_api_key_header(self):
        jwt = make_test_jwt({"sub": "jwt-wins", "exp": _future_exp()}, self.JWT_SECRET)
        provider = self._provider(jwt_secret=self.JWT_SECRET, api_key=self.API_KEY)
        result = provider.authenticate({
            "Authorization": f"Bearer {jwt}",
            "X-Cortex-API-Key": self.API_KEY,
        })
        assert result.method == "jwt"
        assert result.owner_id == "jwt-wins"

    # ---- Open access mode --------------------------------------------------

    def test_open_access_no_secrets_configured(self):
        """When no secrets are configured, any request is open access."""
        provider = self._provider(jwt_secret=None, api_key=None)
        result = provider.authenticate({})
        assert result.method == "open"
        assert result.owner_id == "default"

    def test_open_access_no_headers_provided(self):
        """Even with secrets configured, missing credentials → open access."""
        provider = self._provider(jwt_secret=self.JWT_SECRET, api_key=self.API_KEY)
        result = provider.authenticate({})
        assert result.method == "open"

    def test_open_access_empty_headers(self):
        provider = self._provider()
        result = provider.authenticate({})
        assert result.method == "open"
        assert result.owner_id == "default"


# ---------------------------------------------------------------------------
# FastAPI integration — http_server + verify_api_key dependency
# ---------------------------------------------------------------------------

pytest.importorskip("fastapi", reason="fastapi not installed")
pytest.importorskip("httpx", reason="httpx not installed")

from contextlib import contextmanager  # noqa: E402

from fastapi.testclient import TestClient  # noqa: E402


@contextmanager
def _make_test_client(
    jwt_secret: Optional[str] = None,
    api_key: Optional[str] = None,
    require_auth: bool = False,
    env_overrides: Optional[dict] = None,
):
    """Create a TestClient with a minimal fake plugin and configurable auth."""
    from cortex.api.http_server import create_app

    class _FakeConfig:
        owner_id: str = "test-owner"
        max_remember_body_bytes: int = 1_000_000

    class _FakeStorage:
        async def initialize(self): pass
        async def close(self): pass
        async def log_request(self, **_): pass

    class _FakePlugin:
        def __init__(self):
            self.storage = _FakeStorage()
            self.config = _FakeConfig()

    env = {}
    if jwt_secret:
        env["SUPABASE_JWT_SECRET"] = jwt_secret
    if api_key:
        env["CORTEX_API_KEY"] = api_key
    if require_auth:
        env["CORTEX_AUTH_REQUIRE"] = "true"
    if env_overrides:
        env.update(env_overrides)

    with patch.dict("os.environ", env, clear=False):
        with patch("cortex.api.plugin.CortexPlugin") as MockPlugin:
            MockPlugin.return_value = _FakePlugin()
            app = create_app()
            with TestClient(app, raise_server_exceptions=False) as client:
                yield client


class TestFastAPIAuthIntegration:
    JWT_SECRET = "integration-test-secret"
    API_KEY = "integration-api-key"

    def test_valid_jwt_bearer_allows_access(self):
        jwt = make_test_jwt({"sub": "int-user-1", "exp": _future_exp()}, self.JWT_SECRET)
        with _make_test_client(jwt_secret=self.JWT_SECRET) as client:
            resp = client.get(
                "/api/v1/health",
                headers={"Authorization": f"Bearer {jwt}"},
            )
            # Must NOT be 401 — auth passes, any other status is from the handler
            assert resp.status_code != 401

    def test_valid_api_key_header_allows_access(self):
        with _make_test_client(api_key=self.API_KEY) as client:
            resp = client.get(
                "/api/v1/health",
                headers={"X-API-Key": self.API_KEY},
            )
            # Must NOT be 401
            assert resp.status_code != 401

    def test_invalid_api_key_returns_401(self):
        with _make_test_client(api_key=self.API_KEY) as client:
            resp = client.get(
                "/api/v1/memories",
                headers={"X-API-Key": "wrong-key"},
            )
            assert resp.status_code == 401

    def test_expired_jwt_returns_401(self):
        jwt = make_test_jwt({"sub": "user-xyz", "exp": _past_exp()}, self.JWT_SECRET)
        with _make_test_client(jwt_secret=self.JWT_SECRET) as client:
            resp = client.get(
                "/api/v1/memories",
                headers={"Authorization": f"Bearer {jwt}"},
            )
            assert resp.status_code == 401

    def test_tampered_jwt_returns_401(self):
        jwt = make_test_jwt({"sub": "user-xyz", "exp": _future_exp()}, self.JWT_SECRET)
        parts = jwt.split(".")
        parts[2] = parts[2][:-1] + ("A" if parts[2][-1] != "A" else "B")
        tampered = ".".join(parts)
        with _make_test_client(jwt_secret=self.JWT_SECRET) as client:
            resp = client.get(
                "/api/v1/memories",
                headers={"Authorization": f"Bearer {tampered}"},
            )
            assert resp.status_code == 401

    def test_open_access_no_secrets_succeeds(self):
        """No secrets configured → open access, any request goes through auth gate."""
        with _make_test_client() as client:
            resp = client.get("/api/v1/health")
            # Should not be 401 — open access
            assert resp.status_code != 401

    def test_public_endpoints_skip_auth(self):
        """Health endpoint is always reachable even without credentials (JWT mode)."""
        # Public endpoints only exempt in JWT mode — use jwt_secret to activate
        with _make_test_client(jwt_secret=self.JWT_SECRET) as client:
            # /api/v1/health is in public_endpoints, no auth header needed
            resp = client.get("/api/v1/health")
            assert resp.status_code != 401

    def test_docs_endpoint_skips_auth(self):
        with _make_test_client(jwt_secret=self.JWT_SECRET) as client:
            resp = client.get("/docs")
            assert resp.status_code != 401

    def test_openapi_json_skips_auth(self):
        with _make_test_client(jwt_secret=self.JWT_SECRET) as client:
            resp = client.get("/openapi.json")
            assert resp.status_code != 401

    def test_x_cortex_api_key_header_parsed(self):
        """New X-Cortex-API-Key header is accepted by the auth gate."""
        # Use JWT mode so we can pair jwt_secret + api_key; api_key-only mode
        # accepts both X-API-Key and X-Cortex-API-Key for backward compat.
        with _make_test_client(jwt_secret=self.JWT_SECRET, api_key=self.API_KEY) as client:
            resp = client.get(
                "/api/v1/memories",
                headers={"X-Cortex-API-Key": self.API_KEY},
            )
            # Either 200 (found) or 422/404 (endpoint variation) — not 401
            assert resp.status_code != 401

    def test_missing_credentials_with_require_auth_returns_401(self):
        """When require_auth=True and no credentials, must return 401."""
        # Requires jwt_secret (JWT mode) + CORTEX_AUTH_REQUIRE=true
        with _make_test_client(
            jwt_secret=self.JWT_SECRET, require_auth=True
        ) as client:
            resp = client.get("/api/v1/memories")  # no auth header
            assert resp.status_code == 401

    def test_detail_message_on_auth_required(self):
        """401 response should include a useful detail message."""
        with _make_test_client(
            jwt_secret=self.JWT_SECRET, require_auth=True
        ) as client:
            resp = client.get("/api/v1/memories")
            assert resp.status_code == 401
            body = resp.json()
            assert "detail" in body


# ---------------------------------------------------------------------------
# b64url helpers (edge cases)
# ---------------------------------------------------------------------------

class TestB64Helpers:
    def test_encode_decode_roundtrip(self):
        original = b"hello world! \x00\xff"
        encoded = _b64url_encode(original)
        decoded = _b64url_decode(encoded)
        assert decoded == original

    def test_no_padding_in_encoded(self):
        encoded = _b64url_encode(b"test")
        assert "=" not in encoded

    def test_decode_with_missing_padding(self):
        """Decoder should handle tokens without padding."""
        data = b"supabase-user-id"
        encoded = base64.urlsafe_b64encode(data).decode()
        # Strip all padding
        stripped = encoded.rstrip("=")
        assert _b64url_decode(stripped) == data
