# SPDX-FileCopyrightText: 2024 UL Research Institutes
# SPDX-License-Identifier: Apache-2.0
"""Integration tests for S3StorageBackend against real AWS S3.

Requires DYFF_TEST_S3_ENDPOINT, DYFF_TEST_S3_BUCKET, DYFF_TEST_S3_ACCESS_KEY, and
DYFF_TEST_S3_SECRET_KEY to be set.
"""

import base64
import hashlib
import os
import types
import uuid

import httpx
import pytest

import dyff.storage.backend.s3.storage as _s3_module
from dyff.schema.platform import Artifact, Digest
from dyff.storage.backend.s3.storage import S3StorageBackend
from dyff.storage.config import S3StorageConfig

_REQUIRED_ENV = (
    "DYFF_TEST_S3_ENDPOINT",
    "DYFF_TEST_S3_BUCKET",
    "DYFF_TEST_S3_ACCESS_KEY",
    "DYFF_TEST_S3_SECRET_KEY",
)

pytestmark = pytest.mark.skipif(
    not all(os.environ.get(v) for v in _REQUIRED_ENV),
    reason=f"set {', '.join(_REQUIRED_ENV)} to run AWS S3 integration tests",
)


def _md5_b64(data: bytes) -> str:
    return base64.b64encode(hashlib.md5(data).digest()).decode()


@pytest.fixture()
def backend(monkeypatch):
    cfg = types.SimpleNamespace(
        storage=types.SimpleNamespace(
            s3=S3StorageConfig(
                endpoint=os.environ["DYFF_TEST_S3_ENDPOINT"],
                access_key=os.environ["DYFF_TEST_S3_ACCESS_KEY"],
                secret_key=os.environ["DYFF_TEST_S3_SECRET_KEY"],
            )
        )
    )
    monkeypatch.setattr(_s3_module, "config", cfg)
    b = S3StorageBackend()
    yield b


@pytest.fixture()
def storage_path(backend):
    bucket = os.environ["DYFF_TEST_S3_BUCKET"]
    prefix = f"dyff-test-{uuid.uuid4().hex}"
    path = f"s3://{bucket}/{prefix}"
    yield path
    try:
        for obj_path in backend.list_dir(f"{path}/", recursive=True):
            obj_bucket, key = obj_path.removeprefix("s3://").split("/", 1)
            _s3_module._get_client().remove_object(obj_bucket, key)
    except Exception:
        pass


def test_upload_succeeds(backend, storage_path):
    content = b"hello s3 " + uuid.uuid4().bytes
    artifact = Artifact(path="upload.bin", digest=Digest(md5=_md5_b64(content)))
    put_url = backend.signed_url_for_artifact_upload(artifact, storage_path)

    headers = {"Content-Length": str(len(content))}
    headers.update(put_url.headers)
    response = httpx.put(put_url.url, content=content, headers=headers)
    assert response.status_code == 200


def test_download_after_upload(backend, storage_path):
    content = b"roundtrip " + uuid.uuid4().bytes
    artifact = Artifact(path="rt.bin", digest=Digest(md5=_md5_b64(content)))
    put_url = backend.signed_url_for_artifact_upload(artifact, storage_path)

    headers = {"Content-Length": str(len(content))}
    headers.update(put_url.headers)
    httpx.put(put_url.url, content=content, headers=headers).raise_for_status()

    _, root_path = storage_path.removeprefix("s3://").split("/", 1)
    art_url = backend._artifact_url_for_downlink(
        root_path=root_path,
        artifact_path=f"{storage_path}/{artifact.path}",
    )
    response = httpx.get(art_url.signedURL.url)
    assert response.status_code == 200
    assert response.content == content


def test_upload_with_content_md5_fails(backend, storage_path):
    """AWS presigned PUTs reject content-md5 because it is not in X-Amz-SignedHeaders.

    If content-md5 is ever re-added to upload requests it will break against AWS S3.
    """
    content = b"negative test " + uuid.uuid4().bytes
    artifact = Artifact(path="neg.bin", digest=Digest(md5=_md5_b64(content)))
    put_url = backend.signed_url_for_artifact_upload(artifact, storage_path)

    headers = {
        "Content-Length": str(len(content)),
        "content-md5": _md5_b64(content),
    }
    headers.update(put_url.headers)
    response = httpx.put(put_url.url, content=content, headers=headers)
    assert response.status_code == 403
