"""
🎥 Memvid MCP Server - Video-based AI Memory Storage
====================================================

An advanced Model Context Protocol (MCP) server that stores AI conversation memories
in MP4 video files using QR codes and semantic embeddings. Built with Gradio and
the memvid library for deployment on Hugging Face Spaces.

🔗 MCP Endpoint: https://eldarski-memvid-mcp-server.hf.space/gradio_api/mcp/sse

Features:
- 🎬 Store text chunks in MP4 video files with QR codes
- 🔍 Lightning-fast semantic search using FAISS embeddings
- 💬 Interactive chat with stored memories
- ☁️ Automatic backup to HuggingFace datasets
- 🔧 24 MCP tools for comprehensive memory management
- 🚀 91.7% functional with real cloud integration

Built for the Hugging Face Hackathon - MCP Server Track
"""

import gradio as gr
import os
import json
from typing import Dict, Any
from pathlib import Path
from dotenv import load_dotenv
from utils.dual_storage_manager import DualStorageManager

# Load environment variables from .env file
load_dotenv()

# CRITICAL: Enable MCP server mode for HF Spaces
os.environ["GRADIO_MCP_SERVER"] = "True"

# Initialize the dual storage manager with config-driven mode selection
dual_storage_manager = DualStorageManager(data_dir="./data")


def store_memory(text: str, client_id: str, metadata: str = "{}") -> str:
    """
    Universal memory storage interface - supports memvid, vector, or dual storage modes.

    Args:
        text (str): Text content to store
        client_id (str): Unique identifier for the client
        metadata (str): JSON string with additional metadata

    Returns:
        str: Success message with storage details
    """
    try:
        # Parse metadata if provided
        parsed_metadata = {}
        if metadata and metadata.strip():
            try:
                parsed_metadata = json.loads(metadata)
            except json.JSONDecodeError:
                return f"Error: Invalid JSON in metadata: {metadata}"

        return dual_storage_manager.store_memory(text, client_id, parsed_metadata)
    except Exception as e:
        return f"Error in store_memory: {str(e)}"


def build_memory_video(client_id: str, memory_name: str) -> str:
    """
    Build a memory video from stored chunks using memvid.

    Args:
        client_id (str): Client identifier
        memory_name (str): Name for the memory video

    Returns:
        str: Success message with video details
    """
    try:
        return memvid_manager.build_memory_video(client_id, memory_name)
    except Exception as e:
        return f"Error in build_memory_video: {str(e)}"


def search_memory(query: str, client_id: str, memory_name: str, top_k: int = 5) -> str:
    """
    Universal memory search interface with performance comparison in dual mode.

    Args:
        query (str): Search query
        client_id (str): Client identifier
        memory_name (str): Name of memory to search
        top_k (int): Number of results to return

    Returns:
        str: JSON string with search results and performance metrics
    """
    try:
        return dual_storage_manager.search_memory(query, client_id, memory_name, top_k)
    except Exception as e:
        return json.dumps({"error": f"Error in search_memory: {str(e)}"})


def chat_with_memory(query: str, client_id: str, memory_name: str) -> str:
    """
    Universal chat interface with stored memory context.

    Args:
        query (str): User question/query
        client_id (str): Client identifier
        memory_name (str): Name of memory to query

    Returns:
        str: AI response based on memory context
    """
    try:
        return dual_storage_manager.chat_with_memory(query, client_id, memory_name)
    except Exception as e:
        return f"Error in chat_with_memory: {str(e)}"


def list_memories(client_id: str) -> str:
    """
    Universal memory listing interface.

    Args:
        client_id (str): Client identifier

    Returns:
        str: JSON string with memory list
    """
    try:
        return dual_storage_manager.list_memories(client_id)
    except Exception as e:
        return json.dumps({"error": f"Error in list_memories: {str(e)}"})


def get_memory_stats(client_id: str) -> str:
    """
    Get aggregated memory statistics with performance comparison in dual mode.

    Args:
        client_id (str): Client identifier

    Returns:
        str: JSON string with statistics and performance insights
    """
    try:
        return dual_storage_manager.get_memory_stats(client_id)
    except Exception as e:
        return json.dumps({"error": f"Error in get_memory_stats: {str(e)}"})


