from moto import mock_aws
from ZODB.MappingStorage import MappingStorage
from ZODB.utils import p64
from zodb_s3blobs.cache import S3BlobCache
from zodb_s3blobs.s3client import S3Client
from zodb_s3blobs.storage import S3BlobStorage

import boto3
import os
import pytest
import stat
import transaction
import ZODB.interfaces


@pytest.fixture
def s3_env():
    with mock_aws():
        boto3.client("s3", region_name="us-east-1").create_bucket(Bucket="test-bucket")
        yield


@pytest.fixture
def base_storage():
    return MappingStorage()


@pytest.fixture
def s3_client(s3_env):
    return S3Client(bucket_name="test-bucket", region_name="us-east-1")


@pytest.fixture
def blob_cache(tmp_path):
    return S3BlobCache(str(tmp_path / "cache"), max_size=10 * 1024 * 1024)


@pytest.fixture
def storage(base_storage, s3_client, blob_cache, tmp_path):
    return S3BlobStorage(
        base_storage, s3_client, blob_cache, temp_dir=str(tmp_path / "staging")
    )


def _make_blob_file(tmp_path, content=b"blob content"):
    """Create a temp blob file and return its path."""
    p = tmp_path / f"blob_{id(content)}.bin"
    p.write_bytes(content)
    return str(p)


class TestProxy:
    def test_getattr_delegates_to_base(self, storage, base_storage):
        assert storage.sortKey() == base_storage.sortKey()

    def test_implements_iblobstorage(self, storage):
        assert ZODB.interfaces.IBlobStorage.providedBy(storage)

    def test_len(self, storage):
        assert len(storage) == len(MappingStorage())

    def test_repr(self, storage):
        r = repr(storage)
        assert "S3BlobStorage" in r
        assert "proxy" in r.lower()

    def test_temporary_directory(self, storage, tmp_path):
        td = storage.temporaryDirectory()
        assert os.path.isdir(td)


class TestStoreBlob:
    def test_store_blob_stages_file(self, storage, tmp_path):
        oid = p64(1)
        blob_path = _make_blob_file(tmp_path)
        txn = transaction.get()
        storage.tpc_begin(txn)
        storage.storeBlob(oid, p64(0), b"pickle", blob_path, "", txn)

        # Blob should be staged (original file consumed)
        assert not os.path.exists(blob_path)

    def test_store_blob_calls_base_store(self, storage, base_storage, tmp_path):
        oid = p64(1)
        blob_path = _make_blob_file(tmp_path)
        txn = transaction.get()
        storage.tpc_begin(txn)
        storage.storeBlob(oid, p64(0), b"pickle", blob_path, "", txn)

        # After tpc_vote + tpc_finish, the object data should be in base storage
        storage.tpc_vote(txn)
        storage.tpc_finish(txn)

        data, _tid = base_storage.load(oid)
        assert data == b"pickle"


