import json

import pytest
from unittest.mock import patch
from pyinfra.storage.utils import (
    download_data_bytes_as_specified_in_message,
    upload_data_as_specified_in_message,
    DownloadedData,
)
from pyinfra.storage.storages.storage import Storage


@pytest.fixture
def mock_storage():
    with patch("pyinfra.storage.utils.Storage") as MockStorage:
        yield MockStorage()


@pytest.fixture(
    params=[
        {
            "raw_payload": {
                "tenantId": "tenant1",
                "dossierId": "dossier1",
                "fileId": "file1",
                "targetFileExtension": "txt",
                "responseFileExtension": "json",
            },
            "expected_result": {
                "data": b'{"key": "value"}',
                "file_path": "tenant1/dossier1/file1.txt"
            }
        },
        {
            "raw_payload": {
                "targetFilePath": "some/path/to/file.txt.gz",
                "responseFilePath": "some/path/to/file.json"
            },
            "expected_result": {
                "data": b'{"key": "value"}',
                "file_path": "some/path/to/file.txt.gz"
            }
        },
        {
            "raw_payload": {
                "targetFilePath": {
                    "file1": "some/path/to/file1.txt.gz",
                    "file2": "some/path/to/file2.txt.gz"
                },
                "responseFilePath": "some/path/to/file.json"
            },
            "expected_result": {
                "file1": {
                    "data": b'{"key": "value"}',
                    "file_path": "some/path/to/file1.txt.gz"
                },
                "file2": {
                    "data": b'{"key": "value"}',
                    "file_path": "some/path/to/file2.txt.gz"
                }
            }
        },
    ]
)
def payload_and_expected_result(request):
    return request.param

def test_download_data_bytes_as_specified_in_message(mock_storage, payload_and_expected_result):
    raw_payload = payload_and_expected_result["raw_payload"]
    expected_result = payload_and_expected_result["expected_result"]
    mock_storage.get_object.return_value = b'{"key": "value"}'

    result = download_data_bytes_as_specified_in_message(mock_storage, raw_payload)

    assert isinstance(result, dict)
    assert result == expected_result
    mock_storage.get_object.assert_called()

def test_upload_data_as_specified_in_message(mock_storage, payload_and_expected_result):
    raw_payload = payload_and_expected_result["raw_payload"]
    data = {"key": "value"}
    upload_data_as_specified_in_message(mock_storage, raw_payload, data)
    mock_storage.put_object.assert_called_once()