"""Shared licensing API rate-limit and trusted-proxy regression tests."""

from __future__ import annotations

import asyncio
from datetime import datetime, timedelta, timezone
import json
from pathlib import Path
import unittest
from unittest.mock import patch

from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
import httpx2
from sqlalchemy import create_engine, select
from sqlalchemy.orm import Session
from sqlalchemy.pool import StaticPool

from licensing_server.app.api import create_app
from licensing_server.app.auth import AuthenticatedPrincipal
from licensing_server.app.config import ServerSettings
from licensing_server.app.database import Base, create_session_factory
from licensing_server.app.errors import ServerErrorCode, ServerLicensingError
from licensing_server.app.middleware import BOUNDED_REQUEST_BODY_STATE_KEY
from licensing_server.app.models import CatalogRelease, RateLimitBucket, User
from licensing_server.app.rate_limits import (
    DatabaseRateLimiter,
    RateLimitQuota,
    RequestRateLimitContext,
    extract_rate_limit_context,
)
from licensing_server.app.security import Ed25519SnapshotSigner


TEST_DIGEST_PEPPER = b"rate-limit-test-pepper-material-32-bytes"
TEST_NOW = datetime(2026, 8, 15, 23, 30, tzinfo=timezone.utc)


class StaticTokenValidator:
    def validate(self, token: str) -> AuthenticatedPrincipal:
        if token != "customer-token":
            raise RuntimeError("unexpected test token")
        return AuthenticatedPrincipal(
            issuer="https://identity.test",
            subject="customer-subject",
            scopes=frozenset(),
            verified_email="customer@example.test",
            authentication_time=TEST_NOW,
        )


def _scope(
    peer: str,
    *,
    path: str = "/v1/serials/redeem",
    body: dict[str, object] | None = None,
    forwarded_for: str | None = None,
) -> dict[str, object]:
    if not isinstance(peer, str) or not peer:
        raise ValueError("peer must be a non-empty string")
    if not isinstance(path, str) or not path.startswith("/"):
        raise ValueError("path must be an absolute path")
    if body is not None and not isinstance(body, dict):
        raise TypeError("body must be a dictionary or None")
    if forwarded_for is not None and not isinstance(forwarded_for, str):
        raise TypeError("forwarded_for must be a string or None")
    headers: list[tuple[bytes, bytes]] = [(b"content-type", b"application/json")]
    if forwarded_for is not None:
        headers.append((b"x-forwarded-for", forwarded_for.encode("ascii")))
    encoded_body = json.dumps(body or {}).encode("utf-8")
    return {
        "type": "http",
        "method": "POST",
        "path": path,
        "client": (peer, 49152),
        "headers": headers,
        "state": {BOUNDED_REQUEST_BODY_STATE_KEY: encoded_body},
    }