class TestTwoPhaseCommit:
    def test_tpc_vote_uploads_to_s3(self, storage, s3_client, tmp_path):
        oid = p64(1)
        blob_path = _make_blob_file(tmp_path, b"vote test data")
        txn = transaction.get()
        storage.tpc_begin(txn)
        storage.storeBlob(oid, p64(0), b"pickle", blob_path, "", txn)
        storage.tpc_vote(txn)

        # After vote, blob should be in S3
        keys = list(s3_client.list_objects("blobs/"))
        assert len(keys) == 1
        assert keys[0].endswith(".blob")

        storage.tpc_finish(txn)

    def test_tpc_finish_populates_cache(self, storage, blob_cache, tmp_path):
        oid = p64(1)
        blob_path = _make_blob_file(tmp_path, b"cache after finish")
        txn = transaction.get()
        storage.tpc_begin(txn)
        storage.storeBlob(oid, p64(0), b"pickle", blob_path, "", txn)
        storage.tpc_vote(txn)
        tid = storage.tpc_finish(txn)

        # After finish, blob should be in cache
        cached = blob_cache.get(oid, tid)
        assert cached is not None
        with open(cached, "rb") as f:
            assert f.read() == b"cache after finish"

    def test_tpc_finish_clears_pending(self, storage, tmp_path):
        oid = p64(1)
        blob_path = _make_blob_file(tmp_path)
        txn = transaction.get()
        storage.tpc_begin(txn)
        storage.storeBlob(oid, p64(0), b"pickle", blob_path, "", txn)
        storage.tpc_vote(txn)
        storage.tpc_finish(txn)

        assert storage._pending_blobs == {}
        assert storage._uploaded_keys == []

    def test_tpc_finish_returns_tid(self, storage, tmp_path):
        oid = p64(1)
        blob_path = _make_blob_file(tmp_path)
        txn = transaction.get()
        storage.tpc_begin(txn)
        storage.storeBlob(oid, p64(0), b"pickle", blob_path, "", txn)
        storage.tpc_vote(txn)
        tid = storage.tpc_finish(txn)

        assert tid is not None
        assert isinstance(tid, bytes)
        assert len(tid) == 8

    def test_tpc_abort_deletes_s3_keys(self, storage, s3_client, tmp_path):
        oid = p64(1)
        blob_path = _make_blob_file(tmp_path, b"abort test")
        txn = transaction.get()
        storage.tpc_begin(txn)
        storage.storeBlob(oid, p64(0), b"pickle", blob_path, "", txn)
        storage.tpc_vote(txn)

        # Before abort, S3 should have the blob
        keys_before = list(s3_client.list_objects("blobs/"))
        assert len(keys_before) == 1

        storage.tpc_abort(txn)

        # After abort, S3 should be clean
        keys_after = list(s3_client.list_objects("blobs/"))
        assert len(keys_after) == 0

    def test_tpc_abort_cleans_staged_files(self, storage, tmp_path):
        oid = p64(1)
        blob_path = _make_blob_file(tmp_path)
        txn = transaction.get()
        storage.tpc_begin(txn)
        storage.storeBlob(oid, p64(0), b"pickle", blob_path, "", txn)

        # Abort before vote (no S3 uploads yet)
        storage.tpc_abort(txn)
        assert storage._pending_blobs == {}

    def test_tpc_abort_best_effort_s3(self, storage, s3_client, tmp_path):
        """S3 delete failure during abort should not raise."""
        oid = p64(1)
        blob_path = _make_blob_file(tmp_path)
        txn = transaction.get()
        storage.tpc_begin(txn)
        storage.storeBlob(oid, p64(0), b"pickle", blob_path, "", txn)
        storage.tpc_vote(txn)

        # Monkey-patch delete to fail
        original_delete = s3_client.delete_object
        s3_client.delete_object = lambda key: (_ for _ in ()).throw(
            Exception("S3 down")
        )

        # Abort should not raise despite S3 failure
        storage.tpc_abort(txn)
        assert storage._pending_blobs == {}
        assert storage._uploaded_keys == []

        s3_client.delete_object = original_delete

    def test_multiple_blobs_single_transaction(self, storage, s3_client, tmp_path):
        txn = transaction.get()
        storage.tpc_begin(txn)

        for i in range(3):
            oid = p64(i + 1)
            blob_path = _make_blob_file(tmp_path, f"blob {i}".encode())
            storage.storeBlob(oid, p64(0), b"pickle", blob_path, "", txn)

        storage.tpc_vote(txn)

        keys = list(s3_client.list_objects("blobs/"))
        assert len(keys) == 3

        storage.tpc_finish(txn)


class TestLoadBlob:
    def _store_and_commit(self, storage, oid, blob_content, tmp_path):
        """Helper: store a blob and commit, return tid."""
        blob_path = _make_blob_file(tmp_path, blob_content)
        txn = transaction.get()
        storage.tpc_begin(txn)
        storage.storeBlob(oid, p64(0), b"pickle", blob_path, "", txn)
        storage.tpc_vote(txn)
        return storage.tpc_finish(txn)

    def test_load_blob_from_cache(self, storage, blob_cache, tmp_path):
        oid = p64(1)
        tid = self._store_and_commit(storage, oid, b"cached blob", tmp_path)

        # Should be in cache already (from tpc_finish)
        result = storage.loadBlob(oid, tid)
        assert os.path.exists(result)
        with open(result, "rb") as f:
            assert f.read() == b"cached blob"

    def test_load_blob_from_s3(self, storage, blob_cache, s3_client, tmp_path):
        oid = p64(1)
        tid = self._store_and_commit(storage, oid, b"s3 blob", tmp_path)

        # Remove from cache to force S3 download
        cached = blob_cache.get(oid, tid)
        if cached and os.path.exists(cached):
            os.remove(cached)

        result = storage.loadBlob(oid, tid)
        assert os.path.exists(result)
        with open(result, "rb") as f:
            assert f.read() == b"s3 blob"

    def test_load_blob_not_found(self, storage):
        from ZODB.POSException import POSKeyError

        with pytest.raises(POSKeyError):
            storage.loadBlob(p64(999), p64(999))

    def test_open_committed_blob_file(self, storage, tmp_path):
        oid = p64(1)
        tid = self._store_and_commit(storage, oid, b"open test", tmp_path)

        f = storage.openCommittedBlobFile(oid, tid)
        try:
            assert f.read() == b"open test"
        finally:
            f.close()

    def test_open_committed_blob_file_with_blob(self, storage, tmp_path):
        from ZODB.blob import Blob

        oid = p64(1)
        tid = self._store_and_commit(storage, oid, b"blob file", tmp_path)

        blob = Blob()
        f = storage.openCommittedBlobFile(oid, tid, blob=blob)
        try:
            assert f.read() == b"blob file"
        finally:
            f.close()