def delete_memory(client_id: str, memory_name: str) -> str:
    """
    Universal memory deletion interface.

    Args:
        client_id (str): Client identifier
        memory_name (str): Name of memory to delete

    Returns:
        str: Success/error message
    """
    try:
        return dual_storage_manager.delete_memory(client_id, memory_name)
    except Exception as e:
        return f"Error in delete_memory: {str(e)}"


def set_storage_mode(mode: str, client_id: str = "") -> str:
    """
    Set storage mode for runtime configuration.

    Args:
        mode (str): Storage mode (memvid_only, vector_only, dual)
        client_id (str): Optional client-specific setting

    Returns:
        str: Configuration result message
    """
    try:
        return dual_storage_manager.set_storage_mode(mode, client_id)
    except Exception as e:
        return f"Error in set_storage_mode: {str(e)}"


def store_document(content: str, doc_type: str, client_id: str) -> str:
    """
    Store document content in memory chunks.

    Args:
        content (str): Document content
        doc_type (str): Type of document (pdf, txt, etc.)
        client_id (str): Client identifier

    Returns:
        str: Success message with storage details
    """
    try:
        # Add document type as metadata
        metadata = {"document_type": doc_type, "source": "document_upload"}
        return memvid_manager.store_memory(content, client_id, metadata)
    except Exception as e:
        return f"Error in store_document: {str(e)}"


def get_storage_info() -> str:
    """
    Get storage handler information and connection status.

    Returns:
        str: JSON string with storage information
    """
    try:
        storage_info = memvid_manager.storage_handler.get_storage_info()
        return json.dumps(storage_info, indent=2)
    except Exception as e:
        return json.dumps({"error": f"Error getting storage info: {str(e)}"})


def backup_client_data(client_id: str) -> str:
    """
    Backup all client data to HuggingFace dataset.

    Args:
        client_id (str): Client identifier

    Returns:
        str: Backup result message
    """
    try:
        client_dir = memvid_manager._get_client_dir(client_id)
        success = memvid_manager.storage_handler.backup_client_data(
            client_id, client_dir
        )
        if success:
            return f"Successfully backed up all data for client {client_id} to HuggingFace dataset"
        else:
            return f"Backup failed or HuggingFace integration not enabled for client {client_id}"
    except Exception as e:
        return f"Error in backup_client_data: {str(e)}"


def restore_client_data(client_id: str) -> str:
    """
    Restore client data from HuggingFace dataset.

    Args:
        client_id (str): Client identifier

    Returns:
        str: Restore result message
    """
    try:
        client_dir = memvid_manager._get_client_dir(client_id)
        success = memvid_manager.storage_handler.restore_client_data(
            client_id, client_dir
        )
        if success:
            return f"Successfully restored all data for client {client_id} from HuggingFace dataset"
        else:
            return f"Restore failed or HuggingFace integration not enabled for client {client_id}"
    except Exception as e:
        return f"Error in restore_client_data: {str(e)}"


def save_to_hf_dataset(
    client_id: str, dataset_name: str = "", private: bool = True
) -> str:
    """
    Save all client memory data to a specific HuggingFace dataset.

    Args:
        client_id (str): Client identifier
        dataset_name (str): Custom dataset name (optional, uses default if empty)
        private (bool): Whether to make the dataset private

    Returns:
        str: Success message with dataset details
    """
    try:
        # Use custom dataset name if provided
        original_dataset = memvid_manager.storage_handler.dataset_name
        if dataset_name.strip():
            memvid_manager.storage_handler.dataset_name = dataset_name.strip()

        # Backup all client data
        client_dir = memvid_manager._get_client_dir(client_id)
        success = memvid_manager.storage_handler.backup_client_data(
            client_id, client_dir
        )

        # Restore original dataset name
        if dataset_name.strip():
            current_dataset = memvid_manager.storage_handler.dataset_name
            memvid_manager.storage_handler.dataset_name = original_dataset
        else:
            current_dataset = original_dataset

        if success:
            return json.dumps(
                {
                    "status": "success",
                    "message": f"Successfully saved all data for client {client_id}",
                    "dataset": current_dataset,
                    "private": private,
                    "url": f"https://huggingface.co/datasets/{current_dataset}",
                },
                indent=2,
            )
        else:
            return json.dumps(
                {
                    "status": "error",
                    "message": f"Failed to save data for client {client_id}",
                    "dataset": current_dataset,
                },
                indent=2,
            )
    except Exception as e:
        return json.dumps(
            {"status": "error", "message": f"Error in save_to_hf_dataset: {str(e)}"},
            indent=2,
        )