class LicensingRateLimitTests(unittest.IsolatedAsyncioTestCase):
    def setUp(self) -> None:
        self.engine = create_engine(
            "sqlite+pysqlite:///:memory:",
            connect_args={"check_same_thread": False},
            poolclass=StaticPool,
        )
        Base.metadata.create_all(self.engine)
        self.session_factory = create_session_factory(self.engine)

    def tearDown(self) -> None:
        self.engine.dispose()

    def test_trusted_proxy_walk_and_semantic_keys_are_bounded(self) -> None:
        body = {
            "serial": "MSA1-ABCDEF-SECRET-WITH-IGNORED-TAIL",
            "device": {"deviceKeyThumbprint": "device.thumbprint.test"},
        }
        trusted = extract_rate_limit_context(
            _scope(
                "127.0.0.1",
                body=body,
                forwarded_for="203.0.113.9, 10.2.3.4",
            ),
            ("127.0.0.1/32", "10.0.0.0/8"),
        )
        self.assertEqual(trusted.source_address, "203.0.113.9")
        self.assertEqual(trusted.serial_lookup_prefix, "MSA1ABCDEF")
        self.assertEqual(trusted.device_key_thumbprint, "device.thumbprint.test")

        untrusted = extract_rate_limit_context(
            _scope(
                "198.51.100.10",
                body=body,
                forwarded_for="203.0.113.200",
            ),
            ("127.0.0.1/32",),
        )
        self.assertEqual(untrusted.source_address, "198.51.100.10")

    def test_shared_bucket_blocks_across_instances_without_raw_identifiers(self) -> None:
        current_time = [TEST_NOW]
        quotas = {
            "public.read": (
                RateLimitQuota("test.public.source", "source", 2, 60),
            )
        }
        first = DatabaseRateLimiter(
            self.session_factory,
            TEST_DIGEST_PEPPER,
            request_quotas=quotas,
            principal_quotas={},
            now_provider=lambda: current_time[0],
        )
        second = DatabaseRateLimiter(
            self.session_factory,
            TEST_DIGEST_PEPPER,
            request_quotas=quotas,
            principal_quotas={},
            now_provider=lambda: current_time[0],
        )
        context = RequestRateLimitContext("public.read", "203.0.113.41")
        first.enforce_request(context)
        second.enforce_request(context)
        with self.assertRaises(ServerLicensingError) as caught:
            first.enforce_request(context)
        self.assertEqual(caught.exception.code, ServerErrorCode.RATE_LIMITED)
        self.assertEqual(caught.exception.status_code, 429)
        self.assertEqual(caught.exception.retry_after_seconds, 60)

        with Session(self.engine) as session:
            row = session.scalar(select(RateLimitBucket))
            self.assertIsNotNone(row)
            assert row is not None
            self.assertEqual(row.request_count, 3)
            persisted = json.dumps(
                {
                    "id": row.id,
                    "policy": row.policy,
                    "requestCount": row.request_count,
                },
                sort_keys=True,
            )
            self.assertNotIn(context.source_address, persisted)
            self.assertNotIn("customer-token", persisted)

        current_time[0] += timedelta(seconds=61)
        second.enforce_request(context)
        with Session(self.engine) as session:
            row = session.scalar(select(RateLimitBucket))
            assert row is not None
            self.assertEqual(row.request_count, 1)

    def test_device_bucket_uses_only_verified_thumbprint(self) -> None:
        limiter = DatabaseRateLimiter(
            self.session_factory,
            TEST_DIGEST_PEPPER,
            request_quotas={
                "trial.start": (
                    RateLimitQuota("test.trial.source", "source", 10, 60),
                    RateLimitQuota("test.trial.device", "device", 1, 60),
                )
            },
            principal_quotas={},
            now_provider=lambda: TEST_NOW,
        )
        context = RequestRateLimitContext(
            "trial.start",
            "203.0.113.50",
            device_key_thumbprint="unverified-attacker-value",
        )
        limiter.enforce_request(context)
        limiter.enforce_verified_device(context, "verified-device-thumbprint")
        with self.assertRaises(ServerLicensingError) as caught:
            limiter.enforce_verified_device(context, "verified-device-thumbprint")
        self.assertEqual(caught.exception.code, ServerErrorCode.RATE_LIMITED)
        with Session(self.engine) as session:
            policies = set(session.scalars(select(RateLimitBucket.policy)))
        self.assertEqual(policies, {"test.trial.source", "test.trial.device"})

    def test_missing_shared_table_fails_closed(self) -> None:
        unavailable_engine = create_engine(
            "sqlite+pysqlite:///:memory:",
            connect_args={"check_same_thread": False},
            poolclass=StaticPool,
        )
        self.addCleanup(unavailable_engine.dispose)
        limiter = DatabaseRateLimiter(
            create_session_factory(unavailable_engine),
            TEST_DIGEST_PEPPER,
            request_quotas={
                "public.read": (
                    RateLimitQuota("test.unavailable", "source", 10, 60),
                )
            },
            principal_quotas={},
            now_provider=lambda: TEST_NOW,
        )
        with self.assertRaises(ServerLicensingError) as caught:
            limiter.enforce_request(
                RequestRateLimitContext("public.read", "203.0.113.90")
            )
        self.assertEqual(caught.exception.code, ServerErrorCode.RATE_LIMIT_UNAVAILABLE)
        self.assertEqual(caught.exception.status_code, 503)
        self.assertTrue(caught.exception.retryable)

    async def test_api_returns_stable_retryable_429_and_retry_after(self) -> None:
        quotas = {
            "public.read": (
                RateLimitQuota("test.api.public", "source", 2, 60),
            )
        }
        limiter = DatabaseRateLimiter(
            self.session_factory,
            TEST_DIGEST_PEPPER,
            request_quotas=quotas,
            principal_quotas={},
            now_provider=lambda: TEST_NOW,
        )
        signer = Ed25519SnapshotSigner(
            "license.test.rate_limit",
            Ed25519PrivateKey.generate(),
        )
        settings = ServerSettings(
            database_url="sqlite+pysqlite:///:memory:",
            environment_name="test",
            signing_key_id=signer.key_id,
            signing_private_key_path=Path("unused-test-signing-key"),
            serial_pepper=b"s" * 32,
            fingerprint_pepper=TEST_DIGEST_PEPPER,
            oidc_issuer="https://identity.test",
            oidc_audience="licensing-api",
            oidc_jwks_url="https://identity.test/jwks.json",
        )
        app = create_app(
            settings,
            engine=self.engine,
            signer=signer,
            token_validator=StaticTokenValidator(),
            rate_limiter=limiter,
        )
        with Session(self.engine) as session, session.begin():
            catalog = app.state.catalog
            session.add_all(
                (
                    User(
                        id="user.test.rate_limit_admin",
                        external_issuer="https://identity.test",
                        external_subject="rate-limit-admin",
                        verified_email=None,
                        status="active",
                        is_server_admin=True,
                    ),
                    CatalogRelease(
                        id="catalog_release.test.rate_limit",
                        product_id=catalog.product_id,
                        revision=catalog.revision,
                        document_json=catalog.to_mapping(),
                        document_digest=bytes.fromhex(catalog.sha256()),
                        status="active",
                        staged_by_user_id="user.test.rate_limit_admin",
                        staged_at=TEST_NOW - timedelta(days=1),
                        published_by_user_id="user.test.rate_limit_admin",
                        published_at=TEST_NOW - timedelta(days=1),
                        retired_at=None,
                    ),
                )
            )

        async def direct_test_call(function, *args, **kwargs):
            return function(*args, **kwargs)

        with (
            patch("fastapi.routing.run_in_threadpool", direct_test_call),
            patch(
                "fastapi.dependencies.utils.run_in_threadpool",
                direct_test_call,
            ),
        ):
            transport = httpx2.ASGITransport(app=app)
            async with httpx2.AsyncClient(
                transport=transport,
                base_url="https://licensing.test",
            ) as client:
                first = await asyncio.wait_for(client.get("/v1/catalog/public"), 5)
                second = await asyncio.wait_for(client.get("/v1/catalog/public"), 5)
                denied = await asyncio.wait_for(client.get("/v1/catalog/public"), 5)
        self.assertEqual(first.status_code, 200)
        self.assertEqual(second.status_code, 200)
        self.assertEqual(denied.status_code, 429)
        self.assertEqual(denied.json()["code"], ServerErrorCode.RATE_LIMITED.value)
        self.assertTrue(denied.json()["retryable"])
        self.assertEqual(denied.headers["retry-after"], "60")

    async def test_validated_principal_bucket_is_shared_and_enforced(self) -> None:
        request_quotas = {
            "account.read": (
                RateLimitQuota("test.account.source", "source", 10, 60),
            )
        }
        principal_quotas = {
            "account.read": (
                RateLimitQuota("test.account.principal", "principal", 1, 60),
            )
        }
        limiter = DatabaseRateLimiter(
            self.session_factory,
            TEST_DIGEST_PEPPER,
            request_quotas=request_quotas,
            principal_quotas=principal_quotas,
            now_provider=lambda: TEST_NOW,
        )
        signer = Ed25519SnapshotSigner(
            "license.test.principal_limit",
            Ed25519PrivateKey.generate(),
        )
        settings = ServerSettings(
            database_url="sqlite+pysqlite:///:memory:",
            environment_name="test",
            signing_key_id=signer.key_id,
            signing_private_key_path=Path("unused-test-signing-key"),
            serial_pepper=b"s" * 32,
            fingerprint_pepper=TEST_DIGEST_PEPPER,
            oidc_issuer="https://identity.test",
            oidc_audience="licensing-api",
            oidc_jwks_url="https://identity.test/jwks.json",
        )
        app = create_app(
            settings,
            engine=self.engine,
            signer=signer,
            token_validator=StaticTokenValidator(),
            rate_limiter=limiter,
        )

        async def direct_test_call(function, *args, **kwargs):
            return function(*args, **kwargs)

        with (
            patch("fastapi.routing.run_in_threadpool", direct_test_call),
            patch(
                "fastapi.dependencies.utils.run_in_threadpool",
                direct_test_call,
            ),
        ):
            transport = httpx2.ASGITransport(app=app)
            async with httpx2.AsyncClient(
                transport=transport,
                base_url="https://licensing.test",
            ) as client:
                headers = {"Authorization": "Bearer customer-token"}
                first = await asyncio.wait_for(
                    client.get("/v1/me/licenses", headers=headers),
                    5,
                )
                denied = await asyncio.wait_for(
                    client.get("/v1/me/licenses", headers=headers),
                    5,
                )
        self.assertEqual(first.status_code, 200)
        self.assertEqual(denied.status_code, 429)
        self.assertEqual(denied.json()["code"], ServerErrorCode.RATE_LIMITED.value)


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