class TestLoadPendingBlob:
    """Verify loadBlob returns pending (staged) blobs during a transaction."""

    def test_load_blob_returns_pending_before_vote(self, storage, tmp_path):
        """After storeBlob but before tpc_vote, loadBlob finds the staged file."""
        oid = p64(1)
        blob_content = b"pending blob data"
        blob_path = _make_blob_file(tmp_path, blob_content)

        txn = transaction.get()
        storage.tpc_begin(txn)
        storage.storeBlob(oid, p64(0), b"pickle", blob_path, "", txn)

        # Before tpc_vote: blob is only in _pending_blobs
        result = storage.loadBlob(oid, p64(0))
        assert os.path.exists(result)
        with open(result, "rb") as f:
            assert f.read() == blob_content

        storage.tpc_abort(txn)

    def test_load_blob_returns_pending_after_vote(self, storage, tmp_path):
        """After tpc_vote (blob in S3 AND pending), loadBlob still works."""
        oid = p64(1)
        blob_content = b"voted blob data"
        blob_path = _make_blob_file(tmp_path, blob_content)

        txn = transaction.get()
        storage.tpc_begin(txn)
        storage.storeBlob(oid, p64(0), b"pickle", blob_path, "", txn)
        storage.tpc_vote(txn)

        tid = storage._extract_base_tid()
        result = storage.loadBlob(oid, tid)
        assert os.path.exists(result)
        with open(result, "rb") as f:
            assert f.read() == blob_content

        storage.tpc_finish(txn)

    def test_open_committed_blob_file_with_pending(self, storage, tmp_path):
        """openCommittedBlobFile works with pending blobs."""
        oid = p64(1)
        blob_content = b"open pending blob"
        blob_path = _make_blob_file(tmp_path, blob_content)

        txn = transaction.get()
        storage.tpc_begin(txn)
        storage.storeBlob(oid, p64(0), b"pickle", blob_path, "", txn)

        f = storage.openCommittedBlobFile(oid, p64(0))
        try:
            assert f.read() == blob_content
        finally:
            f.close()

        storage.tpc_abort(txn)


class TestNewInstance:
    def test_new_instance_shares_s3_and_cache(self, storage, s3_client, blob_cache):
        new = storage.new_instance()
        assert new._s3_client is s3_client
        assert new._cache is blob_cache

    def test_new_instance_returns_different_wrapper(self, storage):
        new = storage.new_instance()
        assert new is not storage
        assert isinstance(new, S3BlobStorage)


