# SPDX-FileCopyrightText: 2024 UL Research Institutes
# SPDX-License-Identifier: Apache-2.0
"""Integration tests for quota enforcement."""

from __future__ import annotations

import json
import time
from datetime import datetime, timedelta, timezone

import pytest
from conftest import DATA_DIR

from dyff.audit.local.platform import DyffLocalPlatform
from dyff.client import Client
from dyff.schema import ids
from dyff.schema.platform import *
from dyff.schema.requests import *


@pytest.mark.datafiles(DATA_DIR)
def test_quotas_create(
    pytestconfig, dyffapi: Client | DyffLocalPlatform, ctx, datafiles
):
    if pytestconfig.getoption("skip_quotas"):
        pytest.skip()
    if not pytestconfig.getoption("test_remote"):
        pytest.skip()
    assert isinstance(dyffapi, Client)

    request = QuotaCreateRequest(
        principal="__global__",
        resource="workflows.dyff.io/admitted",
        flavor="Total",
        limit="16",
    )
    quota = dyffapi.quotas.create(request)

    print(f"quota: {quota.id}")
    ctx["quota"] = quota

    time.sleep(10)

    stored_quota = dyffapi.quotas.get(quota.id)
    assert stored_quota.principal == request.principal
    assert stored_quota.resource == request.resource
    assert stored_quota.flavor == request.flavor
    assert stored_quota.limit == request.limit


@pytest.mark.datafiles(DATA_DIR)
@pytest.mark.depends(
    on=[
        "test_quotas_create",
        "tests/test_datasets.py::test_datasets_create",
        "tests/test_inference.py::test_inferenceservices_create_mock",
    ]
)
def test_quotas_global(
    pytestconfig, dyffapi: Client | DyffLocalPlatform, ctx, datafiles
):
    if pytestconfig.getoption("skip_quotas"):
        pytest.skip()
    if not pytestconfig.getoption("test_remote"):
        pytest.skip()
    assert isinstance(dyffapi, Client)

    # We will test that both global and account quotas are enforced by
    # creating 3 accounts with a quota of 1 workflow, and setting the global
    # quota to 2 workflows.
    num_accounts = 3
    num_evaluations_per_account = 2
    global_limit = 2
    account_limit = 1
    max_blocked = num_accounts * num_evaluations_per_account - min(
        global_limit, num_accounts * account_limit
    )

    dataset: Dataset = ctx["dataset"]
    inferenceservice: InferenceService = ctx["inferenceservice_mock"]
    accounts = [f"test-{ids.generate_entity_id()}" for _ in range(num_accounts)]

    quotas = [
        dyffapi.quotas.create(
            QuotaCreateRequest(
                principal="__global__",
                resource="workflows.dyff.io/admitted",
                flavor="Total",
                limit=str(global_limit),
            )
        )
    ]
    for account in accounts:
        quotas.append(
            dyffapi.quotas.create(
                QuotaCreateRequest(
                    principal=f"accounts/{account}",
                    resource="workflows.dyff.io/admitted",
                    flavor="Total",
                    limit=str(account_limit),
                )
            )
        )
    time.sleep(10)
    for q in quotas:
        assert dyffapi.quotas.get(q.id).limit == q.limit

    # Interleave starting evaluations from different accounts so that we
    # push the maximum global parallelism
    evaluations = {}
    for _ in range(num_evaluations_per_account):
        for account in accounts:
            evaluations.setdefault(account, []).append(
                dyffapi.evaluations.create(
                    EvaluationCreateRequest(
                        account=account,
                        dataset=dataset.id,
                        inferenceSession=EvaluationInferenceSessionRequest(
                            inferenceService=inferenceservice.id,
                            expires=datetime.now(timezone.utc) + timedelta(days=1),
                            replicas=1,
                            useSpotPods=False,
                        ),
                        replications=2,
                        workersPerReplica=2,
                    )
                )
            )
    evaluation_ids = {a: [e.id for e in es] for a, es in evaluations.items()}
    print(f"evaluations: {json.dumps(evaluation_ids, indent=2)}")

    observed_blocked = False
    observed_max_utilization = False
    then = datetime.now(timezone.utc)
    timeout = timedelta(minutes=10)
    while True:
        time.sleep(5)
        status: dict[str, list[Evaluation]] = {
            account: [dyffapi.evaluations.get(e.id) for e in es]
            for account, es in evaluations.items()
        }

        total_admitted = set()
        total_blocked = set()
        total_success = set()
        total_failure = set()
        for account, es in status.items():
            blocked = set(
                e.id for e in es if (e.status == "Created" and e.reason == "QuotaLimit")
            )
            admitted = set(e.id for e in es if e.status == "Admitted")
            success = set(e.id for e in es if is_status_success(e.status))
            failure = set(e.id for e in es if is_status_failure(e.status))

            assert (
                len(admitted) <= account_limit
            ), f"account: {account}; admitted: {len(admitted)} -- {{ {admitted} }}"

            total_admitted |= admitted
            total_blocked |= blocked
            total_success |= success
            total_failure |= failure

        if len(total_blocked) > 0:
            observed_blocked = True
        if len(total_admitted) == global_limit:
            observed_max_utilization = True
        assert (
            len(total_admitted) <= global_limit
        ), f"total_admitted: {len(total_admitted)} -- {{ {total_admitted} }}"
        assert (
            len(total_blocked) <= max_blocked
        ), f"total_blocked: {len(total_blocked)} -- {{ {total_blocked} }}"
        assert (
            len(total_failure) == 0
        ), f"total_failure: {len(total_failure)} -- {{ {total_failure} }}"
        if len(total_success) == (num_accounts * num_evaluations_per_account):
            break

        now = datetime.now(timezone.utc)
        if now - then > timeout:
            raise AssertionError("timeout")
        then = now

    assert observed_blocked
    assert observed_max_utilization


