persistence_manager.py

# File: persistence_manager.py
# Location: /persistence_manager.py
# Description: Manages persistent storage for book writing projects

import os
import json
import streamlit as st
import sqlite3
from typing import Dict, Any, List
import uuid
from datetime import datetime

class ProjectPersistenceManager:
    """
    Manages project persistence using filesystem and SQLite
    """
    def __init__(self, base_path: str = '/data'):
        """
        Initialize persistence manager
        
        Args:
            base_path (str): Base directory for storing project data
        """
        # Ensure base path exists
        self.base_path = base_path
        self._ensure_directories()
        
        # Setup database path
        self.db_path = os.path.join(self.base_path, 'book_projects.db')
        
        # Initialize database connection
        self.conn = self._initialize_database()

    def _ensure_directories(self):
        """
        Create necessary directories for project storage
        """
        directories = [
            self.base_path,
            os.path.join(self.base_path, 'projects'),
            os.path.join(self.base_path, 'chapters')
        ]
        
        for dir_path in directories:
            os.makedirs(dir_path, exist_ok=True)

    def _initialize_database(self):
        """
        Initialize SQLite database for project tracking
        
        Returns:
            sqlite3.Connection: Database connection
        """
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            # Create projects table
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS projects (
                    project_id TEXT PRIMARY KEY,
                    title TEXT,
                    genre TEXT,
                    total_chapters INTEGER,
                    created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                    last_modified DATETIME DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            
            # Create chapters table
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS chapters (
                    project_id TEXT,
                    chapter_number INTEGER,
                    content TEXT,
                    status TEXT DEFAULT 'DRAFT',
                    created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                    FOREIGN KEY(project_id) REFERENCES projects(project_id),
                    PRIMARY KEY(project_id, chapter_number)
                )
            ''')
            
            conn.commit()
            return conn
        except Exception as e:
            st.error(f"Database initialization error: {e}")
            return None

    def save_project(self, project_data: Dict[str, Any]) -> str:
        """
        Save project data to both filesystem and database
        
        Args:
            project_data (Dict): Project information to save
        
        Returns:
            str: Project ID
        """
        # Generate project ID if not provided
        project_id = project_data.get('project_id') or str(uuid.uuid4())
        
        try:
            # Save to JSON file
            project_file = os.path.join(
                self.base_path, 
                'projects', 
                f'{project_id}.json'
            )
            
            # Prepare save data
            save_data = {
                **project_data,
                'project_id': project_id,
                'last_modified': datetime.now().isoformat()
            }
            
            # Write JSON file
            with open(project_file, 'w') as f:
                json.dump(save_data, f, indent=4)
            
            # Save to database
            if self.conn:
                cursor = self.conn.cursor()
                
                # Insert/update project metadata
                cursor.execute('''
                    INSERT OR REPLACE INTO projects 
                    (project_id, title, genre, total_chapters) 
                    VALUES (?, ?, ?, ?)
                ''', (
                    project_id, 
                    save_data.get('title', 'Untitled'),
                    save_data.get('genre', 'Unspecified'),
                    save_data.get('total_chapters', 0)
                ))
                
                # Save chapters
                if 'chapters' in save_data:
                    for chapter_num, chapter_content in save_data['chapters'].items():
                        cursor.execute('''
                            INSERT OR REPLACE INTO chapters 
                            (project_id, chapter_number, content) 
                            VALUES (?, ?, ?)
                        ''', (project_id, chapter_num, chapter_content))
                
                self.conn.commit()
            
            return project_id
        
        except Exception as e:
            st.error(f"Error saving project: {e}")
            return None

    def load_project(self, project_id: str) -> Dict[str, Any]:
        """
        Load project data from filesystem or database
        
        Args:
            project_id (str): Project identifier
        
        Returns:
            Dict: Project data or None
        """
        try:
            # Try loading from JSON file first
            project_file = os.path.join(
                self.base_path, 
                'projects', 
                f'{project_id}.json'
            )
            
            if os.path.exists(project_file):
                with open(project_file, 'r') as f:
                    return json.load(f)
            
            # Fallback to database
            if self.conn:
                cursor = self.conn.cursor()
                
                # Fetch project details
                cursor.execute('SELECT * FROM projects WHERE project_id = ?', (project_id,))
                project = cursor.fetchone()
                
                if project:
                    # Fetch chapters
                    cursor.execute('SELECT * FROM chapters WHERE project_id = ?', (project_id,))
                    chapters = cursor.fetchall()
                    
                    return {
                        'project_id': project[0],
                        'title': project[1],
                        'genre': project[2],
                        'total_chapters': project[3],
                        'chapters': {
                            chapter[1]: chapter[2] for chapter in chapters
                        }
                    }
            
            return None
        
        except Exception as e:
            st.error(f"Error loading project: {e}")
            return None

    def list_projects(self) -> List[Dict[str, Any]]:
        """
        List all saved projects
        
        Returns:
            List of project metadata
        """
        projects = []
        
        try:
            # Check JSON files
            projects_dir = os.path.join(self.base_path, 'projects')
            
            # Scan JSON files
            for filename in os.listdir(projects_dir):
                if filename.endswith('.json'):
                    try:
                        with open(os.path.join(projects_dir, filename), 'r') as f:
                            project = json.load(f)
                            projects.append({
                                'project_id': project.get('project_id'),
                                'title': project.get('title', 'Untitled'),
                                'last_modified': project.get('last_modified')
                            })
                    except Exception as e:
                        st.warning(f"Could not load project {filename}: {e}")
        
        except Exception as e:
            st.error(f"Error listing projects: {e}")
        
        return projects

def main():
    """
    Demonstration of persistence manager
    """
    # Initialize persistence manager
    pm = ProjectPersistenceManager()
    
    # Example: Save a project
    project_data = {
        'title': 'My First Book',
        'genre': 'Science Fiction',
        'total_chapters': 5,
        'chapters': {
            1: 'First chapter content',
            2: 'Second chapter content'
        }
    }
    
    # Save project
    project_id = pm.save_project(project_data)
    print(f"Saved project with ID: {project_id}")
    
    # Load project
    loaded_project = pm.load_project(project_id)
    print("Loaded Project:", loaded_project)

if __name__ == "__main__":
    main()