utils/pdfutils.py

import io
import re  # Add missing import
from typing import Union, Dict
from reportlab.lib.pagesizes import letter
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib.enums import TA_JUSTIFY, TA_LEFT, TA_CENTER
from reportlab.platypus import Table, TableStyle
from reportlab.lib import colors
import os
import json  # Add missing import
from datetime import datetime

class PDFGenerator:
    def __init__(self):
        self.styles = getSampleStyleSheet()
        self._setup_styles()

    def _setup_styles(self):
        """Setup custom styles for PDF generation"""
        # Add custom styles
        self.styles.add(
            ParagraphStyle(
                name='CustomTitle',
                parent=self.styles['Heading1'],
                fontSize=24,
                spaceAfter=30,
                alignment=TA_CENTER,
                textColor='navy'
            )
        )
        
        # Add SectionHeader style
        self.styles.add(
            ParagraphStyle(
                name='SectionHeader',
                parent=self.styles['Heading2'],
                fontSize=16,
                spaceBefore=20,
                spaceAfter=12,
                textColor='navy',
                alignment=TA_LEFT
            )
        )
        
        self.styles.add(
            ParagraphStyle(
                name='CustomContent',
                parent=self.styles['Normal'],
                fontSize=12,
                spaceAfter=12,
                alignment=TA_LEFT,
                leading=16
            )
        )

    def generate_pdf_from_text(self, text: str, filename: str, output_dir: str = "discharge_papers") -> str:
            """Generate a PDF file from plain text and return the file path."""
            try:
                if not os.path.exists(output_dir):
                    os.makedirs(output_dir)
                filepath = os.path.join(output_dir, filename)
                doc = SimpleDocTemplate(filepath, pagesize=letter, rightMargin=72, leftMargin=72, topMargin=72, bottomMargin=72)
                elements = []

                # Add title
                elements.append(Paragraph("Discharge Summary", self.styles['CustomTitle']))
                elements.append(Spacer(1, 20))

                # Split text into lines and format
                lines = text.split('\n')
                for line in lines:
                    if line.strip():
                        if ':' in line and not line.strip().startswith('-'):
                            elements.append(Paragraph(line.strip(), self.styles['SectionHeader']))
                            elements.append(Spacer(1, 12))
                        else:
                            elements.append(Paragraph(line.strip(), self.styles['CustomContent']))
                            elements.append(Spacer(1, 8))

                # Build PDF
                doc.build(elements)
                return filepath
            except Exception as e:
                print(f"Error generating PDF from text: {str(e)}")
                return None

    def generate_pdf(self, content: Union[str, Dict]) -> io.BytesIO:
        """Generate PDF with improved formatting"""
        try:
            buffer = io.BytesIO()
            doc = SimpleDocTemplate(
                buffer, 
                pagesize=letter,
                rightMargin=72,
                leftMargin=72,
                topMargin=72,
                bottomMargin=72
            )
            elements = []

            # Process title and content
            if isinstance(content, dict):
                title = content.get('title', 'Medical Document')
                content_text = content.get('content', '')
                if isinstance(content_text, dict):
                    content_text = json.dumps(content_text, indent=2)
            else:
                title = "Medical Document"
                content_text = str(content)

            # Add title
            elements.append(Paragraph(title, self.styles['CustomTitle']))
            elements.append(Spacer(1, 20))

            # Process content sections
            sections = content_text.split('\n')
            for section in sections:
                if section.strip():
                    # Check if this is a header
                    if section.startswith('##') or section.startswith('# '):
                        header_text = section.lstrip('#').strip()
                        elements.append(Paragraph(header_text, self.styles['SectionHeader']))
                        elements.append(Spacer(1, 12))
                    else:
                        # Handle bullet points and regular text
                        if section.strip().startswith('-'):
                            text = '•' + section.strip()[1:]
                        else:
                            text = section.strip()
                        elements.append(Paragraph(text, self.styles['CustomContent']))
                        elements.append(Spacer(1, 8))

            # Build PDF
            doc.build(elements)
            buffer.seek(0)
            return buffer

        except Exception as e:
            print(f"PDF Generation error: {str(e)}")
            # Create error PDF
            buffer = io.BytesIO()
            doc = SimpleDocTemplate(buffer, pagesize=letter)
            elements = [
                Paragraph("Error Generating Document", self.styles['CustomTitle']),
                Spacer(1, 12),
                Paragraph(f"An error occurred: {str(e)}", self.styles['CustomContent'])
            ]
            doc.build(elements)
            buffer.seek(0)
            return buffer

    def _format_content(self, content: Union[str, Dict]) -> str:
        """Format content for PDF generation"""
        if isinstance(content, str):
            return content
        elif isinstance(content, dict):
            try:
                # Handle nested content
                if 'content' in content:
                    return self._format_content(content['content'])
                # Format dictionary as string
                return "\n".join(f"{k}: {v}" for k, v in content.items())
            except Exception as e:
                return f"Error formatting content: {str(e)}"
        else:
            return str(content)
        
    def generate_discharge_form(
        self,
        patient_info: dict,
        discharge_info: dict,
        diagnosis_info: dict,
        medication_info: dict,
        prepared_by: dict
    ) -> io.BytesIO:
        """
        Generate a PDF that replicates the 'Patient Discharge Form' layout.
        patient_info: {
        "first_name": "...",
        "last_name": "...",
        "dob": "YYYY-MM-DD",
        "age": "...",
        "sex": "...",
        "mobile": "...",
        "address": "...",
        "city": "...",
        "state": "...",
        "zip": "..."
        }
        discharge_info: {
        "date_of_admission": "...",
        "date_of_discharge": "...",
        "source_of_admission": "...",
        "mode_of_admission": "...",
        "discharge_against_advice": "Yes/No"
        }
        diagnosis_info: {
        "diagnosis": "...",
        "operation_procedure": "...",
        "treatment": "...",
        "follow_up": "..."
        }
        medication_info: {
        "medications": [ "Med1", "Med2", ...],
        "instructions": "..."
        }
        prepared_by: {
        "name": "...",
        "title": "...",
        "signature": "..."
        }
        """
        buffer = io.BytesIO()
        doc = SimpleDocTemplate(
            buffer,
            pagesize=letter,
            rightMargin=72,
            leftMargin=72,
            topMargin=72,
            bottomMargin=72
        )
        elements = []

        # Title
        elements.append(Paragraph("Patient Discharge Form", self.styles['CustomTitle']))
        elements.append(Spacer(1, 20))

        # ---- Patient Details ----
        elements.append(Paragraph("Patient Details", self.styles['SectionHeader']))
        patient_data_table = [
            ["First Name", patient_info.get("first_name", ""), "Last Name", patient_info.get("last_name", "")],
            ["Date of Birth", patient_info.get("dob", ""), "Age", patient_info.get("age", "")],
            ["Sex", patient_info.get("sex", ""), "Mobile", patient_info.get("mobile", "")],
            ["Address", patient_info.get("address", ""), "City", patient_info.get("city", "")],
            ["State", patient_info.get("state", ""), "Zip", patient_info.get("zip", "")]
        ]
        table_style = TableStyle([
            ('GRID', (0,0), (-1,-1), 0.5, colors.grey),
            ('BACKGROUND', (0,0), (-1,0), colors.whitesmoke),
            ('VALIGN', (0,0), (-1,-1), 'TOP'),
            ('FONTNAME', (0,0), (-1,0), 'Helvetica-Bold'),
        ])
        pt = Table(patient_data_table, colWidths=[100,150,100,150])
        pt.setStyle(table_style)
        elements.append(pt)
        elements.append(Spacer(1, 16))

        # ---- Admission and Discharge Details ----
        elements.append(Paragraph("Admission and Discharge Details", self.styles['SectionHeader']))
        ad_data_table = [
            ["Date of Admission", discharge_info.get("date_of_admission", ""), "Date of Discharge", discharge_info.get("date_of_discharge", "")],
            ["Source of Admission", discharge_info.get("source_of_admission", ""), "Mode of Admission", discharge_info.get("mode_of_admission", "")],
            ["Discharge Against Advice", discharge_info.get("discharge_against_advice", "No"), "", ""]
        ]
        ad_table = Table(ad_data_table, colWidths=[120,130,120,130])
        ad_table.setStyle(table_style)
        elements.append(ad_table)
        elements.append(Spacer(1, 16))

        # ---- Diagnosis & Procedures ----
        elements.append(Paragraph("Diagnosis & Procedures", self.styles['SectionHeader']))
        diag_table_data = [
            ["Diagnosis", diagnosis_info.get("diagnosis", "")],
            ["Operation / Procedure", diagnosis_info.get("operation_procedure", "")],
            ["Treatment", diagnosis_info.get("treatment", "")],
            ["Follow-up", diagnosis_info.get("follow_up", "")]
        ]
        diag_table = Table(diag_table_data, colWidths=[150, 330])
        diag_table.setStyle(table_style)
        elements.append(diag_table)
        elements.append(Spacer(1, 16))

        # ---- Medication Details ----
        elements.append(Paragraph("Medication Details", self.styles['SectionHeader']))
        meds_joined = ", ".join(medication_info.get("medications", []))
        med_table_data = [
            ["Medications", meds_joined],
            ["Instructions", medication_info.get("instructions", "")]
        ]
        med_table = Table(med_table_data, colWidths=[100, 380])
        med_table.setStyle(table_style)
        elements.append(med_table)
        elements.append(Spacer(1, 16))

        # ---- Prepared By ----
        elements.append(Paragraph("Prepared By", self.styles['SectionHeader']))
        prepared_table_data = [
            ["Name", prepared_by.get("name", ""), "Title", prepared_by.get("title", "")],
            ["Signature", prepared_by.get("signature", ""), "", ""]
        ]
        prepared_table = Table(prepared_table_data, colWidths=[80,180,80,180])
        prepared_table.setStyle(table_style)
        elements.append(prepared_table)
        elements.append(Spacer(1, 16))

        # Build PDF
        doc.build(elements)
        buffer.seek(0)
        return buffer

