# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
#
# Copyright (c) 2026 Dmitrii Gagarin aka madgagarin

import os
import tempfile
from unittest.mock import AsyncMock, MagicMock, patch

import pytest
from obstore.store import S3Store
from rxon.exceptions import IntegrityError, RxonError
from rxon.models import FileMetadata

from avtomatika_worker.config import WorkerConfig
from avtomatika_worker.s3 import S3Manager


@pytest.fixture
def s3_manager():
    """Provides an S3Manager instance for testing."""
    config = WorkerConfig()
    config.S3_ENDPOINT_URL = "http://localhost:5000"
    config.S3_ACCESS_KEY = "testing"
    config.S3_SECRET_KEY = "testing"
    config.S3_DEFAULT_BUCKET = "test-bucket"
    config.TASK_FILES_DIR = tempfile.gettempdir()

    manager = S3Manager(config)
    return manager


# Helper classes for mocking obstore


class MockGetResult:
    def __init__(self, data=b"test content"):
        self.data = data
        self.meta = MagicMock(size=len(data), e_tag='"hash"')

    async def stream(self):
        yield self.data


class MockObjectMeta:
    def __init__(self, key):
        self.key = key


class MockAsyncIterator:
    def __init__(self, items):
        self.items = items

    def __aiter__(self):
        return self

    async def __anext__(self):
        if not self.items:
            raise StopAsyncIteration
        return self.items.pop(0)


@pytest.mark.asyncio
async def test_get_store_caching(s3_manager):
    """Tests that _get_store caches S3Store instances."""
    # We want S3Store() to return a new mock each time it's called as a constructor
    with patch(
        "avtomatika_worker.s3.S3Store", side_effect=lambda *args, **kwargs: MagicMock(spec=S3Store)
    ) as mock_s3_store:
        store1 = s3_manager._get_store("bucket1")
        store2 = s3_manager._get_store("bucket1")
        store3 = s3_manager._get_store("bucket2")

        assert store1 is store2
        assert store1 is not store3

        assert mock_s3_store.call_count == 2


@pytest.mark.asyncio
async def test_process_s3_uri_file(s3_manager):
    """Tests that _process_s3_uri downloads a single file from S3."""
    task_id = "task-123"

    with patch("avtomatika_worker.s3.obstore_get", new_callable=AsyncMock) as mock_get:
        mock_get.return_value = MockGetResult()

        with patch.object(s3_manager, "_get_store", return_value=MagicMock()) as mock_get_store:
            local_path = await s3_manager._process_s3_uri("s3://test-bucket/test-file.txt", task_id)

            # Path should include task_id
            expected_suffix = os.path.join(task_id, "test-file.txt")
            assert local_path.endswith(expected_suffix)

            mock_get_store.assert_called_with("test-bucket")
            mock_get.assert_awaited_once()

            # Check content
            with open(local_path, "rb") as f:
                assert f.read() == b"test content"


@pytest.mark.asyncio
async def test_process_s3_uri_folder(s3_manager):
    """Tests that _process_s3_uri downloads a folder (prefix) from S3."""
    task_id = "task-folder"

    mock_files = [MockObjectMeta("data/file1.txt"), MockObjectMeta("data/sub/file2.txt")]

    # obstore.list is called as: async for obj in obstore.list(...)
    # So obstore.list itself should return the async iterator.
    with (
        patch("avtomatika_worker.s3.obstore_list", return_value=MockAsyncIterator(mock_files)) as mock_list,
        patch("avtomatika_worker.s3.obstore_get", new_callable=AsyncMock) as mock_get,
    ):
        mock_get.return_value = MockGetResult()

        with patch.object(s3_manager, "_get_store", return_value=MagicMock()):
            local_path = await s3_manager._process_s3_uri("s3://test-bucket/data/", task_id)

    assert local_path.endswith(os.path.join(task_id, "data"))

    # Verify list called with correct prefix
    mock_list.assert_called_once()
    assert mock_list.call_args[1]["prefix"] == "data/"

    # Verify download calls
    assert mock_get.await_count == 2


