"""Tests for M4: Enforcer — admission control, fairness, rate limiting, admin API."""

from __future__ import annotations

from radix_edge.enforcer.admission import check_admission, get_user_quota, set_user_quota
from radix_edge.enforcer.fairness import get_user_usage, get_all_user_usage, get_user_priority
from radix_edge.enforcer.rate_limiter import RateLimiter


# --- Admission Control ---

class TestAdmission:
    def test_admit_within_limits(self, db):
        allowed, reason = check_admission("user-a", 1, db)
        assert allowed is True
        assert reason == "admitted"

    def test_reject_concurrent_limit(self, db):
        # Default max_concurrent = 2
        db.execute("INSERT INTO jobs (job_id, user_id, image, status) VALUES ('j1', 'user-a', 'img', 'running')")
        db.execute("INSERT INTO jobs (job_id, user_id, image, status) VALUES ('j2', 'user-a', 'img', 'queued')")
        db.commit()

        allowed, reason = check_admission("user-a", 1, db)
        assert allowed is False
        assert "Concurrent job limit" in reason

    def test_reject_gpu_limit(self, db):
        # Default max_gpus_total = 4
        db.execute("INSERT INTO jobs (job_id, user_id, image, status, num_gpus) VALUES ('j1', 'user-a', 'img', 'running', 3)")
        db.commit()

        # 3 + 2 = 5 > 4
        allowed, reason = check_admission("user-a", 2, db)
        assert allowed is False
        assert "GPU limit" in reason

    def test_admit_with_custom_quota(self, db):
        set_user_quota("user-a", max_concurrent=5, max_gpus=8, db_conn=db)
        db.commit()

        # Add 4 jobs
        for i in range(4):
            db.execute(f"INSERT INTO jobs (job_id, user_id, image, status) VALUES ('j{i}', 'user-a', 'img', 'running')")
        db.commit()

        # 4 < 5 concurrent limit
        allowed, _ = check_admission("user-a", 1, db)
        assert allowed is True

    def test_other_users_dont_count(self, db):
        db.execute("INSERT INTO jobs (job_id, user_id, image, status) VALUES ('j1', 'user-b', 'img', 'running')")
        db.execute("INSERT INTO jobs (job_id, user_id, image, status) VALUES ('j2', 'user-b', 'img', 'running')")
        db.commit()

        # user-a has no jobs
        allowed, _ = check_admission("user-a", 1, db)
        assert allowed is True


class TestQuotaManagement:
    def test_get_default_quota(self, db):
        quota = get_user_quota("new-user", db)
        assert quota["max_concurrent_gpu_jobs"] == 2
        assert quota["max_gpus_total"] == 4
        assert quota["fair_share_weight"] == 1.0

    def test_set_and_get_quota(self, db):
        set_user_quota("user-a", max_concurrent=10, max_gpus=16, fair_share_weight=2.0, db_conn=db)
        db.commit()

        quota = get_user_quota("user-a", db)
        assert quota["max_concurrent_gpu_jobs"] == 10
        assert quota["max_gpus_total"] == 16
        assert quota["fair_share_weight"] == 2.0

    def test_update_partial_quota(self, db):
        set_user_quota("user-a", max_concurrent=5, db_conn=db)
        db.commit()

        # Update only one field
        set_user_quota("user-a", max_gpus=20, db_conn=db)
        db.commit()

        quota = get_user_quota("user-a", db)
        assert quota["max_concurrent_gpu_jobs"] == 5  # Unchanged
        assert quota["max_gpus_total"] == 20  # Updated


# --- Fairness ---