@pytest.mark.datafiles(DATA_DIR)
@pytest.mark.depends(
    on=[
        "test_quotas_create",
        "test_quotas_global",
        "tests/test_datasets.py::test_datasets_create",
        "tests/test_inference.py::test_inferenceservices_create_mock",
    ]
)
def test_quotas_increase_unblocks_waiting(
    pytestconfig, dyffapi: Client | DyffLocalPlatform, ctx, datafiles
):
    if pytestconfig.getoption("skip_quotas"):
        pytest.skip()
    if not pytestconfig.getoption("test_remote"):
        pytest.skip()
    assert isinstance(dyffapi, Client)

    # We will create more workflows than our quota allows, then increase the
    # quota and see if more workflows get admitted
    num_evaluations = 2
    big_limit = 2
    small_limit = 1

    dataset: Dataset = ctx["dataset"]
    inferenceservice: InferenceService = ctx["inferenceservice_mock"]
    account = f"test-{ids.generate_entity_id()}"

    quotas = [
        dyffapi.quotas.create(
            QuotaCreateRequest(
                principal="__global__",
                resource="workflows.dyff.io/admitted",
                flavor="Total",
                limit=str(big_limit),
            )
        ),
        dyffapi.quotas.create(
            QuotaCreateRequest(
                principal=f"accounts/{account}",
                resource="workflows.dyff.io/admitted",
                flavor="Total",
                limit=str(small_limit),
            )
        ),
    ]
    time.sleep(10)
    for q in quotas:
        assert dyffapi.quotas.get(q.id).limit == q.limit

    evaluations = [
        dyffapi.evaluations.create(
            EvaluationCreateRequest(
                account=account,
                dataset=dataset.id,
                inferenceSession=EvaluationInferenceSessionRequest(
                    inferenceService=inferenceservice.id,
                    expires=datetime.now(timezone.utc) + timedelta(days=1),
                    replicas=1,
                    useSpotPods=False,
                ),
                replications=2,
                workersPerReplica=2,
            )
        )
        for _ in range(num_evaluations)
    ]

    # Wait until 1 evaluation is admitted and 1 is blocked
    then = datetime.now(timezone.utc)
    timeout = timedelta(minutes=5)
    while True:
        time.sleep(5)
        status: list[Evaluation] = [dyffapi.evaluations.get(e.id) for e in evaluations]

        blocked = set(
            e.id for e in status if (e.status == "Created" and e.reason == "QuotaLimit")
        )
        admitted = set(e.id for e in status if e.status == "Admitted")
        success = set(e.id for e in status if is_status_success(e.status))
        failure = set(e.id for e in status if is_status_failure(e.status))

        assert (
            len(admitted) <= small_limit
        ), f"account: {account}; admitted: {len(admitted)} -- {{ {admitted} }}"
        assert len(failure) == 0, f"failure: {len(failure)} -- {{ {failure} }}"

        if len(admitted) == 1 and len(blocked) == 1:
            break

        now = datetime.now(timezone.utc)
        if now - then > timeout:
            raise AssertionError("timeout")
        then = now

    # Increase the quota
    dyffapi.quotas.create(
        QuotaCreateRequest(
            principal=f"accounts/{account}",
            resource="workflows.dyff.io/admitted",
            flavor="Total",
            limit=str(big_limit),
        )
    )

    # Confirm that both workflows get unblocked
    observed_unblocked = False
    then = datetime.now(timezone.utc)
    timeout = timedelta(minutes=5)
    while True:
        time.sleep(5)
        status: list[Evaluation] = [dyffapi.evaluations.get(e.id) for e in evaluations]

        blocked = set(
            e.id for e in status if (e.status == "Created" and e.reason == "QuotaLimit")
        )
        admitted = set(e.id for e in status if e.status == "Admitted")
        success = set(e.id for e in status if is_status_success(e.status))
        failure = set(e.id for e in status if is_status_failure(e.status))

        assert (
            len(admitted) <= big_limit
        ), f"account: {account}; admitted: {len(admitted)} -- {{ {admitted} }}"
        assert len(blocked) <= (
            num_evaluations - small_limit
        ), f"account: {account}; blocked: {len(blocked)} -- {{ {blocked} }}"
        assert len(failure) == 0, f"failure: {len(failure)} -- {{ {failure} }}"

        if len(blocked) == 0:
            observed_unblocked = True

        if len(success) == num_evaluations:
            break

        now = datetime.now(timezone.utc)
        if now - then > timeout:
            raise AssertionError("timeout")
        then = now

    assert observed_unblocked


