File size: 4,885 Bytes
8126b08
 
 
 
 
 
 
 
ec8d47f
 
 
 
 
 
 
8126b08
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
# Live PDF Generator Module
# Enables real-time PDF updates as content changes

import asyncio
import io
import tempfile
import os
from typing import Dict, Any, Optional, AsyncGenerator
try:
    from fpdf2 import FPDF
except ImportError:
    try:
        from fpdf import FPDF
    except ImportError:
        raise ImportError("Neither fpdf2 nor fpdf package found. Please install fpdf2.")
import base64

import sys
import os
sys.path.append(os.path.dirname(__file__))
from pdf_generator import create_pdf_document
from models import ExportRequest

class LivePDFGenerator:
    """Generates and streams PDF updates in real-time"""
    
    def __init__(self):
        self.active_sessions: Dict[str, 'PDFSession'] = {}
    
    async def create_live_session(self, business_plan_id: str, initial_request: ExportRequest) -> str:
        """Create a new live PDF session"""
        session_id = f"pdf_session_{business_plan_id}_{int(asyncio.get_event_loop().time())}"
        
        session = PDFSession(session_id, business_plan_id, initial_request)
        self.active_sessions[session_id] = session
        
        # Generate initial PDF
        await session.update_pdf(initial_request)
        
        return session_id
    
    async def update_session(self, session_id: str, updated_request: ExportRequest) -> bool:
        """Update an existing PDF session with new content"""
        if session_id not in self.active_sessions:
            return False
        
        session = self.active_sessions[session_id]
        await session.update_pdf(updated_request)
        return True
    
    async def stream_pdf_updates(self, session_id: str) -> AsyncGenerator[bytes, None]:
        """Stream PDF updates as they occur"""
        if session_id not in self.active_sessions:
            return
        
        session = self.active_sessions[session_id]
        async for pdf_chunk in session.stream_updates():
            yield pdf_chunk
    
    def close_session(self, session_id: str):
        """Close and cleanup a PDF session"""
        if session_id in self.active_sessions:
            del self.active_sessions[session_id]

class PDFSession:
    """Individual PDF session that manages real-time updates"""
    
    def __init__(self, session_id: str, business_plan_id: str, initial_request: ExportRequest):
        self.session_id = session_id
        self.business_plan_id = business_plan_id
        self.current_request = initial_request
        self.update_queue = asyncio.Queue()
        self.is_active = True
        self.last_pdf_bytes: Optional[bytes] = None
        
    async def update_pdf(self, new_request: ExportRequest):
        """Update the PDF with new content"""
        self.current_request = new_request
        
        try:
            # Generate new PDF
            pdf_bytes, filename = create_pdf_document(new_request)
            self.last_pdf_bytes = pdf_bytes
            
            # Notify subscribers of update
            await self.update_queue.put({
                'type': 'pdf_update',
                'data': base64.b64encode(pdf_bytes).decode('utf-8'),
                'filename': filename,
                'timestamp': asyncio.get_event_loop().time()
            })
            
        except Exception as e:
            # Notify subscribers of error
            await self.update_queue.put({
                'type': 'pdf_error',
                'error': str(e),
                'timestamp': asyncio.get_event_loop().time()
            })
    
    async def stream_updates(self) -> AsyncGenerator[bytes, None]:
        """Stream PDF updates as they occur"""
        while self.is_active:
            try:
                # Wait for updates with timeout
                update = await asyncio.wait_for(self.update_queue.get(), timeout=30.0)
                
                if update['type'] == 'pdf_update':
                    # Send SSE event for PDF update
                    event_data = f"event: pdf_update\ndata: {{\"data\": \"{update['data']}\", \"filename\": \"{update['filename']}\"}}\n\n"
                    yield event_data.encode('utf-8')
                    
                elif update['type'] == 'pdf_error':
                    # Send SSE event for error
                    error_data = f"event: pdf_error\ndata: {{\"error\": \"{update['error']}\"}}\n\n"
                    yield error_data.encode('utf-8')
                    
            except asyncio.TimeoutError:
                # Send keep-alive
                yield b": keep-alive\n\n"
            except Exception as e:
                # Send error event
                error_data = f"event: stream_error\ndata: {str(e)}\n\n"
                yield error_data.encode('utf-8')
                break
    
    def close(self):
        """Close this session"""
        self.is_active = False

# Global instance
live_pdf_generator = LivePDFGenerator()