services/pdf_report.py

# /home/user/app/services/pdf_report.py
from reportlab.lib.pagesizes import letter
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Image
from reportlab.lib.styles import getSampleStyleSheet
from reportlab.lib.units import inch
from io import BytesIO
from typing import List, Optional # Optional for tool_name
from pathlib import Path

# Assuming ChatMessage is defined in models.py or models/chat.py
# Make sure ChatMessage has 'role', 'content', and optionally 'tool_name' attributes.
# Example ChatMessage structure (Pydantic or dataclass):
# class ChatMessage(BaseModel):
#     role: str  # "user", "assistant", "tool"
#     content: str
#     tool_name: Optional[str] = None
from models import ChatMessage # Or from models.chat import ChatMessage
from config.settings import settings # For APP_TITLE
from assets.logo import get_logo_path # Corrected import path
from services.logger import app_logger # For logging issues

def generate_pdf_report(chat_messages: List[ChatMessage], patient_name: str = "Patient") -> BytesIO:
    buffer = BytesIO()
    # Reduce margins for more content space
    doc = SimpleDocTemplate(buffer, pagesize=letter,
                            leftMargin=0.75*inch, rightMargin=0.75*inch,
                            topMargin=0.75*inch, bottomMargin=0.75*inch)
    styles = getSampleStyleSheet()
    story = []

    # Style adjustments
    styles['h1'].alignment = 1 # Center align H1
    styles['Normal'].spaceBefore = 6
    styles['Normal'].spaceAfter = 6
    code_style = styles['Code'] # For AI and Tool messages
    code_style.spaceBefore = 6
    code_style.spaceAfter = 6
    code_style.leftIndent = 10 # Indent AI/Tool messages slightly

    # 1. Logo (optional)
    logo_path_str = get_logo_path()
    if logo_path_str:
        logo_file = Path(logo_path_str) # Convert to Path object
        if logo_file.exists():
            try:
                # Adjust width/height as needed, preserve aspect ratio if possible
                img = Image(logo_file, width=1.0*inch, height=1.0*inch, preserveAspectRatio=True)
                img.hAlign = 'LEFT' # Align logo to the left
                story.append(img)
                story.append(Spacer(1, 0.2*inch))
            except Exception as e:
                app_logger.warning(f"Could not add logo to PDF from path '{logo_path_str}': {e}")
        else:
            app_logger.warning(f"Logo path '{logo_path_str}' from get_logo_path() does not exist.")

    # 2. Title
    title_text = f"{settings.APP_TITLE} - Consultation Report"
    title = Paragraph(title_text, styles['h1'])
    story.append(title)
    story.append(Spacer(1, 0.2*inch))

    # 3. Patient Info
    patient_info_text = f"<b>Patient:</b> {patient_name}" # Make "Patient:" bold
    patient_info = Paragraph(patient_info_text, styles['h2'])
    story.append(patient_info)
    story.append(Spacer(1, 0.3*inch)) # More space after patient info

    # 4. Chat Transcript
    story.append(Paragraph("<u>Consultation Transcript:</u>", styles['h3'])) # Underline transcript title
    story.append(Spacer(1, 0.1*inch))

    for msg in chat_messages:
        prefix = ""
        current_style = styles['Normal']

        if msg.role == 'assistant':
            prefix = "<b>AI Assistant:</b> "
            current_style = code_style
        elif msg.role == 'user':
            prefix = "<b>You:</b> "
            current_style = styles['Normal']
        elif msg.role == 'tool':
            tool_name = getattr(msg, 'tool_name', 'UnknownTool') # Handle if tool_name is missing
            prefix = f"<b>Tool ({tool_name}):</b> "
            current_style = code_style
        else:
            prefix = f"<b>{msg.role.capitalize()}:</b> " # Fallback for other roles

        # Sanitize content for ReportLab Paragraph (handles HTML-like tags)
        # Replace newlines with <br/> for PDF line breaks
        content = msg.content.replace('\n', '<br/>\n')
        # Escape < and > that are not part of <br/>
        content = content.replace("<", "<").replace(">", ">").replace("<br/>", "<br/>")

        try:
            story.append(Paragraph(prefix + content, current_style))
            story.append(Spacer(1, 0.05*inch)) # Smaller spacer between messages
        except Exception as e:
            app_logger.error(f"Error adding message to PDF: {prefix}{content}. Error: {e}")
            story.append(Paragraph(f"<i>Error rendering message: {e}</i>", styles['Italic']))


    try:
        doc.build(story)
        buffer.seek(0)
        app_logger.info(f"PDF report generated successfully for patient: {patient_name}")
    except Exception as e:
        app_logger.error(f"Failed to build PDF document: {e}", exc_info=True)
        # Return an empty or error buffer if build fails
        buffer = BytesIO() # Reset buffer
        error_doc = SimpleDocTemplate(buffer, pagesize=letter)
        error_story = [Paragraph("Error generating PDF report. Please check logs.", styles['h1'])]
        try:
            error_doc.build(error_story)
        except: pass # If even error doc fails, just return empty buffer
        buffer.seek(0)

    return buffer