class DischargeDocumentCreator:
    def __init__(self, output_dir='discharge_papers'):
        self.output_dir = output_dir
        self.styles = getSampleStyleSheet()
        self.title_style = ParagraphStyle(
            'TitleStyle', 
            parent=self.styles['Heading1'],
            alignment=1,  # Center alignment
            spaceAfter=12
        )
        
        # Ensure output directory exists
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
    
    def generate_discharge_paper(self, patient_data, llm_content):
        """
        Generate a discharge paper with patient data and LLM-generated content
        
        Args:
            patient_data (dict): Patient information including name, DOB, admission date, etc.
            llm_content (dict): LLM-generated content for different sections
        
        Returns:
            str: Path to the generated PDF
        """
        # Create filename based on patient name and current date
        filename = f"{patient_data['patient_id']}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.pdf"
        filepath = os.path.join(self.output_dir, filename)
        
        # Create the document
        doc = SimpleDocTemplate(filepath, pagesize=letter,
                               rightMargin=72, leftMargin=72,
                               topMargin=72, bottomMargin=72)
        
        # Build content
        content = []
        
        # Title
        content.append(Paragraph("HOSPITAL DISCHARGE SUMMARY", self.title_style))
        content.append(Spacer(1, 12))
        
        # Patient information table
        patient_info = [
            ["Patient Name:", patient_data.get('name', 'N/A')],
            ["Date of Birth:", patient_data.get('dob', 'N/A')],
            ["Patient ID:", patient_data.get('patient_id', 'N/A')],
            ["Admission Date:", patient_data.get('admission_date', 'N/A')],
            ["Discharge Date:", datetime.now().strftime("%Y-%m-%d")],
            ["Attending Physician:", patient_data.get('physician', 'N/A')]
        ]
        
        t = Table(patient_info, colWidths=[150, 350])
        t.setStyle(TableStyle([
            ('GRID', (0, 0), (-1, -1), 0.5, colors.grey),
            ('BACKGROUND', (0, 0), (0, -1), colors.lightgrey),
            ('VALIGN', (0, 0), (-1, -1), 'MIDDLE'),
            ('PADDING', (0, 0), (-1, -1), 6)
        ]))
        content.append(t)
        content.append(Spacer(1, 20))
        
        # Add LLM-generated sections
        sections = [
            ("Diagnosis", llm_content.get('diagnosis', 'No diagnosis provided.')),
            ("Treatment Summary", llm_content.get('treatment', 'No treatment summary provided.')),
            ("Medications", llm_content.get('medications', 'No medications listed.')),
            ("Follow-up Instructions", llm_content.get('follow_up', 'No follow-up instructions provided.')),
            ("Special Instructions", llm_content.get('special_instructions', 'No special instructions provided.'))
        ]
        
        for title, content_text in sections:
            content.append(Paragraph(title, self.styles['Heading2']))
            content.append(Paragraph(content_text, self.styles['Normal']))
            content.append(Spacer(1, 12))
        
        # Build the document
        doc.build(content)
        
        return filepath

