user_storage.py

"""
user_storage.py - User-Specific Storage Management
Handles all file operations with user isolation
"""

import os
import json
import shutil
from pathlib import Path
from datetime import datetime
import uuid
import logging

logger = logging.getLogger(__name__)

# ===========================================
# CONFIGURATION
# ===========================================

# Base directory for all user data
# On HF Spaces, use /tmp for ephemeral storage (or persistent storage if mounted)
DATA_ROOT = os.environ.get('USER_DATA_ROOT', '/data/users')

# Ensure data root exists
os.makedirs(DATA_ROOT, exist_ok=True)


# ===========================================
# PATH HELPERS
# ===========================================

def get_user_dir(user_id):
    """Get the base directory for a user"""
    user_dir = os.path.join(DATA_ROOT, str(user_id))
    os.makedirs(user_dir, exist_ok=True)
    return user_dir


def get_conversations_path(user_id):
    """Get path to user's conversations file"""
    return os.path.join(get_user_dir(user_id), 'conversations.json')


def get_workspace_path(user_id):
    """Get path to user's workspace file"""
    return os.path.join(get_user_dir(user_id), 'workspace.json')


def get_user_profile_path(user_id):
    """Get path to user's profile file"""
    return os.path.join(get_user_dir(user_id), 'profile.json')


def get_user_uploads_dir(user_id):
    """Get path to user's uploads directory"""
    uploads_dir = os.path.join(get_user_dir(user_id), 'uploads')
    os.makedirs(uploads_dir, exist_ok=True)
    return uploads_dir


def get_user_generated_images_dir(user_id):
    """Get path to user's generated images directory"""
    images_dir = os.path.join(get_user_dir(user_id), 'generated_images')
    os.makedirs(images_dir, exist_ok=True)
    return images_dir


def get_user_documents_dir(user_id):
    """Get path to user's documents directory"""
    docs_dir = os.path.join(get_user_dir(user_id), 'documents')
    os.makedirs(docs_dir, exist_ok=True)
    return docs_dir


# ===========================================
# CONVERSATIONS
# ===========================================

def load_conversations(user_id):
    """
    Load user's conversations from JSON file.
    Returns empty dict if file doesn't exist.
    """
    conv_path = get_conversations_path(user_id)
    
    if not os.path.exists(conv_path):
        return {}
    
    try:
        with open(conv_path, 'r', encoding='utf-8') as f:
            conversations = json.load(f)
            
            # Ensure all conversation objects have required fields
            for conv_id, conv_data in conversations.items():
                if 'versions' not in conv_data:
                    conv_data['versions'] = {}
                if 'current_version_index' not in conv_data:
                    conv_data['current_version_index'] = {}
                if 'branch_root' not in conv_data:
                    conv_data['branch_root'] = None
                if 'version_branches' not in conv_data:
                    conv_data['version_branches'] = {}
                if 'pending_attachments' not in conv_data:
                    conv_data['pending_attachments'] = []
            
            return conversations
    except Exception as e:
        logger.error(f"Error loading conversations for user {user_id}: {e}")
        return {}


def save_conversations(user_id, conversations):
    """
    Save user's conversations to JSON file.
    """
    conv_path = get_conversations_path(user_id)
    
    try:
        with open(conv_path, 'w', encoding='utf-8') as f:
            json.dump(conversations, f, indent=2, ensure_ascii=False)
        return True
    except Exception as e:
        logger.error(f"Error saving conversations for user {user_id}: {e}")
        return False


def get_conversation(user_id, conversation_id):
    """Get a single conversation by ID"""
    conversations = load_conversations(user_id)
    return conversations.get(conversation_id)


def save_conversation(user_id, conversation_id, conversation_data):
    """Save a single conversation"""
    conversations = load_conversations(user_id)
    conversations[conversation_id] = conversation_data
    return save_conversations(user_id, conversations)