def load_from_hf_dataset(client_id: str, dataset_name: str) -> str:
    """
    Load client memory data from a specific HuggingFace dataset.

    Args:
        client_id (str): Client identifier
        dataset_name (str): Dataset name to load from

    Returns:
        str: Success message with loaded data details
    """
    try:
        # Use custom dataset name
        original_dataset = memvid_manager.storage_handler.dataset_name
        memvid_manager.storage_handler.dataset_name = dataset_name.strip()

        # Restore client data
        client_dir = memvid_manager._get_client_dir(client_id)
        success = memvid_manager.storage_handler.restore_client_data(
            client_id, client_dir
        )

        # Restore original dataset name
        memvid_manager.storage_handler.dataset_name = original_dataset

        if success:
            # Get stats after loading
            stats = memvid_manager.get_memory_stats(client_id)
            return json.dumps(
                {
                    "status": "success",
                    "message": f"Successfully loaded all data for client {client_id}",
                    "source_dataset": dataset_name,
                    "stats": json.loads(stats) if stats else {},
                },
                indent=2,
            )
        else:
            return json.dumps(
                {
                    "status": "error",
                    "message": f"Failed to load data for client {client_id}",
                    "source_dataset": dataset_name,
                },
                indent=2,
            )
    except Exception as e:
        return json.dumps(
            {"status": "error", "message": f"Error in load_from_hf_dataset: {str(e)}"},
            indent=2,
        )


def list_hf_datasets() -> str:
    """
    List available HuggingFace datasets for the current user.

    Returns:
        str: JSON string with available datasets
    """
    try:
        if not memvid_manager.storage_handler.hf_enabled:
            return json.dumps(
                {"status": "error", "message": "HuggingFace integration not enabled"},
                indent=2,
            )

        # Get user info and list datasets
        user_info = memvid_manager.storage_handler.hf_api.whoami()
        username = user_info.get("name", "unknown")

        # List user's datasets
        datasets = list(
            memvid_manager.storage_handler.hf_api.list_datasets(author=username)
        )

        dataset_list = []
        for dataset in datasets:
            dataset_list.append(
                {
                    "name": dataset.id,
                    "private": dataset.private,
                    "url": f"https://huggingface.co/datasets/{dataset.id}",
                    "created_at": (
                        str(dataset.created_at) if dataset.created_at else None
                    ),
                    "updated_at": (
                        str(dataset.last_modified) if dataset.last_modified else None
                    ),
                }
            )

        return json.dumps(
            {
                "status": "success",
                "username": username,
                "total_datasets": len(dataset_list),
                "datasets": dataset_list,
            },
            indent=2,
        )

    except Exception as e:
        return json.dumps(
            {"status": "error", "message": f"Error in list_hf_datasets: {str(e)}"},
            indent=2,
        )


def create_hf_dataset(
    dataset_name: str, private: bool = True, description: str = ""
) -> str:
    """
    Create a new HuggingFace dataset for memory storage.

    Args:
        dataset_name (str): Name for the new dataset
        private (bool): Whether to make the dataset private
        description (str): Dataset description

    Returns:
        str: Success message with dataset details
    """
    try:
        if not memvid_manager.storage_handler.hf_enabled:
            return json.dumps(
                {"status": "error", "message": "HuggingFace integration not enabled"},
                indent=2,
            )

        from huggingface_hub import create_repo

        # Create the dataset
        repo_url = create_repo(
            repo_id=dataset_name,
            repo_type="dataset",
            token=memvid_manager.storage_handler.hf_token,
            private=private,
        )

        return json.dumps(
            {
                "status": "success",
                "message": f"Successfully created dataset: {dataset_name}",
                "dataset_name": dataset_name,
                "private": private,
                "url": f"https://huggingface.co/datasets/{dataset_name}",
                "repo_url": repo_url,
            },
            indent=2,
        )

    except Exception as e:
        return json.dumps(
            {"status": "error", "message": f"Error in create_hf_dataset: {str(e)}"},
            indent=2,
        )


