# SPDX-FileCopyrightText: 2024 UL Research Institutes
# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

# mypy: disable-error-code="import-untyped"
import base64
import json
from datetime import datetime, timedelta, timezone

import httpx
import pytest
from conftest import DATA_DIR

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


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

    account_name = ctx["account_non_exempt"]
    account = dyffapi.accounts.create(AccountCreateRequest(name=account_name))
    print(f"account: {account.id} {account.name}")
    ctx["account_non_exempt_object"] = account


@pytest.mark.datafiles(DATA_DIR)
@pytest.mark.depends(
    on=[
        "test_accounts_create",
    ]
)
def test_accounts_get(
    pytestconfig,
    dyffapi: Client | DyffLocalPlatform,
    ctx,
    datafiles,
):
    if not pytestconfig.getoption("test_remote"):
        pytest.skip()
    if pytestconfig.getoption("skip_accounts"):
        pytest.skip()
    assert isinstance(dyffapi, Client)

    account: Account = ctx["account_non_exempt_object"]
    assert account.id is not None
    stored_account = dyffapi.accounts.get(account.id)
    assert stored_account.id == account.id
    assert stored_account.name == account.name
    assert stored_account.model_dump()["auth"] == account.model_dump()["auth"]


@pytest.mark.datafiles(DATA_DIR)
@pytest.mark.depends(
    on=[
        "test_accounts_create",
    ]
)
def test_accounts_set_permissions(
    pytestconfig,
    dyffapi: Client | DyffLocalPlatform,
    ctx,
    datafiles,
):
    if not pytestconfig.getoption("test_remote"):
        pytest.skip()
    if pytestconfig.getoption("skip_accounts"):
        pytest.skip()
    assert isinstance(dyffapi, Client)

    account: Account = ctx["account_non_exempt_object"]
    assert account.id is not None
    dyffapi.accounts.set_permissions(
        account.id,
        AccountPermissionsSetRequest(
            permissions=[
                AccessGrant(
                    resources=[Resources.Dataset.value, Resources.Evaluation.value],
                    functions=[APIFunctions.get],
                    accounts=[ctx["account"]],
                )
            ]
        ),
    )


@pytest.mark.datafiles(DATA_DIR)
@pytest.mark.depends(
    on=[
        "test_accounts_set_permissions",
    ]
)
def test_accounts_get_permissions(
    pytestconfig,
    dyffapi: Client | DyffLocalPlatform,
    ctx,
    datafiles,
):
    if not pytestconfig.getoption("test_remote"):
        pytest.skip()
    if pytestconfig.getoption("skip_accounts"):
        pytest.skip()
    assert isinstance(dyffapi, Client)

    account: Account = ctx["account_non_exempt_object"]
    assert account.id is not None
    stored_permissions = dyffapi.accounts.get_permissions(account.id)
    assert [p.model_dump(mode="json") for p in stored_permissions] == [
        AccessGrant(
            resources=[Resources.Dataset.value, Resources.Evaluation.value],
            functions=[APIFunctions.get],
            accounts=[ctx["account"]],
        ).model_dump(mode="json")
    ]


