# tests/integration/core/test_dependency_injection.py

import msgpack
from typing import Iterator
from beautyspot import Spot
from beautyspot.storage import BlobStorageBase, LocalStorage, ReadableBuffer
from beautyspot.db import TaskDBBase, SQLiteTaskDB


class MockStorage(BlobStorageBase):
    def __init__(self):
        self.data = {}

    def save(self, key: str, data: ReadableBuffer) -> str:
        self.data[key] = bytes(data)
        return f"mock://{key}"

    def load(self, location: str) -> bytes:
        key = location.replace("mock://", "")
        if key not in self.data:
            raise FileNotFoundError(key)
        return self.data[key]

    def delete(self, location: str) -> None:
        key = location.replace("mock://", "")
        if key in self.data:
            del self.data[key]

    def list_keys(self) -> Iterator[str]:
        """
        Yields location identifiers (mock://...) for all stored data.
        Required by BlobStorageBase interface.
        """
        for key in self.data:
            yield f"mock://{key}"

    def get_mtime(self, location: str) -> float:
        """Required by BlobStorageBase interface."""
        import time

        return time.time()


class MockDB(TaskDBBase):
    def __init__(self):
        self.store = {}

    def init_schema(self):
        pass

    def get(self, cache_key: str, *, include_expired: bool = False):
        return self.store.get(cache_key)

    def save(
        self,
        cache_key,
        func_name,
        func_identifier,
        input_id,
        version,
        result_type,
        content_type,
        result_value=None,
        result_data=None,
        expires_at=None,
    ):
        self.store[cache_key] = {
            "func_name": func_name,
            "result_value": result_value,
            "result_data": result_data,
            "result_type": result_type,
        }

    def delete(self, cache_key: str) -> bool:
        if cache_key in self.store:
            del self.store[cache_key]
            return True
        return False

    def get_history(self, limit=1000):
        import pandas as pd

        return pd.DataFrame(list(self.store.values()))


def test_custom_storage_injection(tmp_path):
    """Test injecting a custom storage backend."""
    storage = MockStorage()
    project = Spot(
        name="di_test",
        db=SQLiteTaskDB(str(tmp_path / "test.db")),
        storage_backend=storage,
    )

    @project.mark(save_blob=True)
    def blob_task():
        return "blob_data"

    res = blob_task()
    assert res == "blob_data"

    # Verify data ended up in our mock storage
    assert len(storage.data) == 1
    # Key is md5 hash, so we just check if any key exists
    key = list(storage.data.keys())[0]
    assert (
        b"blob_data" in storage.data[key]
    )  # msgpack serialized string might contain the string itself


def test_custom_db_injection(tmp_path):
    """Test injecting a custom DB backend."""

    db = MockDB()
    project = Spot(
        name="di_test", db=db, storage_backend=LocalStorage(tmp_path / "blobs")
    )

    @project.mark
    def simple_task():
        return 123

    res = simple_task()
    assert res == 123

    # Verify metadata ended up in our mock DB
    assert len(db.store) == 1
    key = list(db.store.keys())[0]

    # DIRECT_B64 ではなく DIRECT_BLOB を確認
    assert db.store[key]["result_type"] == "DIRECT_BLOB"

    # result_value (Base64) ではなく result_data (Bytes) を確認
    stored_bytes = db.store[key]["result_data"]
    assert msgpack.unpackb(stored_bytes) == 123