def generate_discharge_summary(patient_data, llm_content, output_dir='discharge_papers'):
    """
    Wrapper function to generate a discharge summary document
    
    Args:
        patient_data (dict): Patient information
        llm_content (dict): LLM-generated content for discharge summary
        output_dir (str): Directory to save the PDF file
        
    Returns:
        str: Path to the generated PDF
    """
    creator = DischargeDocumentCreator(output_dir=output_dir)
    return creator.generate_discharge_paper(patient_data, llm_content)

if __name__ == "__main__":
    # Test the PDF generator with different methods
    pdf_gen = PDFGenerator()
    
    # Test data for generate_discharge_form
    patient_info = {
        "first_name": "John",
        "last_name": "Doe",
        "dob": "1980-05-15",
        "age": "43",
        "sex": "Male",
        "mobile": "555-123-4567",
        "address": "123 Main Street",
        "city": "Anytown",
        "state": "CA",
        "zip": "12345"
    }
    
    discharge_info = {
        "date_of_admission": "2023-10-01",
        "date_of_discharge": "2023-10-10",
        "source_of_admission": "Emergency",
        "mode_of_admission": "Ambulance",
        "discharge_against_advice": "No"
    }
    
    diagnosis_info = {
        "diagnosis": "Appendicitis with successful appendectomy",
        "operation_procedure": "Laparoscopic appendectomy",
        "treatment": "Antibiotics, pain management, IV fluids",
        "follow_up": "Follow up with Dr. Smith in 2 weeks"
    }
    
    medication_info = {
        "medications": ["Amoxicillin 500mg 3x daily for 7 days", "Ibuprofen 400mg as needed for pain"],
        "instructions": "Take antibiotics with food. Avoid driving while taking pain medication."
    }
    
    prepared_by = {
        "name": "Dr. Jane Smith",
        "title": "Attending Physician",
        "signature": "J. Smith, MD"
    }
    
    # Generate discharge form
    discharge_form_pdf = pdf_gen.generate_discharge_form(
        patient_info,
        discharge_info,
        diagnosis_info,
        medication_info,
        prepared_by
    )
    
    # Save discharge form to file
    with open("discharge_form_sample.pdf", "wb") as f:
        f.write(discharge_form_pdf.read())
    print("Discharge form saved as discharge_form_sample.pdf")
    
    # Test data for generate_discharge_summary
    patient_data = {
        "name": "John Doe",
        "dob": "1980-05-15",
        "patient_id": "P12345",
        "admission_date": "2023-10-01",
        "physician": "Dr. Jane Smith"
    }
    
    llm_content = {
        "diagnosis": "Acute appendicitis requiring surgical intervention.",
        "treatment": "Patient underwent successful laparoscopic appendectomy on 2023-10-02. Post-operative recovery was uneventful with good pain control and return of bowel function.",
        "medications": "1. Amoxicillin 500mg capsules, take 1 capsule 3 times daily for 7 days\n2. Ibuprofen 400mg tablets, take 1-2 tablets every 6 hours as needed for pain",
        "follow_up": "Please schedule a follow-up appointment with Dr. Smith in 2 weeks. Return sooner if experiencing fever, increasing pain, or wound drainage.",
        "special_instructions": "Keep incision sites clean and dry. No heavy lifting (>10 lbs) for 4 weeks. May shower 24 hours after surgery."
    }
    
    # Generate discharge summary
    summary_path = generate_discharge_summary(patient_data, llm_content)
    print(f"Discharge summary saved as {summary_path}")