# 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()