class TestFairness:
    def test_no_usage(self, db):
        usage = get_user_usage("new-user", db_conn=db)
        assert usage["gpu_seconds"] == 0.0
        assert usage["job_count"] == 0

    def test_usage_tracking(self, db):
        db.execute(
            """INSERT INTO jobs (job_id, user_id, image, status, runtime_seconds, num_gpus, completed_at)
            VALUES ('j1', 'user-a', 'img', 'succeeded', 100.0, 2, datetime('now'))"""
        )
        db.execute(
            """INSERT INTO jobs (job_id, user_id, image, status, runtime_seconds, num_gpus, completed_at)
            VALUES ('j2', 'user-a', 'img', 'succeeded', 50.0, 1, datetime('now'))"""
        )
        db.commit()

        usage = get_user_usage("user-a", db_conn=db)
        assert usage["gpu_seconds"] == 250.0  # 100*2 + 50*1
        assert usage["job_count"] == 2

    def test_all_user_usage(self, db):
        db.execute(
            "INSERT INTO jobs (job_id, user_id, image, status, runtime_seconds, num_gpus, completed_at) VALUES ('j1', 'a', 'img', 'succeeded', 100, 1, datetime('now'))"
        )
        db.execute(
            "INSERT INTO jobs (job_id, user_id, image, status, runtime_seconds, num_gpus, completed_at) VALUES ('j2', 'b', 'img', 'succeeded', 200, 1, datetime('now'))"
        )
        db.commit()

        all_usage = get_all_user_usage(db_conn=db)
        assert len(all_usage) == 2
        # b has more usage, should be listed first (DESC)
        assert all_usage[0]["user_id"] == "b"

    def test_user_priority_no_usage(self, db):
        priority = get_user_priority("new-user", db)
        assert priority == 2.0  # High priority for unused users


# --- Rate Limiter ---

class TestRateLimiter:
    def test_allows_within_limit(self):
        limiter = RateLimiter()
        limiter.max_requests = 5

        for i in range(5):
            allowed, remaining = limiter.check("user-a")
            assert allowed is True
            assert remaining == 5 - i - 1

    def test_rejects_over_limit(self):
        limiter = RateLimiter()
        limiter.max_requests = 2

        limiter.check("user-a")
        limiter.check("user-a")
        allowed, remaining = limiter.check("user-a")
        assert allowed is False
        assert remaining == 0

    def test_independent_per_user(self):
        limiter = RateLimiter()
        limiter.max_requests = 1

        limiter.check("user-a")
        # user-a exhausted, but user-b should still work
        allowed, _ = limiter.check("user-b")
        assert allowed is True

    def test_reset(self):
        limiter = RateLimiter()
        limiter.max_requests = 1

        limiter.check("user-a")
        limiter.reset("user-a")
        allowed, _ = limiter.check("user-a")
        assert allowed is True


# --- Admin API ---

class TestAdminApi:
    def test_get_quota(self, client, db):
        resp = client.get("/v1/config/quotas/test-user")
        assert resp.status_code == 200
        data = resp.json()
        assert data["max_concurrent_gpu_jobs"] == 2

    def test_update_quota(self, client, db):
        resp = client.put("/v1/config/quotas/test-user", json={
            "max_concurrent_gpu_jobs": 10,
            "max_gpus_total": 32,
        })
        assert resp.status_code == 200
        data = resp.json()
        assert data["max_concurrent_gpu_jobs"] == 10
        assert data["max_gpus_total"] == 32

    def test_get_brain_params(self, client, db):
        resp = client.get("/v1/config/brain")
        assert resp.status_code == 200
        data = resp.json()
        assert data["tau"] == 0.8
        # w_queue_depth may have been modified by FEP during lifespan startup
        assert 0.0 < data["w_queue_depth"] <= 1.0

    def test_list_users(self, client, db):
        # Add a quota
        client.put("/v1/config/quotas/user-a", json={"max_concurrent_gpu_jobs": 5})

        resp = client.get("/v1/config/users")
        assert resp.status_code == 200
        users = resp.json()["users"]
        assert len(users) >= 1
        assert any(u["user_id"] == "user-a" for u in users)

    def test_get_usage(self, client, db):
        resp = client.get("/v1/config/usage/test-user")
        assert resp.status_code == 200
        data = resp.json()
        assert data["gpu_seconds"] == 0.0
