from __future__ import annotations

from typing import Any

from storage.object_store import LocalObjectStore, S3ObjectStore


def test_local_object_store_scopes_keys_to_workspace_prefix(tmp_path):
    workspace = tmp_path / "workspace"
    snapshot_dir = workspace / "snapshot"
    increment_dir = workspace / "increment" / "test"
    snapshot_dir.mkdir(parents=True)
    increment_dir.mkdir(parents=True)
    (snapshot_dir / "table.csv").write_text("id\n1\n", encoding="utf-8")
    (increment_dir / "schema.json").write_text("{}", encoding="utf-8")

    store = LocalObjectStore(str(tmp_path), workspace_prefix="/workspace/")

    assert store.read_bytes("snapshot/table.csv") == b"id\n1\n"
    assert store.read_text("snapshot/table.csv") == "id\n1\n"
    assert store.exists("snapshot/table.csv") is True
    assert store.prefix_exists("increment") is True
    assert store.list_keys("snapshot") == ["snapshot/table.csv"]
    assert store.list_common_prefixes("increment") == ["increment/test/"]


def test_s3_object_store_read_text_buffers_and_closes_body(monkeypatch):
    import boto3

    captured: dict[str, str] = {}

    class FakeBody:
        closed = False

        def read(self) -> bytes:
            return b"id\n1\n"

        def close(self) -> None:
            self.closed = True

    body = FakeBody()

    class FakeClient:
        def get_object(self, **kwargs: str) -> dict[str, Any]:
            captured.update(kwargs)
            return {"Body": body}

    class FakeSession:
        def __init__(self, **kwargs: str) -> None:
            pass

        def client(self, service_name: str, **kwargs: Any) -> FakeClient:
            return FakeClient()

    monkeypatch.setattr(boto3, "Session", FakeSession)

    store = S3ObjectStore(bucket="my-bucket", prefix="task1")

    assert store.read_text("snapshot/table.csv") == "id\n1\n"

    assert body.closed is True
    assert captured == {"Bucket": "my-bucket", "Key": "task1/snapshot/table.csv"}


def test_s3_object_store_read_bytes_buffers_and_closes_body(monkeypatch):
    import boto3

    captured: dict[str, str] = {}

    class FakeBody:
        closed = False

        def read(self) -> bytes:
            return b"id\n1\n"

        def close(self) -> None:
            self.closed = True

    body = FakeBody()

    class FakeClient:
        def get_object(self, **kwargs: str) -> dict[str, Any]:
            captured.update(kwargs)
            return {"Body": body}

    class FakeSession:
        def __init__(self, **kwargs: str) -> None:
            pass

        def client(self, service_name: str, **kwargs: Any) -> FakeClient:
            return FakeClient()

    monkeypatch.setattr(boto3, "Session", FakeSession)

    store = S3ObjectStore(bucket="my-bucket", prefix="task1")

    assert store.read_bytes("snapshot/table.csv") == b"id\n1\n"

    assert body.closed is True
    assert captured == {"Bucket": "my-bucket", "Key": "task1/snapshot/table.csv"}