# File: data_persistence.py
# Location: /data_persistence.py
# Description: Robust data persistence and serialization for book writing project

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

class ProjectPersistenceManager:
    def __init__(self, project_id: str = None):
        """
        Initialize persistent storage for book writing projects
        
        Args:
            project_id (str, optional): Unique identifier for the project
        """
        # Ensure data directory exists
        os.makedirs('project_data', exist_ok=True)
        
        # SQLite Database Setup
        self.db_path = 'project_data/book_projects.db'
        self.conn = sqlite3.connect(self.db_path)
        self._create_tables()
    
    def _create_tables(self):
        """
        Create necessary tables for project persistence
        """
        cursor = self.conn.cursor()
        
        # Project Metadata 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
            )
        ''')
        
        # Chapters Table
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS chapters (
                project_id TEXT,
                chapter_number INTEGER,
                content TEXT,
                status TEXT,
                generated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                PRIMARY KEY (project_id, chapter_number),
                FOREIGN KEY (project_id) REFERENCES projects (project_id)
            )
        ''')
        
        # Book Concept Table
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS book_concepts (
                project_id TEXT PRIMARY KEY,
                concept_data TEXT,
                FOREIGN KEY (project_id) REFERENCES projects (project_id)
            )
        ''')
        
        self.conn.commit()
    
    def save_project_metadata(self, project_data: Dict[str, Any]):
        """
        Save project metadata to persistent storage
        
        Args:
            project_data (Dict): Project metadata dictionary
        """
        cursor = self.conn.cursor()
        
        # Upsert project metadata
        cursor.execute('''
            INSERT OR REPLACE INTO projects 
            (project_id, title, genre, total_chapters) 
            VALUES (?, ?, ?, ?)
        ''', (
            project_data.get('project_id', 'unknown'),
            project_data.get('title', 'Untitled'),
            project_data.get('genre', 'Unspecified'),
            project_data.get('total_chapters', 0)
        ))
        
        self.conn.commit()
    
    def save_book_concept(self, project_id: str, concept_data: Dict[str, Any]):
        """
        Save book concept to persistent storage
        
        Args:
            project_id (str): Unique project identifier
            concept_data (Dict): Book concept details
        """
        cursor = self.conn.cursor()
        
        # Upsert book concept
        cursor.execute('''
            INSERT OR REPLACE INTO book_concepts 
            (project_id, concept_data) 
            VALUES (?, ?)
        ''', (
            project_id, 
            json.dumps(concept_data)
        ))
        
        self.conn.commit()
    
    def save_chapter(
        self, 
        project_id: str, 
        chapter_number: int, 
        content: str, 
        status: str = 'Generated'
    ):
        """
        Save chapter content to persistent storage
        
        Args:
            project_id (str): Unique project identifier
            chapter_number (int): Chapter number
            content (str): Chapter content
            status (str): Chapter status
        """
        cursor = self.conn.cursor()
        
        # Upsert chapter
        cursor.execute('''
            INSERT OR REPLACE INTO chapters 
            (project_id, chapter_number, content, status) 
            VALUES (?, ?, ?, ?)
        ''', (
            project_id, 
            chapter_number, 
            content, 
            status
        ))
        
        self.conn.commit()
    
    def load_project_metadata(self, project_id: str) -> Dict[str, Any]:
        """
        Load project metadata from persistent storage
        
        Args:
            project_id (str): Unique project identifier
        
        Returns:
            Dict containing project metadata
        """
        cursor = self.conn.cursor()
        cursor.execute('SELECT * FROM projects WHERE project_id = ?', (project_id,))
        
        project = cursor.fetchone()
        if project:
            return {
                'project_id': project[0],
                'title': project[1],
                'genre': project[2],
                'total_chapters': project[3]
            }
        
        return {}
    
    def load_book_concept(self, project_id: str) -> Dict[str, Any]:
        """
        Load book concept from persistent storage
        
        Args:
            project_id (str): Unique project identifier
        
        Returns:
            Dict containing book concept
        """
        cursor = self.conn.cursor()
        cursor.execute('SELECT concept_data FROM book_concepts WHERE project_id = ?', (project_id,))
        
        concept = cursor.fetchone()
        return json.loads(concept[0]) if concept else {}
    
    def load_chapters(self, project_id: str) -> Dict[int, Dict]:
        """
        Load all chapters for a project
        
        Args:
            project_id (str): Unique project identifier
        
        Returns:
            Dict of chapters with chapter number as key
        """
        cursor = self.conn.cursor()
        cursor.execute('SELECT chapter_number, content, status FROM chapters WHERE project_id = ?', (project_id,))
        
        chapters = {}
        for chapter_num, content, status in cursor.fetchall():
            chapters[chapter_num] = {
                'content': content,
                'status': status
            }
        
        return chapters
    
    def close(self):
        """
        Close database connection
        """
        self.conn.close()

# Test functionality
def test_persistence():
    """
    Run comprehensive tests for data persistence
    """
    # Initialize persistence manager
    pm = ProjectPersistenceManager()
    
    # Test project metadata
    test_project = {
        'project_id': 'test_project_1',
        'title': 'Test Book',
        'genre': 'Science Fiction',
        'total_chapters': 10
    }
    
    # Save project metadata
    pm.save_project_metadata(test_project)
    
    # Load and verify project metadata
    loaded_project = pm.load_project_metadata('test_project_1')
    assert loaded_project == test_project, "Project metadata not saved/loaded correctly"
    
    # Test book concept
    test_concept = {
        'narrative_premise': 'A story about AI and humanity',
        'core_themes': ['Technology', 'Ethics']
    }
    
    pm.save_book_concept('test_project_1', test_concept)
    
    # Load and verify book concept
    loaded_concept = pm.load_book_concept('test_project_1')
    assert loaded_concept == test_concept, "Book concept not saved/loaded correctly"
    
    # Test chapter saving and loading
    pm.save_chapter('test_project_1', 1, "Chapter 1 content", "Generated")
    pm.save_chapter('test_project_1', 2, "Chapter 2 content", "In Review")
    
    # Load chapters
    loaded_chapters = pm.load_chapters('test_project_1')
    assert len(loaded_chapters) == 2, "Chapters not saved/loaded correctly"
    assert loaded_chapters[1]['content'] == "Chapter 1 content", "Chapter 1 content mismatch"
    
    print("All persistence tests passed successfully!")

# Run tests if script is executed directly
if __name__ == "__main__":
    test_persistence()