def _key_from_token(token: str) -> dict:
    _, token_body, _ = token.split(".")
    body_bytes = token_body.encode("utf-8")
    padding = 4 - (len(body_bytes) % 4)
    if padding == 4:
        padding = 0
    padded = body_bytes + b"=" * padding
    return json.loads(base64.urlsafe_b64decode(padded))


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

    account: Account = ctx["account_non_exempt_object"]
    assert account.id is not None
    token = dyffapi.accounts.create_token(
        account.id,
        AccountTokenCreateRequest(
            grants=[
                AccessGrant(
                    resources=[Resources.Dataset.value],
                    functions=[APIFunctions.get],
                    accounts=[ctx["account"]],
                )
            ],
            expires=(datetime.now(tz=timezone.utc) + timedelta(days=7)),
        ),
    )

    api_key_json = _key_from_token(token)
    assert api_key_json["grants"] == [
        {
            "resources": ["datasets"],
            "functions": ["get"],
            "accounts": [ctx["account"]],
            "entities": [],
        }
    ]

    # See if the token works
    dataset: Dataset = ctx["dataset"]
    new_client = Client(
        api_token=token,
        endpoint=dyffapi.endpoint,
        registry_url=dyffapi.registry_url,
        insecure=dyffapi.insecure,
    )
    new_client.datasets.get(dataset.id)


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

    account: Account = ctx["account_non_exempt_object"]
    assert account.id is not None
    tokens = [
        dyffapi.accounts.create_token(
            account.id,
            AccountTokenCreateRequest(
                grants=[
                    AccessGrant(
                        resources=[Resources.Dataset.value],
                        functions=[APIFunctions.get],
                        accounts=[ctx["account"]],
                    )
                ],
                expires=(datetime.now(tz=timezone.utc) + timedelta(days=7)),
            ),
        )
        for _ in range(2)
    ]
    print(f"tokens: {[_key_from_token(t)['jti'] for t in tokens]}")

    dyffapi.accounts.revoke_token(account.id, _key_from_token(tokens[0])["jti"])

    dataset: Dataset = ctx["dataset"]
    # Revoked token should not work
    with pytest.raises(HTTPError) as exc_info:
        new_client = Client(
            api_token=tokens[0],
            endpoint=dyffapi.endpoint,
            registry_url=dyffapi.registry_url,
            insecure=dyffapi.insecure,
        )
        new_client.datasets.get(dataset.id)
    assert exc_info.value.status_code == 401
    # Non-revoked token should work
    new_client = Client(
        api_token=tokens[1],
        endpoint=dyffapi.endpoint,
        registry_url=dyffapi.registry_url,
        insecure=dyffapi.insecure,
    )
    new_client.datasets.get(dataset.id)


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

    account: Account = ctx["account_non_exempt_object"]
    assert account.id is not None
    tokens = [
        dyffapi.accounts.create_token(
            account.id,
            AccountTokenCreateRequest(
                grants=[
                    AccessGrant(
                        resources=[Resources.Dataset.value],
                        functions=[APIFunctions.get],
                        accounts=[ctx["account"]],
                    )
                ],
                expires=(datetime.now(tz=timezone.utc) + timedelta(days=7)),
            ),
        )
        for _ in range(2)
    ]
    print(f"tokens: {[_key_from_token(t)['jti'] for t in tokens]}")

    dyffapi.accounts.revoke_all_tokens(account.id)

    dataset: Dataset = ctx["dataset"]
    # Neither token should work
    for token in tokens:
        with pytest.raises(HTTPError) as exc_info:
            new_client = Client(
                api_token=token,
                endpoint=dyffapi.endpoint,
                registry_url=dyffapi.registry_url,
                insecure=dyffapi.insecure,
            )
            new_client.datasets.get(dataset.id)
        assert exc_info.value.status_code == 401


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

    account: Account = ctx["account_non_exempt_object"]
    assert account.id is not None

    # No adding resources
    with pytest.raises(httpx.HTTPStatusError) as exc_info:
        dyffapi.accounts.create_token(
            account.id,
            AccountTokenCreateRequest(
                grants=[
                    AccessGrant(
                        resources=[Resources.Dataset.value, Resources.SafetyCase.value],
                        functions=[APIFunctions.get],
                        accounts=[ctx["account"]],
                    )
                ],
                expires=(datetime.now(tz=timezone.utc) + timedelta(days=7)),
            ),
        )
    assert exc_info.value.response.status_code == 403

    # No adding functions
    with pytest.raises(httpx.HTTPStatusError) as exc_info:
        dyffapi.accounts.create_token(
            account.id,
            AccountTokenCreateRequest(
                grants=[
                    AccessGrant(
                        resources=[Resources.Dataset.value],
                        functions=[APIFunctions.get, APIFunctions.create],
                        accounts=[ctx["account"]],
                    )
                ],
                expires=(datetime.now(tz=timezone.utc) + timedelta(days=7)),
            ),
        )
    assert exc_info.value.response.status_code == 403

    # No adding accounts
    with pytest.raises(httpx.HTTPStatusError) as exc_info:
        dyffapi.accounts.create_token(
            account.id,
            AccountTokenCreateRequest(
                grants=[
                    AccessGrant(
                        resources=[Resources.Dataset.value],
                        functions=[APIFunctions.get],
                        accounts=[ctx["account"], ids.generate_entity_id()],
                    )
                ],
                expires=(datetime.now(tz=timezone.utc) + timedelta(days=7)),
            ),
        )
    assert exc_info.value.response.status_code == 403

    # No adding entities
    with pytest.raises(httpx.HTTPStatusError) as exc_info:
        dyffapi.accounts.create_token(
            account.id,
            AccountTokenCreateRequest(
                grants=[
                    AccessGrant(
                        resources=[Resources.Dataset.value],
                        functions=[APIFunctions.get],
                        accounts=[ctx["account"]],
                        entities=[ids.generate_entity_id()],
                    )
                ],
                expires=(datetime.now(tz=timezone.utc) + timedelta(days=7)),
            ),
        )
    assert exc_info.value.response.status_code == 403

    # No using '*'
    with pytest.raises(httpx.HTTPStatusError) as exc_info:
        dyffapi.accounts.create_token(
            account.id,
            AccountTokenCreateRequest(
                grants=[
                    AccessGrant(
                        resources=["*"],
                        functions=[APIFunctions.get],
                        accounts=[ctx["account"]],
                    )
                ],
                expires=(datetime.now(tz=timezone.utc) + timedelta(days=7)),
            ),
        )
    assert exc_info.value.response.status_code == 403