def delete_conversation(user_id, conversation_id):
    """Delete a conversation"""
    conversations = load_conversations(user_id)
    if conversation_id in conversations:
        del conversations[conversation_id]
        return save_conversations(user_id, conversations)
    return False


def get_user_conversations_list(user_id):
    """Get list of conversations (id, title, last_updated) for sidebar"""
    conversations = load_conversations(user_id)
    
    conv_list = []
    for conv_id, conv_data in conversations.items():
        conv_list.append({
            'id': conv_id,
            'title': conv_data.get('title', 'New Chat'),
            'message_count': len(conv_data.get('messages', [])),
            'last_updated': conv_data.get('last_updated', datetime.now().isoformat())
        })
    
    # Sort by last_updated descending (newest first)
    conv_list.sort(key=lambda x: x['last_updated'], reverse=True)
    return conv_list


# ===========================================
# WORKSPACE
# ===========================================

def load_workspace(user_id):
    """
    Load user's workspace structure from file.
    Returns default workspace structure if file doesn't exist.
    """
    workspace_path = get_workspace_path(user_id)
    
    if not os.path.exists(workspace_path):
        return {"folders": []}
    
    try:
        with open(workspace_path, 'r', encoding='utf-8') as f:
            return json.load(f)
    except Exception as e:
        logger.error(f"Error loading workspace for user {user_id}: {e}")
        return {"folders": []}


def save_workspace(user_id, workspace):
    """
    Save user's workspace structure to file.
    """
    workspace_path = get_workspace_path(user_id)
    
    try:
        with open(workspace_path, 'w', encoding='utf-8') as f:
            json.dump(workspace, f, indent=2, ensure_ascii=False)
        return True
    except Exception as e:
        logger.error(f"Error saving workspace for user {user_id}: {e}")
        return False


def get_file_content(user_id, file_path):
    """
    Get content of a specific file in user's workspace.
    file_path is relative to workspace root.
    """
    workspace = load_workspace(user_id)
    
    parts = file_path.split('/')
    file_name = parts[-1]
    folder_parts = parts[:-1] if len(parts) > 1 else []
    
    # Navigate to the correct folder
    current = workspace
    for folder_name in folder_parts:
        found = False
        for folder in current.get('folders', []):
            if folder.get('name') == folder_name:
                current = folder
                found = True
                break
        if not found:
            return None
    
    # Find the file
    for file_item in current.get('files', []):
        if file_item.get('name') == file_name:
            return file_item.get('content', '')
    
    return None


def save_file_content(user_id, file_path, content):
    """
    Save content to a specific file in user's workspace.
    """
    workspace = load_workspace(user_id)
    
    parts = file_path.split('/')
    file_name = parts[-1]
    folder_parts = parts[:-1] if len(parts) > 1 else []
    
    # Navigate to the correct folder
    current = workspace
    for folder_name in folder_parts:
        found = False
        for folder in current.get('folders', []):
            if folder.get('name') == folder_name:
                current = folder
                found = True
                break
        if not found:
            return False
    
    # Find and update the file
    for file_item in current.get('files', []):
        if file_item.get('name') == file_name:
            file_item['content'] = content
            file_item['updated_at'] = datetime.now().isoformat()
            return save_workspace(user_id, workspace)
    
    return False


def delete_file(user_id, file_path):
    """
    Delete a file from user's workspace.
    """
    workspace = load_workspace(user_id)
    
    parts = file_path.split('/')
    file_name = parts[-1]
    folder_parts = parts[:-1] if len(parts) > 1 else []
    
    # Navigate to the correct folder
    current = workspace
    for folder_name in folder_parts:
        found = False
        for folder in current.get('folders', []):
            if folder.get('name') == folder_name:
                current = folder
                found = True
                break
        if not found:
            return False
    
    # Find and delete the file
    files = current.get('files', [])
    for i, file_item in enumerate(files):
        if file_item.get('name') == file_name:
            del files[i]
            return save_workspace(user_id, workspace)
    
    return False


