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

# %% auto 0
__all__ = ['MANIFEST_FILENAME', 'Metadata', 'get_storage_client', 'store_blob', 'get_blob', 'list_blobs']

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

import dotenv
from google.cloud import storage
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
    additional_info: Dict[str, Any] = Field(default_factory=dict)

# %% ../nbs/00_core.ipynb 6
def get_storage_client():
    # TODO shy: mock storage provider with local file system
    return storage.Client()

# %% ../nbs/00_core.ipynb 7
def store_blob(client_id: str, blob_name: str, data: bytes, metadata: Metadata) -> 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)
    
    # 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_blob(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 = Metadata(**metadata_dict[blob_name])
    
    return data, metadata

def list_blobs(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]]
