# storage_handler.py
import os
import base64
import logging
from typing import Optional, BinaryIO, Dict, Any

from huggingface_hub import HfFileSystem, HfApi
from cryptography.hazmat.primitives.ciphers.aead import AESGCM

logger = logging.getLogger("StorageHandler")

# ------------------------------------------------------------------------
# Configuration – all from environment variables
# ------------------------------------------------------------------------
HF_TOKEN = os.environ.get("HF_TOKEN")
if not HF_TOKEN:
    raise ValueError("HF_TOKEN environment variable is required")

BUCKET_ID = os.environ.get("INFINITY_CHAT_BUCKET", "infinitychat-data")
FILE_ENCRYPTION_KEY_B64 = os.environ.get("FILE_ENCRYPTION_KEY", None)
if FILE_ENCRYPTION_KEY_B64 is None:
    from cryptography.fernet import Fernet
    FILE_ENCRYPTION_KEY_B64 = Fernet.generate_key().decode()
    logger.warning("FILE_ENCRYPTION_KEY not set – generated random key. Set it permanently!")
FILE_ENCRYPTION_KEY = base64.urlsafe_b64decode(FILE_ENCRYPTION_KEY_B64.encode())
assert len(FILE_ENCRYPTION_KEY) == 32, "AES‑256 requires exactly 32 bytes"

# Get the bucket owner (your HF username)
_api = HfApi(token=HF_TOKEN)
try:
    OWNER = _api.whoami()["name"]
except Exception:
    # Fallback: use the bucket ID prefix (if it contains a slash)
    if "/" in BUCKET_ID:
        OWNER, BUCKET_ID = BUCKET_ID.split("/", 1)
    else:
        raise ValueError("Cannot determine bucket owner. Set INFINITY_CHAT_BUCKET as 'owner/bucket-name'")

# Build the proper bucket URI prefix (Xet-native)
BUCKET_URI_PREFIX = f"hf://buckets/{OWNER}/{BUCKET_ID}"

# ------------------------------------------------------------------------
# HfFileSystem for low-level bucket operations (Xet-native, NO Git/LFS)
# ------------------------------------------------------------------------
_fs = HfFileSystem(token=HF_TOKEN)

# ------------------------------------------------------------------------
# Bucket initialization (using the modern Hugging Face Hub API)
# ------------------------------------------------------------------------
def ensure_bucket():
    """
    Ensure the bucket exists using the native HF Bucket API.
    The HfApi.create_bucket call does NOT use 'name' – it uses 'bucket_id'.
    """
    api = HfApi(token=HF_TOKEN)
    
    # Modern API: create_bucket(bucket_id=..., private=True, exist_ok=True)
    try:
        api.create_bucket(
            bucket_id=f"{OWNER}/{BUCKET_ID}",
            private=True,
            exist_ok=True
        )
        logger.info(f"✅ Bucket '{OWNER}/{BUCKET_ID}' is ready (private).")
    except Exception as e:
        # Some versions require the full ID
        if "already exists" in str(e).lower():
            logger.info(f"Bucket '{OWNER}/{BUCKET_ID}' already exists.")
        # Try without the owner prefix
        else:
            try:
                api.create_bucket(
                    bucket_id=BUCKET_ID,
                    private=True,
                    exist_ok=True
                )
                logger.info(f"✅ Bucket '{BUCKET_ID}' ready.")
                global BUCKET_URI_PREFIX
                BUCKET_URI_PREFIX = f"hf://buckets/{BUCKET_ID}"
            except Exception as e2:
                if "already exists" in str(e2).lower():
                    logger.info(f"Bucket '{BUCKET_ID}' already exists.")
                else:
                    logger.error(f"Failed to create bucket: {e2}")
                    raise

# Initialize bucket
ensure_bucket()

# ------------------------------------------------------------------------
# Encryption / Decryption
# ------------------------------------------------------------------------
def encrypt_file(data: bytes) -> bytes:
    """Encrypt binary data using AES‑256‑GCM. Output: nonce (12 B) || ciphertext."""
    aesgcm = AESGCM(FILE_ENCRYPTION_KEY)
    nonce = os.urandom(12)
    ciphertext = aesgcm.encrypt(nonce, data, None)
    return nonce + ciphertext

