"""Tests for API authentication, JWT tokens, and rate limiting."""

from __future__ import annotations

import time

from orxaq.enterprise.auth import (
    APIKeyStore,
    AuthMiddleware,
    RateLimiter,
    TokenProvider,
)


class TestAPIKeyStore:
    def test_generate_and_validate(self):
        store = APIKeyStore()
        raw_key, record = store.generate_key(name="test")
        assert raw_key.startswith("orxaq_")
        assert record.name == "test"
        assert store.validate(raw_key) is not None

    def test_invalid_key_rejected(self):
        store = APIKeyStore()
        store.generate_key(name="test")
        assert store.validate("orxaq_invalid_key_here") is None

    def test_revoke_key(self):
        store = APIKeyStore()
        raw_key, record = store.generate_key(name="to-revoke")
        assert store.validate(raw_key) is not None
        assert store.revoke(record.key_id)
        assert store.validate(raw_key) is None

    def test_revoke_nonexistent(self):
        store = APIKeyStore()
        assert not store.revoke("nonexistent")

    def test_expired_key_rejected(self):
        store = APIKeyStore()
        raw_key, record = store.generate_key(name="expiring", ttl_seconds=0.01)
        time.sleep(0.02)
        assert store.validate(raw_key) is None

    def test_scopes_preserved(self):
        store = APIKeyStore()
        scopes = frozenset({"read", "write"})
        raw_key, record = store.generate_key(name="scoped", scopes=scopes)
        validated = store.validate(raw_key)
        assert validated is not None
        assert validated.scopes == scopes

    def test_list_keys(self):
        store = APIKeyStore()
        store.generate_key(name="key1")
        store.generate_key(name="key2")
        keys = store.list_keys()
        assert len(keys) == 2
        names = {k.name for k in keys}
        assert names == {"key1", "key2"}

    def test_key_count(self):
        store = APIKeyStore()
        assert store.key_count() == 0
        store.generate_key(name="a")
        assert store.key_count() == 1


class TestTokenProvider:
    def test_issue_and_validate(self):
        tp = TokenProvider(secret="test-secret")
        token = tp.issue("user123", scopes=frozenset({"read"}))
        claims = tp.validate(token)
        assert claims is not None
        assert claims.subject == "user123"
        assert "read" in claims.scopes

    def test_expired_token_rejected(self):
        tp = TokenProvider(secret="test-secret")
        token = tp.issue("user", ttl_seconds=0.01)
        time.sleep(0.02)
        assert tp.validate(token) is None

    def test_tampered_token_rejected(self):
        tp = TokenProvider(secret="test-secret")
        token = tp.issue("user")
        # Tamper with the signature.
        parts = token.split(".")
        parts[2] = "tampered_signature"
        tampered = ".".join(parts)
        assert tp.validate(tampered) is None

    def test_wrong_secret_rejected(self):
        tp1 = TokenProvider(secret="secret1")
        tp2 = TokenProvider(secret="secret2")
        token = tp1.issue("user")
        assert tp2.validate(token) is None

    def test_malformed_token_rejected(self):
        tp = TokenProvider(secret="test")
        assert tp.validate("not.a.valid.token.format") is None
        assert tp.validate("") is None
        assert tp.validate("single") is None

    def test_extra_claims(self):
        tp = TokenProvider(secret="test")
        token = tp.issue("user", extra={"role": "admin"})
        claims = tp.validate(token)
        assert claims is not None
        assert claims.extra["role"] == "admin"

    def test_token_format(self):
        tp = TokenProvider(secret="test")
        token = tp.issue("user")
        parts = token.split(".")
        assert len(parts) == 3  # header.payload.signature


class TestRateLimiter:
    def test_allows_within_capacity(self):
        rl = RateLimiter(capacity=5, refill_rate=1.0)
        for _ in range(5):
            assert rl.allow("client1")

    def test_blocks_over_capacity(self):
        rl = RateLimiter(capacity=3, refill_rate=0.0)  # No refill.
        assert rl.allow("client1")
        assert rl.allow("client1")
        assert rl.allow("client1")
        assert not rl.allow("client1")

    def test_refills_over_time(self):
        rl = RateLimiter(capacity=2, refill_rate=100.0)  # Fast refill.
        assert rl.allow("client1")
        assert rl.allow("client1")
        time.sleep(0.02)  # Should refill ~2 tokens at 100/s.
        assert rl.allow("client1")

    def test_separate_clients(self):
        rl = RateLimiter(capacity=1, refill_rate=0.0)
        assert rl.allow("client1")
        assert not rl.allow("client1")
        assert rl.allow("client2")  # Different client, own bucket.

    def test_remaining(self):
        rl = RateLimiter(capacity=10, refill_rate=0.0)
        assert rl.remaining("new_client") == 10
        rl.allow("new_client")
        assert rl.remaining("new_client") < 10

    def test_reset(self):
        rl = RateLimiter(capacity=1, refill_rate=0.0)
        assert rl.allow("client1")
        assert not rl.allow("client1")
        rl.reset("client1")
        assert rl.allow("client1")


class TestAuthMiddleware:
    def setup_method(self):
        self.key_store = APIKeyStore()
        self.token_provider = TokenProvider(secret="test-secret")
        self.rate_limiter = RateLimiter(capacity=100, refill_rate=10)
        self.mw = AuthMiddleware(
            key_store=self.key_store,
            token_provider=self.token_provider,
            rate_limiter=self.rate_limiter,
        )

    def test_excluded_paths_bypass_auth(self):
        result = self.mw.authenticate({}, "/health")
        assert result.authenticated
        assert result.identity == "anonymous"

    def test_api_key_auth(self):
        raw_key, _ = self.key_store.generate_key(name="test")
        result = self.mw.authenticate({"X-API-Key": raw_key}, "/api/data")
        assert result.authenticated
        assert result.error == ""

    def test_invalid_api_key(self):
        result = self.mw.authenticate({"X-API-Key": "bad_key"}, "/api/data")
        assert not result.authenticated
        assert result.error == "invalid_api_key"

    def test_bearer_token_auth(self):
        token = self.token_provider.issue("admin", scopes=frozenset({"admin"}))
        result = self.mw.authenticate({"Authorization": f"Bearer {token}"}, "/api/data")
        assert result.authenticated
        assert result.identity == "admin"
        assert "admin" in result.scopes

    def test_invalid_bearer_token(self):
        result = self.mw.authenticate({"Authorization": "Bearer invalid.token.here"}, "/api/data")
        assert not result.authenticated
        assert result.error == "invalid_token"

    def test_no_credentials(self):
        result = self.mw.authenticate({}, "/api/data")
        assert not result.authenticated
        assert result.error == "no_credentials"

    def test_rate_limiting(self):
        mw = AuthMiddleware(
            key_store=self.key_store,
            token_provider=self.token_provider,
            rate_limiter=RateLimiter(capacity=2, refill_rate=0.0),
        )
        raw_key, _ = self.key_store.generate_key(name="limited")
        headers = {"X-API-Key": raw_key}

        r1 = mw.authenticate(headers, "/api/data")
        assert r1.authenticated and not r1.rate_limited

        r2 = mw.authenticate(headers, "/api/data")
        assert r2.authenticated and not r2.rate_limited

        r3 = mw.authenticate(headers, "/api/data")
        assert r3.authenticated and r3.rate_limited

    def test_defaults_work(self):
        mw = AuthMiddleware()
        result = mw.authenticate({}, "/health")
        assert result.authenticated