class TestPack:
    def _store_blob_and_commit(self, storage, oid, blob_content, tmp_path):
        blob_path = _make_blob_file(tmp_path, blob_content)
        txn = transaction.get()
        storage.tpc_begin(txn)
        storage.storeBlob(oid, p64(0), b"pickle", blob_path, "", txn)
        storage.tpc_vote(txn)
        return storage.tpc_finish(txn)

    def _store_root(self, storage):
        """Store a root object (oid 0) required for MappingStorage GC pack."""
        from ZODB.utils import z64

        txn = transaction.get()
        storage.tpc_begin(txn)
        storage.store(z64, z64, b"root", "", txn)
        storage.tpc_vote(txn)
        storage.tpc_finish(txn)

    def test_pack_delegates_to_base(self, storage, tmp_path):
        import time

        self._store_root(storage)
        oid = p64(1)
        self._store_blob_and_commit(storage, oid, b"pack test", tmp_path)

        # Pack should not raise
        storage.pack(time.time(), lambda p: [])

    def test_pack_keeps_reachable_blobs(self, storage, s3_client, tmp_path):
        import time

        self._store_root(storage)
        oid = p64(1)
        self._store_blob_and_commit(storage, oid, b"keep me", tmp_path)

        # referencesf that returns oid=1 from root, so it stays reachable
        def referencesf(pickle):
            return [p64(1)]

        storage.pack(time.time(), referencesf)

        keys_after = list(s3_client.list_objects("blobs/"))
        assert len(keys_after) == 1  # Blob still there

    def test_pack_gc_cleans_orphaned_keys(self, storage, s3_client, tmp_path):
        """Test S3 GC by manually placing an orphan key in S3."""
        import time

        self._store_root(storage)

        # Manually upload an orphan blob to S3 (oid 999 doesn't exist in base)
        orphan_src = _make_blob_file(tmp_path, b"orphan")
        s3_client.upload_file(orphan_src, "blobs/3e7/1.blob")

        keys_before = list(s3_client.list_objects("blobs/"))
        assert len(keys_before) == 1

        storage.pack(time.time(), lambda p: [])

        # Orphan should be removed by GC
        keys_after = list(s3_client.list_objects("blobs/"))
        assert len(keys_after) == 0


class TestOidFromKey:
    def test_valid_key(self):
        oid = S3BlobStorage._oid_from_key("blobs/1/2.blob")
        assert oid == p64(1)

    def test_short_key_returns_none(self):
        assert S3BlobStorage._oid_from_key("noslash") is None

    def test_non_hex_returns_none(self):
        assert S3BlobStorage._oid_from_key("blobs/notahex/1.blob") is None

    def test_overflow_returns_none(self):
        huge_hex = "f" * 40  # much larger than 8-byte OID can hold
        assert S3BlobStorage._oid_from_key(f"blobs/{huge_hex}/1.blob") is None

    def test_rejects_non_blob_extension(self):
        assert S3BlobStorage._oid_from_key("blobs/1/2.json") is None

    def test_rejects_missing_blobs_prefix(self):
        assert S3BlobStorage._oid_from_key("other/1/2.blob") is None

    def test_rejects_uppercase_hex(self):
        assert S3BlobStorage._oid_from_key("blobs/FF/1.blob") is None

    def test_rejects_extra_segments(self):
        assert S3BlobStorage._oid_from_key("blobs/extra/1/2.blob") is None

    def test_valid_key_with_long_oid(self):
        oid = S3BlobStorage._oid_from_key("blobs/1a2b3c4d5e6f/abc.blob")
        assert oid is not None
        assert oid == p64(0x1A2B3C4D5E6F)


class TestDirectoryPermissions:
    def test_temp_dir_mode(self, storage):
        mode = stat.S_IMODE(os.stat(storage.temporaryDirectory()).st_mode)
        assert mode == 0o700

    def test_cache_dir_mode(self, blob_cache):
        mode = stat.S_IMODE(os.stat(blob_cache.cache_dir).st_mode)
        assert mode == 0o700