@pytest.mark.asyncio
async def test_upload_to_s3_file(s3_manager):
    """Tests that _upload_to_s3 uploads a single file to S3."""
    local_path = os.path.join(s3_manager._config.TASK_FILES_DIR, "test-upload.txt")
    with open(local_path, "w") as f:
        f.write("test content")

    try:
        with (
            patch("avtomatika_worker.s3.obstore_put", new_callable=AsyncMock) as mock_put,
            patch.object(s3_manager, "_get_store", return_value=MagicMock()),
        ):
            mock_put.return_value = {"e_tag": '"etag"'}
            meta = await s3_manager._upload_to_s3(local_path)

        assert meta.uri == "s3://test-bucket/test-upload.txt"
        assert meta.etag == "etag"
        mock_put.assert_awaited_once()
        args = mock_put.await_args[0]
        assert args[1] == "test-upload.txt"  # key
        # Check that a file object was passed
        assert hasattr(args[2], "read")
    finally:
        if os.path.exists(local_path):
            os.remove(local_path)


@pytest.mark.asyncio
async def test_upload_to_s3_folder(s3_manager):
    """Tests that _upload_to_s3 uploads a folder recursively to S3."""
    # Create temp directory structure
    with tempfile.TemporaryDirectory(dir=s3_manager._config.TASK_FILES_DIR) as tmpdir:
        # Create file1
        with open(os.path.join(tmpdir, "file1.txt"), "w") as f:
            f.write("content1")

        dirname = os.path.basename(tmpdir)

        with (
            patch("avtomatika_worker.s3.obstore_put", new_callable=AsyncMock) as mock_put,
            patch.object(s3_manager, "_get_store", return_value=MagicMock()),
        ):
            meta = await s3_manager._upload_to_s3(tmpdir)

        assert meta.uri == f"s3://test-bucket/{dirname}/"
        assert mock_put.await_count == 1


@pytest.mark.asyncio
async def test_process_params(s3_manager):
    """Tests that process_params correctly calls _process_s3_uri with task_id."""
    with patch.object(
        s3_manager,
        "_process_s3_uri",
        side_effect=lambda uri, tid, verify_meta=None: f"/local/{tid}/{uri.split('/')[-1]}",
    ):
        params = {"file": "s3://test-bucket/test-file.txt", "other": "value"}
        processed_params = await s3_manager.process_params(params, "task-1")

        assert processed_params["other"] == "value"
        assert processed_params["file"] == "/local/task-1/test-file.txt"


@pytest.mark.asyncio
async def test_process_result(s3_manager):
    """Tests that process_result correctly calls _upload_to_s3."""
    local_path = os.path.join(s3_manager._config.TASK_FILES_DIR, "output.txt")
    with open(local_path, "w") as f:
        f.write("test content")

    try:
        mock_meta = FileMetadata(uri="s3://bucket/output.txt", size=10, etag="hash")
        with patch.object(s3_manager, "_upload_to_s3", return_value=mock_meta):
            result = {"data": {"output_file": local_path}}
            processed_result, metadata = await s3_manager.process_result(result)

            assert processed_result["data"]["output_file"] == "s3://bucket/output.txt"
            assert metadata["data.output_file"] == mock_meta
    finally:
        if os.path.exists(local_path):
            os.remove(local_path)


@pytest.mark.asyncio
async def test_cleanup(s3_manager):
    """Tests that cleanup removes the task directory."""
    task_id = "cleanup-task"
    task_dir = os.path.join(s3_manager._config.TASK_FILES_DIR, task_id)
    os.makedirs(task_dir, exist_ok=True)
    with open(os.path.join(task_dir, "file.txt"), "w") as f:
        f.write("data")

    assert os.path.exists(task_dir)

    await s3_manager.cleanup(task_id)

    assert not os.path.exists(task_dir)


@pytest.mark.asyncio
async def test_process_s3_uri_integrity_size_mismatch(s3_manager):
    """Tests that _process_s3_uri raises IntegrityError on size mismatch."""
    task_id = "integrity-size"
    verify_meta = FileMetadata(uri="s3://bucket/file", size=100)

    with patch("avtomatika_worker.s3.obstore_get", new_callable=AsyncMock) as mock_get:
        # Mock returns 50 bytes instead of 100
        mock_get.return_value = MockGetResult(data=b"x" * 50)

        with patch.object(s3_manager, "_get_store", return_value=MagicMock()):
            with pytest.raises(IntegrityError) as exc:
                await s3_manager._process_s3_uri("s3://bucket/file", task_id, verify_meta=verify_meta)
            assert "Size mismatch" in str(exc.value)


