storage.py

# ./storage.py

"""

Persistence Layer - Handles the "Save/Load" functionality using Hugging Face Dataset as a database

"""

import json
import os
from datetime import datetime
from huggingface_hub import HfApi, hf_hub_download


# --- CONFIGURATION ---
REPO_ID = "prashantmatlani/chathistorycoderg"
HISTORY_DIR = "./chathistory"

# Internal Session Authentication Toggle
_SESSION_UNLOCKED = False

# Initialize the API with your token
api = HfApi(token=os.getenv("HF_TOKEN"))

def get_secret_password():
    """Dynamically reads the environment secret to prevent cached empty variables at startup."""
    val = os.getenv("APP_PASSWORD")

    print(f"\nlen val: {len(val)}\n")

    return str(val).strip() if val is not None else ""

def verify_and_unlock(user_password_input):
    """Validates input credentials and flips the internal tracking state variable."""
    global _SESSION_UNLOCKED
    target_password = get_secret_password()
    clean_input = str(user_password_input).strip()
    
    # If no master environment password exists, unlock access automatically
    if not target_password:
        _SESSION_UNLOCKED = True
        return True
        
    if clean_input == target_password:
        _SESSION_UNLOCKED = True
        return True
        
    return False

def load_history(user_password=""):
    """Retrieves list of chat IDs from the Hub ONLY if the password matches the secret."""
    target_password = get_secret_password()
    clean_input = str(user_password).strip()

    # ====================================================================
    # 🔐 IF / ELSE SECURITY CONFIGURATION
    # ====================================================================
    
    # Condition 1: If the password field is left blank OR if it doesn't match the secret...
    if not clean_input or clean_input != target_password:
        print("[*] Security Guardrail: Invalid or empty password submitted. Bypassing history retrieval.")
        return []  # --> Then: No history loading
        
    # Condition 2: Otherwise (it matches perfectly), proceed to load history
    print("[+] Security Pass: Credentials verified. Loading environment histories...")
    
    try:
        # 1. Attempt to get rich metadata for chronological sorting
        repo_info = api.repo_info(repo_id=REPO_ID, repo_type="dataset", files_metadata=True)
        
        chat_files_meta = []
        if repo_info and hasattr(repo_info, 'siblings'):
            for f in repo_info.siblings:
                # Safe startswith check
                if f.rfind('chats/') == 0:  
                    filename_idx = f.rfind('/') + 1
                    chat_id = f[filename_idx:].replace(".json", "")
                    
                    # Safely extract last_modified timestamp if it exists
                    last_modified = getattr(f, 'last_modified', None)
                    chat_files_meta.append((chat_id, last_modified))
            
            # Sort by timestamp if available; otherwise fall back to string name
            chat_files_meta.sort(key=lambda x: x[1] if x[1] is not None else x[0], reverse=True)
            return [[item[0]] for item in chat_files_meta]

    except Exception as metadata_error:
        # Log the error so it's visible in your HF Space logs instead of swallowing it
        print(f"[!] Metadata fetch failed, switching to fallback string sort: {metadata_error}")
        
    # 2. Fallback Layer: If metadata calls fail, immediately run original working string method
    try:
        files = api.list_repo_files(repo_id=REPO_ID, repo_type="dataset")
        chat_files = [f.split("/")[-1].replace(".json", "") for f in files if f.startswith("chats/")]
        return [[f] for f in sorted(chat_files, reverse=True)]
    except Exception as fallback_error:
        print(f"[!!] Total storage access failure: {fallback_error}")
        return []

def save_chat(chat_id, history):
    """Saves chat to local subdirectory and syncs to Hugging Face Dataset."""
    if not os.path.exists(HISTORY_DIR):
        os.makedirs(HISTORY_DIR)
    
    if not chat_id:
        chat_id = datetime.now().strftime("%m%d%Y_%H%M%S")
    
    filename = f"{chat_id}.json"
    local_path = os.path.join(HISTORY_DIR, filename)
    
    with open(local_path, "w", encoding="utf-8") as f:
        json.dump(history, f, indent=4)
    
    try:
        api.upload_file(
            path_or_fileobj=local_path,
            path_in_repo=f"chats/{filename}",
            repo_id=REPO_ID,
            repo_type="dataset"
        )
    except Exception as e:
        print(f"Cloud Sync Warning: {e}")
        
    return chat_id

def get_chat_content(chat_id):
    """Loads a specific chat's content from the Hub or local cache."""
    filename = f"chats/{chat_id}.json"
    local_path = os.path.join(HISTORY_DIR, f"{chat_id}.json")
    
    try:
        if not os.path.exists(HISTORY_DIR):
            os.makedirs(HISTORY_DIR)
            
        downloaded_path = hf_hub_download(
            repo_id=REPO_ID, 
            repo_type="dataset", 
            filename=filename,
            token=os.getenv("HF_TOKEN")
        )
        with open(downloaded_path, "r", encoding="utf-8") as f:
            return json.load(f)
    except Exception:
        if os.path.exists(local_path):
            with open(local_path, "r", encoding="utf-8") as f:
                return json.load(f)
    return []