api/routers/export.py

"""
Export Router for Conversation History
"""
import io
import sys
import os
from datetime import datetime
from fastapi import APIRouter, HTTPException
from fastapi.responses import StreamingResponse
from docx import Document
from docx.shared import Pt, RGBColor, Inches
from docx.enum.text import WD_ALIGN_PARAGRAPH
from reportlab.lib.pagesizes import letter
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib.units import inch
from reportlab.lib.colors import Color, HexColor
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, PageBreak
from reportlab.lib.enums import TA_LEFT, TA_RIGHT

# Add src to path for imports
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))

from core.agent import _memory_manager

router = APIRouter(tags=["export"])


def get_conversation_history(session_id: str):
    """
    Retrieve conversation history from memory manager.
    
    Args:
        session_id: Session identifier
        
    Returns:
        List of message dictionaries with 'role' and 'content'
    """
    try:
        memory = _memory_manager.get_memory(session_id)
        messages = []
        
        # Get chat history from memory
        if hasattr(memory, 'chat_memory') and hasattr(memory.chat_memory, 'messages'):
            for msg in memory.chat_memory.messages:
                # LangChain messages have type and content
                role = "User" if msg.type == "human" else "Assistant"
                messages.append({
                    "role": role,
                    "content": msg.content
                })
        
        return messages
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"Error retrieving conversation history: {str(e)}"
        )


def generate_docx(session_id: str, messages: list) -> io.BytesIO:
    """
    Generate a Word document from conversation history.
    
    Args:
        session_id: Session identifier
        messages: List of message dictionaries
        
    Returns:
        BytesIO object containing the Word document
    """
    doc = Document()
    
    # Add title
    title = doc.add_heading('Conversation History', 0)
    title.alignment = WD_ALIGN_PARAGRAPH.CENTER
    
    # Add metadata
    doc.add_paragraph(f'Session ID: {session_id}')
    doc.add_paragraph(f'Export Date: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}')
    doc.add_paragraph('')  # Empty line
    
    # Add separator
    doc.add_paragraph('_' * 80)
    doc.add_paragraph('')
    
    # Add messages
    for msg in messages:
        # Add role heading
        role_para = doc.add_paragraph()
        role_run = role_para.add_run(f"{msg['role']}:")
        role_run.bold = True
        role_run.font.size = Pt(12)
        
        if msg['role'] == 'User':
            role_run.font.color.rgb = RGBColor(255, 107, 53)  # Orange color
        else:
            role_run.font.color.rgb = RGBColor(51, 51, 51)  # Dark gray
        
        # Add message content
        content_para = doc.add_paragraph(msg['content'])
        content_para.paragraph_format.left_indent = Inches(0.25)
        
        # Add spacing
        doc.add_paragraph('')
    
    # Save to BytesIO
    buffer = io.BytesIO()
    doc.save(buffer)
    buffer.seek(0)
    
    return buffer


def generate_pdf(session_id: str, messages: list) -> io.BytesIO:
    """
    Generate a PDF document from conversation history.
    
    Args:
        session_id: Session identifier
        messages: List of message dictionaries
        
    Returns:
        BytesIO object containing the PDF document
    """
    buffer = io.BytesIO()
    doc = SimpleDocTemplate(buffer, pagesize=letter,
                           rightMargin=72, leftMargin=72,
                           topMargin=72, bottomMargin=18)
    
    # Container for the 'Flowable' objects
    elements = []
    
    # Define styles
    styles = getSampleStyleSheet()
    
    # Title style
    title_style = ParagraphStyle(
        'CustomTitle',
        parent=styles['Heading1'],
        fontSize=24,
        textColor=HexColor('#000000'),
        spaceAfter=30,
        alignment=TA_LEFT
    )
    
    # Metadata style
    meta_style = ParagraphStyle(
        'MetaStyle',
        parent=styles['Normal'],
        fontSize=10,
        textColor=HexColor('#666666'),
        spaceAfter=6
    )
    
    # User message style
    user_style = ParagraphStyle(
        'UserStyle',
        parent=styles['Normal'],
        fontSize=11,
        textColor=HexColor('#FF6B35'),
        fontName='Helvetica-Bold',
        spaceAfter=6
    )
    
    # Assistant message style
    assistant_style = ParagraphStyle(
        'AssistantStyle',
        parent=styles['Normal'],
        fontSize=11,
        textColor=HexColor('#333333'),
        fontName='Helvetica-Bold',
        spaceAfter=6
    )
    
    # Content style
    content_style = ParagraphStyle(
        'ContentStyle',
        parent=styles['Normal'],
        fontSize=10,
        leftIndent=20,
        spaceAfter=12
    )
    
    # Add title
    elements.append(Paragraph('Conversation History', title_style))
    elements.append(Spacer(1, 12))
    
    # Add metadata
    elements.append(Paragraph(f'Session ID: {session_id}', meta_style))
    elements.append(Paragraph(f'Export Date: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}', meta_style))
    elements.append(Spacer(1, 20))
    
    # Add separator line
    elements.append(Paragraph('_' * 100, meta_style))
    elements.append(Spacer(1, 12))
    
    # Add messages
    for msg in messages:
        # Add role
        role_style = user_style if msg['role'] == 'User' else assistant_style
        elements.append(Paragraph(f"{msg['role']}:", role_style))
        
        # Add content (escape HTML special characters)
        content = msg['content'].replace('&', '&amp;').replace('<', '&lt;').replace('>', '&gt;')
        elements.append(Paragraph(content, content_style))
        elements.append(Spacer(1, 12))
    
    # Build PDF
    doc.build(elements)
    buffer.seek(0)
    
    return buffer


@router.get("/export/{format}")
async def export_conversation(format: str, session_id: str = "default"):
    """
    Export conversation history as Word (.docx) or PDF (.pdf) file.
    
    Args:
        format: Export format ('docx' or 'pdf')
        session_id: Session identifier (default: "default")
        
    Returns:
        StreamingResponse with the generated document
    """
    # Validate format
    if format.lower() not in ['docx', 'pdf']:
        raise HTTPException(
            status_code=400,
            detail="Invalid format. Use 'docx' or 'pdf'"
        )
    
    # Get conversation history
    messages = get_conversation_history(session_id)
    
    if not messages:
        raise HTTPException(
            status_code=404,
            detail="No conversation history found for this session"
        )
    
    try:
        # Generate document based on format
        if format.lower() == 'docx':
            buffer = generate_docx(session_id, messages)
            media_type = "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
            filename = f"conversation_{session_id}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.docx"
        else:  # pdf
            buffer = generate_pdf(session_id, messages)
            media_type = "application/pdf"
            filename = f"conversation_{session_id}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.pdf"
        
        return StreamingResponse(
            buffer,
            media_type=media_type,
            headers={
                "Content-Disposition": f"attachment; filename={filename}"
            }
        )
        
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"Error generating {format.upper()} document: {str(e)}"
        )