@pytest.mark.asyncio
async def test_process_s3_uri_integrity_etag_mismatch(s3_manager):
    """Tests that _process_s3_uri raises IntegrityError on ETag mismatch."""
    task_id = "integrity-etag"
    verify_meta = FileMetadata(uri="s3://bucket/file", size=12, etag="expected_hash")

    with patch("avtomatika_worker.s3.obstore_get", new_callable=AsyncMock) as mock_get:
        # Mock returns different ETag
        result = MockGetResult()
        result.meta.e_tag = '"actual_hash"'
        mock_get.return_value = result

        with patch.object(s3_manager, "_get_store", return_value=MagicMock()):
            with pytest.raises(IntegrityError) as exc:
                await s3_manager._process_s3_uri("s3://bucket/file", task_id, verify_meta=verify_meta)
            assert "ETag mismatch" in str(exc.value)


@pytest.mark.asyncio
async def test_upload_to_s3_retry_on_transient_error(s3_manager):
    """Tests that _upload_to_s3 retries on transient errors (e.g. 500)."""
    local_path = os.path.join(s3_manager._config.TASK_FILES_DIR, "retry.txt")
    with open(local_path, "w") as f:
        f.write("content")

    try:
        with (
            patch("avtomatika_worker.s3.obstore_put", new_callable=AsyncMock) as mock_put,
            patch.object(s3_manager, "_get_store", return_value=MagicMock()),
            patch("avtomatika_worker.s3.sleep", new_callable=AsyncMock) as mock_sleep,
        ):
            # First call fails, second succeeds
            mock_put.side_effect = [Exception("500 Internal Error"), {"e_tag": '"ok"'}]
            meta = await s3_manager._upload_to_s3(local_path)

            assert meta.etag == "ok"
            assert mock_put.await_count == 2
            mock_sleep.assert_awaited_once()
    finally:
        if os.path.exists(local_path):
            os.remove(local_path)


@pytest.mark.asyncio
async def test_upload_to_s3_permanent_error(s3_manager):
    """Tests that _upload_to_s3 fails fast on permanent errors (e.g. 403 Forbidden)."""
    local_path = os.path.join(s3_manager._config.TASK_FILES_DIR, "fatal.txt")
    with open(local_path, "w") as f:
        f.write("content")

    try:
        with (
            patch("avtomatika_worker.s3.obstore_put", new_callable=AsyncMock) as mock_put,
            patch.object(s3_manager, "_get_store", return_value=MagicMock()),
        ):
            # Error 403 should stop retries
            mock_put.side_effect = Exception("403 Forbidden: Access Denied")
            with pytest.raises(RxonError) as exc:
                await s3_manager._upload_to_s3(local_path)
            assert "Access Denied" in str(exc.value)
            assert mock_put.await_count == 1
    finally:
        if os.path.exists(local_path):
            os.remove(local_path)


@pytest.mark.asyncio
async def test_upload_to_s3_empty_folder(s3_manager):
    """Tests uploading an empty folder."""
    with tempfile.TemporaryDirectory(dir=s3_manager._config.TASK_FILES_DIR) as tmpdir:
        with (
            patch("avtomatika_worker.s3.obstore_put", new_callable=AsyncMock) as mock_put,
            patch.object(s3_manager, "_get_store", return_value=MagicMock()),
        ):
            meta = await s3_manager._upload_to_s3(tmpdir)

        assert meta.size == 0
        assert mock_put.await_count == 0


@pytest.mark.asyncio
async def test_process_params_invalid_uri(s3_manager):
    """Tests that process_params handles invalid URIs by propagating errors."""
    params = {"file": "s3://invalid-uri"}  # missing bucket/key or malformed

    with (
        patch("avtomatika_worker.s3.parse_uri", side_effect=ValueError("Invalid URI")),
        pytest.raises(ValueError, match="Invalid URI"),
    ):
        await s3_manager.process_params(params, "task-1")