# ===========================================
# USER PROFILE
# ===========================================

def load_user_profile(user_id):
    """
    Load user's profile information.
    """
    profile_path = get_user_profile_path(user_id)
    
    if not os.path.exists(profile_path):
        # Return default profile
        return {
            'user_id': user_id,
            'username': f'user_{user_id[:8]}',
            'email': '',
            'avatar': None,
            'theme': 'dark',
            'created_at': datetime.now().isoformat(),
            'last_active': datetime.now().isoformat()
        }
    
    try:
        with open(profile_path, 'r', encoding='utf-8') as f:
            return json.load(f)
    except Exception as e:
        logger.error(f"Error loading profile for user {user_id}: {e}")
        return {'user_id': user_id, 'username': f'user_{user_id[:8]}'}


def save_user_profile(user_id, profile):
    """
    Save user's profile information.
    """
    profile_path = get_user_profile_path(user_id)
    
    try:
        with open(profile_path, 'w', encoding='utf-8') as f:
            json.dump(profile, f, indent=2, ensure_ascii=False)
        return True
    except Exception as e:
        logger.error(f"Error saving profile for user {user_id}: {e}")
        return False


def update_user_last_active(user_id):
    """Update user's last active timestamp"""
    profile = load_user_profile(user_id)
    profile['last_active'] = datetime.now().isoformat()
    return save_user_profile(user_id, profile)


# ===========================================
# IMAGES (Generated)
# ===========================================

def save_generated_image(user_id, image_data, filename):
    """
    Save a generated image for a user.
    Returns the full path to the saved image.
    """
    user_images_dir = get_user_generated_images_dir(user_id)
    filepath = os.path.join(user_images_dir, filename)
    
    with open(filepath, 'wb') as f:
        f.write(image_data)
    
    return filepath


def get_generated_image_path(user_id, filename):
    """
    Get the full path to a user's generated image.
    Returns None if file doesn't exist.
    """
    user_images_dir = get_user_generated_images_dir(user_id)
    filepath = os.path.join(user_images_dir, filename)
    
    if os.path.exists(filepath):
        return filepath
    return None


def delete_generated_image(user_id, filename):
    """Delete a user's generated image"""
    user_images_dir = get_user_generated_images_dir(user_id)
    filepath = os.path.join(user_images_dir, filename)
    
    if os.path.exists(filepath):
        os.remove(filepath)
        return True
    return False


def list_user_generated_images(user_id):
    """List all generated images for a user"""
    user_images_dir = get_user_generated_images_dir(user_id)
    
    if not os.path.exists(user_images_dir):
        return []
    
    images = []
    for filename in os.listdir(user_images_dir):
        if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.webp')):
            filepath = os.path.join(user_images_dir, filename)
            images.append({
                'filename': filename,
                'path': filepath,
                'size': os.path.getsize(filepath),
                'created_at': datetime.fromtimestamp(os.path.getctime(filepath)).isoformat()
            })
    
    # Sort by creation date (newest first)
    images.sort(key=lambda x: x['created_at'], reverse=True)
    return images


# ===========================================
# UPLOADS
# ===========================================

def save_uploaded_file(user_id, file_data, filename):
    """
    Save an uploaded file for a user.
    Returns the full path to the saved file.
    """
    user_uploads_dir = get_user_uploads_dir(user_id)
    
    # Ensure unique filename
    base, ext = os.path.splitext(filename)
    counter = 1
    filepath = os.path.join(user_uploads_dir, filename)
    while os.path.exists(filepath):
        filepath = os.path.join(user_uploads_dir, f"{base}_{counter}{ext}")
        counter += 1
    
    with open(filepath, 'wb') as f:
        f.write(file_data)
    
    return filepath