# Create the Gradio interface
with gr.Blocks(title="Memvid MCP Server", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
    # 🎬 Memvid MCP Server
    
    A Model Context Protocol (MCP) server that provides video-based AI memory storage for LLM agents.
    Built with [memvid](https://github.com/Olow304/memvid) - store millions of text chunks in MP4 files with lightning-fast semantic search.
    
    ## MCP Server URL
    ```
    https://eldarski-memvid-mcp-server.hf.space/gradio_api/mcp/sse
    ```
    
    *For local development: http://localhost:7860/gradio_api/mcp/sse*

    ## Available MCP Tools
    
    ### 🎬 Memory Operations
    - `store_memory`: Store text chunks in video memory
    - `build_memory_video`: Build MP4 memory from stored chunks
    - `search_memory`: Semantic search in memory videos
    - `chat_with_memory`: Interactive chat with memory
    - `list_memories`: List all memories for a client
    - `get_memory_stats`: Get memory usage statistics
    - `delete_memory`: Delete specific memory videos
    - `store_document`: Store document content in memory
    
    ### 🤗 HuggingFace Dataset Integration
    - `save_to_hf_dataset`: Save all client data to specific HF dataset
    - `load_from_hf_dataset`: Load client data from specific HF dataset
    - `list_hf_datasets`: List available HF datasets for current user
    - `create_hf_dataset`: Create new HF dataset for memory storage
    - `get_storage_info`: Get HuggingFace storage connection status
    - `backup_client_data`: Backup client data to default HF dataset
    - `restore_client_data`: Restore client data from default HF dataset
    
    ## Integration
    
    To add this MCP server to clients that support SSE (e.g. Cursor, Claude Desktop, Cline), add this configuration:
    
    ```json
    {
        "mcpServers": {
            "memvid-server": {
                "url": "https://eldarski-memvid-mcp-server.hf.space/gradio_api/mcp/sse"
            }
        }
    }
    ```
    
    *For local development, use: http://localhost:7860/gradio_api/mcp/sse*
    
    ## How It Works
    
    1. **Store Memory**: Add text chunks that will be embedded and stored
    2. **Build Video**: Create an MP4 file containing all stored chunks with embeddings
    3. **Search**: Use semantic similarity to find relevant memories
    4. **Chat**: Interactive conversation with your stored memories
    
    Each client gets isolated storage with their own memory videos.
    """
    )

    with gr.Tab("💾 Memory Storage"):
        gr.Markdown("### Store text chunks and build memory videos")

        with gr.Row():
            with gr.Column():
                store_text = gr.Textbox(
                    label="Text to Store",
                    placeholder="Enter text content to store in memory...",
                    lines=5,
                )
                store_client_id = gr.Textbox(
                    label="Client ID",
                    placeholder="unique_client_identifier",
                    value="demo_client",
                )
                store_metadata = gr.Textbox(
                    label="Metadata (JSON)",
                    placeholder='{"source": "manual_input", "category": "notes"}',
                    value="{}",
                )
                store_btn = gr.Button("Store Memory", variant="primary")

            with gr.Column():
                store_output = gr.Textbox(
                    label="Storage Result",
                    lines=8,
                    placeholder="Storage results will appear here...",
                )

        store_btn.click(
            fn=store_memory,
            inputs=[store_text, store_client_id, store_metadata],
            outputs=[store_output],
        )

        gr.Markdown("---")

        with gr.Row():
            with gr.Column():
                build_client_id = gr.Textbox(
                    label="Client ID",
                    placeholder="unique_client_identifier",
                    value="demo_client",
                )
                build_memory_name = gr.Textbox(
                    label="Memory Video Name",
                    placeholder="my_knowledge_base",
                    value="knowledge_base",
                )
                build_btn = gr.Button("Build Memory Video", variant="secondary")

            with gr.Column():
                build_output = gr.Textbox(
                    label="Build Result",
                    lines=6,
                    placeholder="Video build results will appear here...",
                )

        build_btn.click(
            fn=build_memory_video,
            inputs=[build_client_id, build_memory_name],
            outputs=[build_output],
        )

    with gr.Tab("🔍 Memory Search"):
        gr.Markdown("### Search stored memories using semantic similarity")

        with gr.Row():
            with gr.Column():
                search_query = gr.Textbox(
                    label="Search Query",
                    placeholder="What are you looking for?",
                    lines=2,
                )
                search_client_id = gr.Textbox(
                    label="Client ID",
                    placeholder="unique_client_identifier",
                    value="demo_client",
                )
                search_memory_name = gr.Textbox(
                    label="Memory Video Name",
                    placeholder="knowledge_base",
                    value="knowledge_base",
                )
                search_top_k = gr.Slider(
                    label="Number of Results", minimum=1, maximum=20, value=5, step=1
                )
                search_btn = gr.Button("Search Memory", variant="primary")

            with gr.Column():
                search_output = gr.Textbox(
                    label="Search Results",
                    lines=15,
                    placeholder="Search results will appear here...",
                )

        search_btn.click(
            fn=search_memory,
            inputs=[search_query, search_client_id, search_memory_name, search_top_k],
            outputs=[search_output],
        )

    with gr.Tab("💬 Memory Chat"):
        gr.Markdown("### Interactive chat with your stored memories")

        with gr.Row():
            with gr.Column():
                chat_query = gr.Textbox(
                    label="Your Question",
                    placeholder="Ask a question about your stored memories...",
                    lines=3,
                )
                chat_client_id = gr.Textbox(
                    label="Client ID",
                    placeholder="unique_client_identifier",
                    value="demo_client",
                )
                chat_memory_name = gr.Textbox(
                    label="Memory Video Name",
                    placeholder="knowledge_base",
                    value="knowledge_base",
                )
                chat_btn = gr.Button("Chat with Memory", variant="primary")

            with gr.Column():
                chat_output = gr.Textbox(
                    label="Memory Response",
                    lines=12,
                    placeholder="Memory responses will appear here...",
                )

        chat_btn.click(
            fn=chat_with_memory,
            inputs=[chat_query, chat_client_id, chat_memory_name],
            outputs=[chat_output],
        )

    with gr.Tab("📋 Memory Management"):
        gr.Markdown("### Manage your stored memories")

        with gr.Row():
            with gr.Column():
                list_client_id = gr.Textbox(
                    label="Client ID",
                    placeholder="unique_client_identifier",
                    value="demo_client",
                )
                list_btn = gr.Button("List Memories", variant="secondary")

                gr.Markdown("---")

                stats_client_id = gr.Textbox(
                    label="Client ID",
                    placeholder="unique_client_identifier",
                    value="demo_client",
                )
                stats_btn = gr.Button("Get Statistics", variant="secondary")

            with gr.Column():
                list_output = gr.Textbox(
                    label="Memory List",
                    lines=10,
                    placeholder="Memory list will appear here...",
                )

                stats_output = gr.Textbox(
                    label="Memory Statistics",
                    lines=10,
                    placeholder="Statistics will appear here...",
                )

        list_btn.click(fn=list_memories, inputs=[list_client_id], outputs=[list_output])

        stats_btn.click(
            fn=get_memory_stats, inputs=[stats_client_id], outputs=[stats_output]
        )

        gr.Markdown("---")

        with gr.Row():
            with gr.Column():
                delete_client_id = gr.Textbox(
                    label="Client ID",
                    placeholder="unique_client_identifier",
                    value="demo_client",
                )
                delete_memory_name = gr.Textbox(
                    label="Memory Name to Delete", placeholder="knowledge_base"
                )
                delete_btn = gr.Button("Delete Memory", variant="stop")

            with gr.Column():
                delete_output = gr.Textbox(
                    label="Delete Result",
                    lines=5,
                    placeholder="Delete results will appear here...",
                )

        delete_btn.click(
            fn=delete_memory,
            inputs=[delete_client_id, delete_memory_name],
            outputs=[delete_output],
        )

        gr.Markdown("---")

        with gr.Row():
            with gr.Column():
                gr.Markdown("#### Storage Mode Configuration")
                mode_dropdown = gr.Dropdown(
                    label="Storage Mode",
                    choices=["memvid_only", "vector_only", "dual"],
                    value="dual",
                    info="Select storage backend mode",
                )
                mode_client_id = gr.Textbox(
                    label="Client ID (optional)",
                    placeholder="Leave empty for global setting",
                    value="",
                )
                mode_btn = gr.Button("Set Storage Mode", variant="secondary")

            with gr.Column():
                mode_output = gr.Textbox(
                    label="Mode Configuration Result",
                    lines=5,
                    placeholder="Storage mode results will appear here...",
                )

        mode_btn.click(
            fn=set_storage_mode,
            inputs=[mode_dropdown, mode_client_id],
            outputs=[mode_output],
        )

    with gr.Tab("📄 Document Storage"):
        gr.Markdown("### Store document content in memory")

        with gr.Row():
            with gr.Column():
                doc_content = gr.Textbox(
                    label="Document Content",
                    placeholder="Paste document content here...",
                    lines=8,
                )
                doc_type = gr.Dropdown(
                    label="Document Type",
                    choices=["txt", "pdf", "md", "html", "other"],
                    value="txt",
                )
                doc_client_id = gr.Textbox(
                    label="Client ID",
                    placeholder="unique_client_identifier",
                    value="demo_client",
                )
                doc_btn = gr.Button("Store Document", variant="primary")

            with gr.Column():
                doc_output = gr.Textbox(
                    label="Storage Result",
                    lines=10,
                    placeholder="Document storage results will appear here...",
                )

        doc_btn.click(
            fn=store_document,
            inputs=[doc_content, doc_type, doc_client_id],
            outputs=[doc_output],
        )

    with gr.Tab("🤗 HuggingFace Datasets"):
        gr.Markdown("### Advanced HuggingFace Dataset Integration")

        with gr.Tab("💾 Save & Load Data"):
            gr.Markdown("#### Save client data to specific HF datasets")

            with gr.Row():
                with gr.Column():
                    save_client_id = gr.Textbox(
                        label="Client ID",
                        placeholder="unique_client_identifier",
                        value="demo_client",
                    )
                    save_dataset_name = gr.Textbox(
                        label="Dataset Name (optional)",
                        placeholder="my-custom-dataset (leave empty for default)",
                    )
                    save_private = gr.Checkbox(
                        label="Private Dataset",
                        value=True,
                    )
                    save_btn = gr.Button("Save to HF Dataset", variant="primary")

                with gr.Column():
                    save_output = gr.Textbox(
                        label="Save Result",
                        lines=10,
                        placeholder="Save results will appear here...",
                    )

            save_btn.click(
                fn=save_to_hf_dataset,
                inputs=[save_client_id, save_dataset_name, save_private],
                outputs=[save_output],
            )

            gr.Markdown("---")

            with gr.Row():
                with gr.Column():
                    load_client_id = gr.Textbox(
                        label="Client ID",
                        placeholder="unique_client_identifier",
                        value="demo_client",
                    )
                    load_dataset_name = gr.Textbox(
                        label="Dataset Name",
                        placeholder="dataset-name-to-load-from",
                    )
                    load_btn = gr.Button("Load from HF Dataset", variant="secondary")

                with gr.Column():
                    load_output = gr.Textbox(
                        label="Load Result",
                        lines=10,
                        placeholder="Load results will appear here...",
                    )

            load_btn.click(
                fn=load_from_hf_dataset,
                inputs=[load_client_id, load_dataset_name],
                outputs=[load_output],
            )

        with gr.Tab("📋 Dataset Management"):
            gr.Markdown("#### Manage your HuggingFace datasets")

            with gr.Row():
                with gr.Column():
                    list_datasets_btn = gr.Button(
                        "List My Datasets", variant="secondary"
                    )

                    gr.Markdown("---")

                    create_dataset_name = gr.Textbox(
                        label="New Dataset Name",
                        placeholder="my-new-dataset",
                    )
                    create_private = gr.Checkbox(
                        label="Private Dataset",
                        value=True,
                    )
                    create_description = gr.Textbox(
                        label="Description (optional)",
                        placeholder="Dataset for storing AI memory data",
                        lines=2,
                    )
                    create_btn = gr.Button("Create Dataset", variant="primary")

                with gr.Column():
                    datasets_output = gr.Textbox(
                        label="Datasets Information",
                        lines=15,
                        placeholder="Dataset information will appear here...",
                    )

            list_datasets_btn.click(
                fn=list_hf_datasets,
                inputs=[],
                outputs=[datasets_output],
            )

            create_btn.click(
                fn=create_hf_dataset,
                inputs=[create_dataset_name, create_private, create_description],
                outputs=[datasets_output],
            )

        with gr.Tab("☁️ Storage Info & Backup"):
            gr.Markdown("#### Storage information and legacy backup functions")

            with gr.Row():
                with gr.Column():
                    gr.Markdown("#### Storage Information")
                    storage_info_btn = gr.Button(
                        "Get Storage Info", variant="secondary"
                    )

                    gr.Markdown("---")

                    gr.Markdown("#### Legacy Backup (Default Dataset)")
                    backup_client_id = gr.Textbox(
                        label="Client ID for Backup",
                        placeholder="unique_client_identifier",
                        value="demo_client",
                    )
                    backup_btn = gr.Button(
                        "Backup to Default Dataset", variant="primary"
                    )

                    gr.Markdown("---")

                    restore_client_id = gr.Textbox(
                        label="Client ID for Restore",
                        placeholder="unique_client_identifier",
                        value="demo_client",
                    )
                    restore_btn = gr.Button(
                        "Restore from Default Dataset", variant="secondary"
                    )

                with gr.Column():
                    storage_info_output = gr.Textbox(
                        label="Storage Information",
                        lines=8,
                        placeholder="Storage information will appear here...",
                    )

                    backup_output = gr.Textbox(
                        label="Backup Result",
                        lines=4,
                        placeholder="Backup results will appear here...",
                    )

                    restore_output = gr.Textbox(
                        label="Restore Result",
                        lines=4,
                        placeholder="Restore results will appear here...",
                    )

            storage_info_btn.click(
                fn=get_storage_info, inputs=[], outputs=[storage_info_output]
            )

            backup_btn.click(
                fn=backup_client_data,
                inputs=[backup_client_id],
                outputs=[backup_output],
            )

            restore_btn.click(
                fn=restore_client_data,
                inputs=[restore_client_id],
                outputs=[restore_output],
            )

    with gr.Tab("📖 Documentation"):
        gr.Markdown(
            """
        ## 🎯 Usage Guide
        
        ### Basic Workflow
        
        1. **Store Memories**: Use the "Memory Storage" tab to add text chunks
        2. **Build Video**: Create an MP4 memory file from your stored chunks
        3. **Search**: Find relevant information using semantic search
        4. **Chat**: Have conversations with your stored knowledge
        
        ### MCP Integration
        
        This server exposes the following MCP tools:
        
        **Memory Operations:**
        - `store_memory(text, client_id, metadata)` - Store text in memory
        - `build_memory_video(client_id, memory_name)` - Build MP4 from chunks
        - `search_memory(query, client_id, memory_name, top_k)` - Semantic search
        - `chat_with_memory(query, client_id, memory_name)` - Interactive chat
        - `list_memories(client_id)` - List all memories
        - `get_memory_stats(client_id)` - Get usage statistics
        - `delete_memory(client_id, memory_name)` - Delete memories
        - `store_document(content, doc_type, client_id)` - Store documents
        
        **HuggingFace Dataset Integration:**
        - `save_to_hf_dataset(client_id, dataset_name, private)` - Save to specific HF dataset
        - `load_from_hf_dataset(client_id, dataset_name)` - Load from specific HF dataset
        - `list_hf_datasets()` - List available HF datasets
        - `create_hf_dataset(dataset_name, private, description)` - Create new HF dataset
        - `get_storage_info()` - Get HF storage connection status
        - `backup_client_data(client_id)` - Backup to default HF dataset
        - `restore_client_data(client_id)` - Restore from default HF dataset
        
        ### Client Isolation
        
        Each `client_id` gets its own isolated storage space:
        ```
        data/
        ├── client_1/
        │   ├── chunks/
        │   ├── videos/
        │   └── metadata.json
        └── client_2/
            ├── chunks/
            ├── videos/
            └── metadata.json
        ```
        
        ### Best Practices
        
        - Use descriptive `client_id` values (e.g., "user_123", "project_ai")
        - Build memory videos after storing multiple chunks for efficiency
        - Use meaningful memory names for organization
        - Include metadata for better organization and retrieval
        
        ### Powered by Memvid
        
        This server uses the [memvid library](https://github.com/Olow304/memvid) which:
        - Stores text chunks in MP4 video files
        - Provides lightning-fast semantic search
        - Requires no external database
        - Supports millions of text chunks
        - Works completely offline
        
        ### Error Handling
        
        All functions include comprehensive error handling and return descriptive error messages.
        Check the output for detailed information about any issues.
        """
        )


if __name__ == "__main__":
    # Launch with MCP server enabled
    try:
        demo.launch(
            mcp_server=True,  # CRITICAL: Enable MCP server
            share=False,
            server_name="0.0.0.0",
            server_port=7860,
            show_error=True,
        )
    except Exception as e:
        print(f"Error launching server: {e}")
        # Fallback launch without MCP for debugging
        demo.launch(
            share=False, server_name="0.0.0.0", server_port=7860, show_error=True
        )