# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_core.ipynb.

# %% auto 0
__all__ = ['MANIFEST_FILENAME', 'Metadata', 'get_storage_client', 'Sojourner']

# %% ../nbs/00_core.ipynb 3
from datetime import datetime
import hashlib
import os
from pathlib import Path
import json
from typing import Any, Dict, Tuple, List

import dotenv
from google.cloud import storage
from google.oauth2 import service_account
from pydantic import BaseModel, Field

# %% ../nbs/00_core.ipynb 4
dotenv.load_dotenv()

MANIFEST_FILENAME = "manifest.json"

# %% ../nbs/00_core.ipynb 5
class Metadata(BaseModel):
    manifest: str
    timestamp: float = Field(
        default_factory=lambda: int(datetime.now().timestamp()), allow_mutation=False
    )
    content_hash: str
    additional_info: Dict[str, Any] = Field(default_factory=dict)

    def __str__(self):
        return f"Metadata(manifest: {self.manifest}, timestamp: {datetime.fromtimestamp(self.timestamp).strftime('%Y-%m-%d %H:%M:%S')}, content_hash: {self.content_hash[:8]}..., additional_info: {json.dumps(self.additional_info)})"

# %% ../nbs/00_core.ipynb 6
def get_storage_client(
    credentials_path=os.environ.get("GOOGLE_APPLICATION_CREDENTIALS"),
):
    # TODO shy: mock storage provider with local file system
    credentials = service_account.Credentials.from_service_account_file(
        credentials_path
    )
    return storage.Client(credentials=credentials)

# %% ../nbs/00_core.ipynb 7
class Sojourner:
    def __init__(self, credentials_path=None):
        if credentials_path:
            self.client = gtet_storage_client(credentials_path)
        else:
            self.client = get_storage_client()

    def store(
        self, client_id: str, blob_name: str, data: bytes, manifest: str, **kwargs
    ) -> bool:
        """
        Store a new blob with associated metadata.
        Returns true if successful, false otherwise.
        Overwriting a blob is not allowed.
        """
        client = get_storage_client()
        bucket = client.bucket(os.environ["BUCKET"])

        client_folder = Path(client_id)
        blob_path = client_folder / blob_name
        manifest_path = client_folder / MANIFEST_FILENAME

        # Create the client folder if it doesn't exist
        if not list(bucket.list_blobs(prefix=str(client_folder))):
            bucket.blob(f"{client_folder}/").upload_from_string("")

        # Check if the blob already exists
        blob = bucket.blob(str(blob_path))
        if blob.exists():
            return False  # Blob already exists, don't overwrite

        # Store the blob
        blob.upload_from_string(data)

        # Construct metadata
        metadata = Metadata(
            manifest=manifest,
            content_hash=hashlib.sha256(data).hexdigest(),
            additional_info=kwargs,
        )

        # Store metadata
        metadata_blob = bucket.blob(str(manifest_path))
        try:
            existing_metadata = json.loads(metadata_blob.download_as_text())
        except:
            existing_metadata = {}

        existing_metadata[blob_name] = metadata.model_dump()
        metadata_blob.upload_from_string(json.dumps(existing_metadata))

        return True

    def get(self, client_id: str, blob_name: str) -> Tuple[bytes, Metadata]:
        """
        Retrieve a blob and its metadata.
        Returns a tuple of (blob_data, metadata).
        """
        client = get_storage_client()
        bucket = client.bucket(os.environ["BUCKET"])

        client_folder = Path(client_id)
        blob_path = client_folder / blob_name
        manifest_path = client_folder / MANIFEST_FILENAME

        # Get the blob
        blob = bucket.blob(str(blob_path))
        data = blob.download_as_bytes()

        # Get metadata
        metadata_blob = bucket.blob(str(manifest_path))
        metadata_dict = json.loads(metadata_blob.download_as_text())
        metadata_entry = metadata_dict.get(blob_name, {})
        metadata = Metadata(**metadata_entry)

        return data, metadata

    def list(self, client_id: str) -> List[str]:
        """
        List all blob names for a given client.
        """
        client = get_storage_client()
        bucket = client.bucket(os.environ["BUCKET"])

        client_folder = Path(client_id)
        blobs = bucket.list_blobs(prefix=str(client_folder))
        return [
            Path(blob.name).name
            for blob in blobs
            if Path(blob.name).name not in ["", client_id, MANIFEST_FILENAME]
        ]

    def list_all_directories(self) -> List[str]:
        """
        List all directories and subdirectories in the bucket.
        """
        client = get_storage_client()
        bucket = client.bucket(os.environ["BUCKET"])

        blobs = bucket.list_blobs()
        directories = set()

        for blob in blobs:
            path = Path(blob.name)
            while len(path.parts) > 1:
                directories.add(str(path.parent))
                path = path.parent

        return sorted(list(directories))

    def list_parent_directories(self) -> List[str]:
        """
        List only the top-level (parent) directories in the bucket.
        """
        client = get_storage_client()
        bucket = client.bucket(os.environ["BUCKET"])

        blobs = bucket.list_blobs()
        parent_dirs = set()

        for blob in blobs:
            path = Path(blob.name)
            if len(path.parts) > 1:
                parent_dirs.add(path.parts[0])

        return sorted(list(parent_dirs))