def get_uploaded_file_path(user_id, filename):
    """Get the full path to a user's uploaded file"""
    user_uploads_dir = get_user_uploads_dir(user_id)
    filepath = os.path.join(user_uploads_dir, filename)
    
    if os.path.exists(filepath):
        return filepath
    return None


def delete_uploaded_file(user_id, filename):
    """Delete a user's uploaded file"""
    user_uploads_dir = get_user_uploads_dir(user_id)
    filepath = os.path.join(user_uploads_dir, filename)
    
    if os.path.exists(filepath):
        os.remove(filepath)
        return True
    return False


# ===========================================
# DATA MIGRATION (from legacy files)
# ===========================================

def migrate_legacy_data(user_id, legacy_conversations_path, legacy_workspace_path=None):
    """
    Migrate existing legacy JSON data to a user account.
    Use this when converting from single-user to multi-user.
    
    Args:
        user_id: The user ID to assign the data to
        legacy_conversations_path: Path to old conversations.json
        legacy_workspace_path: Path to old workspace.json (optional)
    """
    migrated = {'conversations': 0, 'workspace': False}
    
    # Migrate conversations
    if os.path.exists(legacy_conversations_path):
        try:
            with open(legacy_conversations_path, 'r', encoding='utf-8') as f:
                legacy_conversations = json.load(f)
            
            # Save to user's conversations
            if save_conversations(user_id, legacy_conversations):
                migrated['conversations'] = len(legacy_conversations)
                logger.info(f"Migrated {len(legacy_conversations)} conversations to user {user_id}")
        except Exception as e:
            logger.error(f"Error migrating conversations: {e}")
    
    # Migrate workspace
    if legacy_workspace_path and os.path.exists(legacy_workspace_path):
        try:
            with open(legacy_workspace_path, 'r', encoding='utf-8') as f:
                legacy_workspace = json.load(f)
            
            if save_workspace(user_id, legacy_workspace):
                migrated['workspace'] = True
                logger.info(f"Migrated workspace to user {user_id}")
        except Exception as e:
            logger.error(f"Error migrating workspace: {e}")
    
    return migrated


def get_user_storage_stats(user_id):
    """Get storage statistics for a user"""
    stats = {
        'conversations': 0,
        'workspace_files': 0,
        'uploaded_files': 0,
        'generated_images': 0,
        'total_size_bytes': 0
    }
    
    # Count conversations
    conversations = load_conversations(user_id)
    stats['conversations'] = len(conversations)
    
    # Count workspace files
    workspace = load_workspace(user_id)
    
    def count_files(node):
        count = 0
        for folder in node.get('folders', []):
            count += count_files(folder)
        count += len(node.get('files', []))
        return count
    
    stats['workspace_files'] = count_files(workspace)
    
    # Count uploaded files
    user_uploads_dir = get_user_uploads_dir(user_id)
    if os.path.exists(user_uploads_dir):
        files = [f for f in os.listdir(user_uploads_dir) if os.path.isfile(os.path.join(user_uploads_dir, f))]
        stats['uploaded_files'] = len(files)
        for f in files:
            stats['total_size_bytes'] += os.path.getsize(os.path.join(user_uploads_dir, f))
    
    # Count generated images
    stats['generated_images'] = len(list_user_generated_images(user_id))
    
    return stats


# ===========================================
# CLEANUP UTILITIES
# ===========================================

def delete_user_data(user_id):
    """
    Delete all data for a user (admin function).
    """
    user_dir = get_user_dir(user_id)
    
    if os.path.exists(user_dir):
        shutil.rmtree(user_dir)
        logger.info(f"Deleted all data for user {user_id}")
        return True
    return False


def get_all_user_ids():
    """Get list of all user IDs with data"""
    if not os.path.exists(DATA_ROOT):
        return []
    
    user_ids = []
    for item in os.listdir(DATA_ROOT):
        item_path = os.path.join(DATA_ROOT, item)
        if os.path.isdir(item_path):
            user_ids.append(item)
    
    return user_ids