class TestRelStorageTidExtraction:
    """Test TID extraction from RelStorage-like base storages."""

    def _make_relstorage_mock(self, tmp_path):
        """Create a mock storage that behaves like RelStorage after tpc_vote.

        RelStorage stores TID in _tpc_phase.committing_tid_lock.tid
        instead of _tid.  We simulate this by moving the TID out of _tid
        during tpc_vote, and restoring it just before tpc_finish so
        MappingStorage's own tpc_finish still works.
        """
        base = MappingStorage()
        original_tpc_vote = base.tpc_vote
        original_tpc_finish = base.tpc_finish

        class FakeCommittingTidLock:
            def __init__(self, tid):
                self.tid = tid
                self.tid_int = int.from_bytes(tid)

        class FakeTPCPhase:
            committing_tid_lock = None

        base._tpc_phase = FakeTPCPhase()

        def patched_tpc_vote(transaction):
            original_tpc_vote(transaction)
            tid = base._tid
            base._tpc_phase.committing_tid_lock = FakeCommittingTidLock(tid)
            del base._tid

        def patched_tpc_finish(transaction, func=lambda tid: None):
            # Restore _tid so MappingStorage.tpc_finish can work
            base._tid = base._tpc_phase.committing_tid_lock.tid
            return original_tpc_finish(transaction, func)

        base.tpc_vote = patched_tpc_vote
        base.tpc_finish = patched_tpc_finish
        return base

    def test_tpc_vote_works_with_relstorage_like_base(
        self, s3_env, s3_client, blob_cache, tmp_path
    ):
        base = self._make_relstorage_mock(tmp_path)
        storage = S3BlobStorage(
            base, s3_client, blob_cache, temp_dir=str(tmp_path / "staging")
        )

        oid = p64(1)
        blob_path = _make_blob_file(tmp_path, b"relstorage test")
        txn = transaction.get()
        storage.tpc_begin(txn)
        storage.storeBlob(oid, p64(0), b"pickle", blob_path, "", txn)
        storage.tpc_vote(txn)

        keys = list(s3_client.list_objects("blobs/"))
        assert len(keys) == 1

        tid = storage.tpc_finish(txn)
        assert tid is not None

    def test_tpc_vote_still_works_with_basestorage(self, storage, s3_client, tmp_path):
        """Verify MappingStorage (BaseStorage) still works after the change."""
        oid = p64(1)
        blob_path = _make_blob_file(tmp_path, b"basestorage test")
        txn = transaction.get()
        storage.tpc_begin(txn)
        storage.storeBlob(oid, p64(0), b"pickle", blob_path, "", txn)
        storage.tpc_vote(txn)

        keys = list(s3_client.list_objects("blobs/"))
        assert len(keys) == 1

        tid = storage.tpc_finish(txn)
        assert tid is not None

    def test_lock_early_forced_for_relstorage(self, s3_env, tmp_path):
        """When base storage is RelStorage, LOCK_EARLY must be forced."""
        try:
            import relstorage.storage.tpc.vote as vote_mod
        except ImportError:
            pytest.skip("RelStorage not installed")

        from zodb_s3blobs.storage import _ensure_relstorage_lock_early

        import zodb_s3blobs.storage as storage_mod

        original = vote_mod.LOCK_EARLY
        original_flag = storage_mod._relstorage_lock_early_applied
        try:
            vote_mod.LOCK_EARLY = False
            storage_mod._relstorage_lock_early_applied = False
            _ensure_relstorage_lock_early()
            assert vote_mod.LOCK_EARLY is True
        finally:
            vote_mod.LOCK_EARLY = original
            storage_mod._relstorage_lock_early_applied = original_flag

    def test_extract_base_tid_raises_on_unknown_storage(
        self, s3_env, s3_client, blob_cache, tmp_path
    ):
        """If neither _tid nor RelStorage path exists, raise RuntimeError."""
        base = MappingStorage()
        storage = S3BlobStorage(
            base, s3_client, blob_cache, temp_dir=str(tmp_path / "staging")
        )

        oid = p64(1)
        blob_path = _make_blob_file(tmp_path, b"unknown storage")
        txn = transaction.get()
        storage.tpc_begin(txn)
        storage.storeBlob(oid, p64(0), b"pickle", blob_path, "", txn)

        original_vote = base.tpc_vote

        def vote_no_tid(transaction):
            original_vote(transaction)
            del base._tid

        base.tpc_vote = vote_no_tid

        with pytest.raises(RuntimeError, match="Cannot determine TID"):
            storage.tpc_vote(txn)

        storage.tpc_abort(txn)


class TestClose:
    def test_close(self, storage, base_storage):
        storage.close()
        # MappingStorage.close() sets _is_open to False
        # After closing, operations should fail
        assert not base_storage.opened()

    def test_close_cleans_temp_dir(self, storage):
        temp_dir = storage.temporaryDirectory()
        assert os.path.isdir(temp_dir)
        storage.close()
        assert not os.path.exists(temp_dir)

    def test_close_calls_cache_close(self, base_storage, s3_client, tmp_path):
        """close() should call cache.close() if available."""
        cache = S3BlobCache(str(tmp_path / "cache_close"), max_size=10 * 1024 * 1024)
        store = S3BlobStorage(
            base_storage, s3_client, cache, temp_dir=str(tmp_path / "staging_close")
        )
        close_called = []
        original_close = cache.close
        cache.close = lambda: (close_called.append(True), original_close())

        store.close()
        assert len(close_called) == 1