def decrypt_file(encrypted_blob: bytes) -> bytes:
    """Decrypt data produced by encrypt_file()."""
    nonce = encrypted_blob[:12]
    ciphertext = encrypted_blob[12:]
    aesgcm = AESGCM(FILE_ENCRYPTION_KEY)
    return aesgcm.decrypt(nonce, ciphertext, None)

# ------------------------------------------------------------------------
# Core Bucket Operations (Xet-native via HfFileSystem)
# ------------------------------------------------------------------------
def _bucket_path(remote_path: str) -> str:
    """Construct the full hf://buckets/... URI."""
    return f"{BUCKET_URI_PREFIX}/{remote_path}"

def store_file(remote_path: str, data: bytes) -> str:
    """
    Encrypt data and write it to the bucket using HfFileSystem.
    Returns: the remote path (for DB storage).
    """
    encrypted_data = encrypt_file(data)
    full_uri = _bucket_path(remote_path)
    
    with _fs.open(full_uri, "wb") as f:
        f.write(encrypted_data)
    
    logger.debug(f"Stored file: {full_uri}")
    return remote_path

def retrieve_file(remote_path: str) -> bytes:
    """
    Read from bucket using HfFileSystem and decrypt.
    """
    full_uri = _bucket_path(remote_path)
    
    with _fs.open(full_uri, "rb") as f:
        encrypted_data = f.read()
    
    return decrypt_file(encrypted_data)

def delete_file(remote_path: str):
    """
    Delete a file from the bucket using HfFileSystem.
    """
    full_uri = _bucket_path(remote_path)
    try:
        _fs.rm(full_uri)
        logger.debug(f"Deleted file: {full_uri}")
    except FileNotFoundError:
        logger.warning(f"File not found for deletion: {full_uri}")
    except Exception as e:
        logger.error(f"Failed to delete file {full_uri}: {e}")

def list_files(prefix: str = "") -> list:
    """
    List files in the bucket under a given prefix.
    Returns list of full URIs.
    """
    search_path = f"{BUCKET_URI_PREFIX}/{prefix}" if prefix else BUCKET_URI_PREFIX
    try:
        return _fs.ls(search_path)
    except FileNotFoundError:
        return []

def file_exists(remote_path: str) -> bool:
    """Check if a file exists in the bucket."""
    full_uri = _bucket_path(remote_path)
    return _fs.exists(full_uri)

def get_file_size(remote_path: str) -> Optional[int]:
    """Get the size of a file in bytes (compressed)."""
    full_uri = _bucket_path(remote_path)
    try:
        info = _fs.info(full_uri)
        return info.get("size")
    except Exception:
        return None

# ------------------------------------------------------------------------
# Stream-based store (for reassembled chunks)
# ------------------------------------------------------------------------
def store_file_stream(remote_path: str, data_stream: BinaryIO) -> str:
    """
    Read the entire stream, encrypt, and write to bucket.
    """
    data = data_stream.read()
    return store_file(remote_path, data)

# ------------------------------------------------------------------------
# Bucket info (for monitoring / health checks)
# ------------------------------------------------------------------------
def get_bucket_info() -> Dict[str, Any]:
    """Return information about the bucket."""
    api = HfApi(token=HF_TOKEN)
    try:
        # Try the modern API first
        info = api.bucket_info(bucket_id=f"{OWNER}/{BUCKET_ID}")
        return {
            "id": f"{OWNER}/{BUCKET_ID}",
            "exists": True,
            "private": info.get("private", True),
            "total_files": len(list_files()),
        }
    except Exception:
        return {
            "id": f"{OWNER}/{BUCKET_ID}",
            "exists": False,
            "note": "Cannot fetch bucket info"
        }

# ------------------------------------------------------------------------
# Clean shutdown
# ------------------------------------------------------------------------
def close():
    """Clean up resources (HfFileSystem)."""
    if _fs:
        _fs.close()