@pytest.mark.datafiles(DATA_DIR)
@pytest.mark.depends(
    on=[
        "test_quotas_create",
        "test_quotas_global",
        "tests/test_datasets.py::test_datasets_create",
        "tests/test_inference.py::test_inferenceservices_create_mock",
    ]
)
def test_quotas_account_rate_quota(
    pytestconfig, dyffapi: Client | DyffLocalPlatform, ctx, datafiles
):
    if pytestconfig.getoption("skip_quotas"):
        pytest.skip()
    if not pytestconfig.getoption("test_remote"):
        pytest.skip()
    assert isinstance(dyffapi, Client)

    dataset: Dataset = ctx["dataset_tiny"]
    inferenceservice: InferenceService = ctx["inferenceservice_mock"]
    account = f"test-{ids.generate_entity_id()}"
    num_evaluations = 3
    rate_limit = 1

    quota = dyffapi.quotas.create(
        QuotaCreateRequest(
            principal=f"accounts/{account}",
            resource="evaluations.dyff.io/created",
            flavor="Rate",
            limit=str(rate_limit),
            duration=timedelta(minutes=10),
        )
    )

    time.sleep(10)

    evaluations = [
        dyffapi.evaluations.create(
            EvaluationCreateRequest(
                account=account,
                dataset=dataset.id,
                inferenceSession=EvaluationInferenceSessionRequest(
                    inferenceService=inferenceservice.id,
                    expires=datetime.now(timezone.utc) + timedelta(days=1),
                    replicas=1,
                    useSpotPods=False,
                ),
                replications=2,
                workersPerReplica=2,
            )
        )
        for _ in range(num_evaluations)
    ]

    then = datetime.now(timezone.utc)
    timeout = timedelta(minutes=5)
    while True:
        time.sleep(5)
        num_rejected = 0
        num_accepted = 0
        for e in evaluations:
            current = dyffapi.evaluations.get(e.id)
            if current.reason == "RateLimitExceeded":
                num_rejected += 1
            elif current.status in ["Admitted", "Complete"]:
                num_accepted += 1
        if num_accepted == rate_limit and num_rejected == (
            num_evaluations - rate_limit
        ):
            break
        assert num_rejected <= (num_evaluations - rate_limit)
        assert num_accepted <= rate_limit

        now = datetime.now(timezone.utc)
        if now - then > timeout:
            raise AssertionError("timeout")
        then = now
