# /home/user/app/services/pdf_report.py
from reportlab.lib.pagesizes import letter
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Image, Table, TableStyle, PageBreak
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib.units import inch
from reportlab.lib import colors
from io import BytesIO
from typing import List, Dict, Any # For the input data structure
from pathlib import Path
from datetime import datetime

# from models import ChatMessage # Not strictly needed if we use dicts/mock objects
from config.settings import settings
from assets.logo import get_logo_path
from services.logger import app_logger

class MockChatMessage: # Helper class if generate_pdf_report expects objects with attributes
    def __init__(self, role: str, content: str, timestamp: datetime, tool_name: Optional[str] = None, **kwargs):
        self.role = role
        self.content = content
        self.timestamp = timestamp
        self.tool_name = tool_name
        # Allow other attributes like source_references, confidence_score if added
        for key, value in kwargs.items():
            setattr(self, key, value)


def generate_pdf_report(report_data: Dict[str, Any]) -> BytesIO:
    buffer = BytesIO()
    doc = SimpleDocTemplate(buffer, pagesize=letter,
                            leftMargin=0.75*inch, rightMargin=0.75*inch,
                            topMargin=0.75*inch, bottomMargin=0.75*inch)
    styles = getSampleStyleSheet()
    story = []

    # Custom Styles
    styles.add(ParagraphStyle(name='Justify', alignment=4, parent=styles['Normal']))
    styles.add(ParagraphStyle(name='Disclaimer', parent=styles['Italic'], fontSize=8, leading=10))
    styles.add(ParagraphStyle(name='SectionHeader', parent=styles['h2'], spaceBefore=12, spaceAfter=6, keepWithNext=1))
    styles.add(ParagraphStyle(name='SubHeader', parent=styles['h3'], spaceBefore=6, spaceAfter=3, keepWithNext=1))
    styles.add(ParagraphStyle(name='ListItem', parent=styles['Normal'], leftIndent=0.25*inch, bulletIndent=0.1*inch))
    
    user_msg_style = ParagraphStyle(name='UserMessage', parent=styles['Normal'], spaceBefore=3, spaceAfter=3, leftIndent=0.1*inch, backColor=colors.whitesmoke)
    ai_msg_style = ParagraphStyle(name='AIMessage', parent=styles['Normal'], spaceBefore=3, spaceAfter=3, leftIndent=0.1*inch, backColor=colors.lightcyan)
    tool_msg_style = ParagraphStyle(name='ToolMessage', parent=styles['Code'], spaceBefore=3, spaceAfter=3, leftIndent=0.1*inch, backColor=colors.lightgrey, textColor=colors.darkblue)
    system_msg_style = ParagraphStyle(name='SystemMessage', parent=styles['Italic'], spaceBefore=3, spaceAfter=3, leftIndent=0.1*inch, backColor=colors.beige, fontSize=9)


    # --- Extract data from report_data dictionary ---
    patient_name = report_data.get("patient_name", "N/A")
    session_id = report_data.get("session_id", "N/A")
    session_title = report_data.get("session_title", "Untitled Consultation")
    session_start_time_obj = report_data.get("session_start_time")
    patient_context_summary = report_data.get("patient_context_summary", "No specific patient context was provided for this session.")
    
    # Convert message dictionaries to MockChatMessage objects if needed, or use dicts directly
    raw_messages = report_data.get("messages", [])
    messages = [MockChatMessage(**msg_data) for msg_data in raw_messages] # If generate_pdf_report expects objects


    # 1. Logo and Document Header
    logo_path_str = get_logo_path()
    if logo_path_str and Path(logo_path_str).exists():
        try:
            img = Image(logo_path_str, width=0.8*inch, height=0.8*inch, preserveAspectRatio=True)
            img.hAlign = 'LEFT'
            story.append(img)
            story.append(Spacer(1, 0.1*inch))
        except Exception as e:
            app_logger.warning(f"PDF Report: Could not add logo: {e}")
    
    story.append(Paragraph(f"{settings.APP_TITLE}", styles['h1']))
    story.append(Paragraph("AI-Assisted Consultation Summary", styles['h2']))
    story.append(Spacer(1, 0.2*inch))

    # 2. Report Metadata Table
    report_date_str = datetime.now().strftime('%Y-%m-%d %H:%M UTC')
    session_start_time_str = session_start_time_obj.strftime('%Y-%m-%d %H:%M UTC') if session_start_time_obj else "N/A"
    
    meta_data = [
        [Paragraph("<b>Report Generated:</b>", styles['Normal']), Paragraph(report_date_str, styles['Normal'])],
        [Paragraph("<b>Clinician:</b>", styles['Normal']), Paragraph(patient_name, styles['Normal'])], # "patient_name" is the clinician's username
        [Paragraph("<b>Consultation Session ID:</b>", styles['Normal']), Paragraph(str(session_id), styles['Normal'])],
        [Paragraph("<b>Session Title:</b>", styles['Normal']), Paragraph(session_title, styles['Normal'])],
        [Paragraph("<b>Session Start Time:</b>", styles['Normal']), Paragraph(session_start_time_str, styles['Normal'])],
    ]
    meta_table = Table(meta_data, colWidths=[1.8*inch, None])
    meta_table.setStyle(TableStyle([
        ('GRID', (0,0), (-1,-1), 0.5, colors.grey),
        ('VALIGN', (0,0), (-1,-1), 'TOP'),
        ('BACKGROUND', (0,0), (0,-1), colors.whitesmoke)
    ]))
    story.append(meta_table)
    story.append(Spacer(1, 0.3*inch))

    # 3. Disclaimer
    story.append(Paragraph("<b>Important Disclaimer:</b>", styles['SubHeader']))
    story.append(Paragraph(settings.MAIN_DISCLAIMER_LONG, styles['Disclaimer']))
    story.append(Paragraph(settings.SIMULATION_DISCLAIMER, styles['Disclaimer']))
    story.append(Spacer(1, 0.3*inch))

    # 4. Patient Context Provided (if any)
    if patient_context_summary and patient_context_summary != "No specific patient context provided.":
        story.append(Paragraph("Patient Context Provided (Simulated Data):", styles['SectionHeader']))
        # Assuming patient_context_summary is a string. If it's structured, format it nicely.
        # Example: if it's "Key1: Value1; Key2: Value2", split and list.
        context_parts = patient_context_summary.replace("Patient Context: ", "").split(';')
        for part in context_parts:
            if part.strip():
                 story.append(Paragraph(f"• {part.strip()}", styles['ListItem']))
        story.append(Spacer(1, 0.2*inch))

    # 5. Consultation Transcript
    story.append(Paragraph("Consultation Transcript:", styles['SectionHeader']))
    story.append(Spacer(1, 0.1*inch))

    for msg in messages:
        timestamp_str = msg.timestamp.strftime('%Y-%m-%d %H:%M:%S UTC') if msg.timestamp else "N/A"
        prefix = ""
        current_style = styles['Normal']

        if msg.role == 'assistant':
            prefix = f"AI Assistant ({timestamp_str}):"
            current_style = ai_msg_style
        elif msg.role == 'user':
            prefix = f"Clinician ({timestamp_str}):"
            current_style = user_msg_style
        elif msg.role == 'tool':
            tool_name = getattr(msg, 'tool_name', 'Tool')
            prefix = f"{tool_name.capitalize()} ({timestamp_str}):"
            current_style = tool_msg_style
        elif msg.role == 'system': # Typically system context, might not always be included
            prefix = f"System Context ({timestamp_str}):"
            current_style = system_msg_style
        else: # Fallback for any other roles
            prefix = f"{msg.role.capitalize()} ({timestamp_str}):"
            current_style = styles['Normal']
        
        # Sanitize content for ReportLab Paragraph (handles HTML-like tags for <br/>)
        content_for_pdf = msg.content.replace('\n', '<br/>\n')
        content_for_pdf = content_for_pdf.replace("<", "<").replace(">", ">").replace("<br/>", "<br/>")

        story.append(Paragraph(f"<b>{prefix}</b><br/>{content_for_pdf}", current_style))
        # story.append(Spacer(1, 0.05*inch)) # Removed for tighter message packing

    # Footer on each page (if needed, more complex - requires onLaterPages in SimpleDocTemplate)
    # For simplicity, a final "End of Report"
    story.append(Spacer(1, 0.5*inch))
    story.append(Paragraph("--- End of Report ---", styles['Italic']))

    try:
        doc.build(story)
        buffer.seek(0)
        app_logger.info(f"PDF report generated for session ID: {session_id}")
    except Exception as e:
        app_logger.error(f"Failed to build PDF document for session ID {session_id}: {e}", exc_info=True)
        buffer = BytesIO() # Return empty buffer on error
        # Optionally, create a simple error PDF
        error_styles = getSampleStyleSheet()
        error_doc = SimpleDocTemplate(buffer, pagesize=letter)
        error_story = [Paragraph("Error: Could not generate PDF report.", error_styles['h1']),
                       Paragraph(f"Details: {str(e)[:500]}", error_styles['Normal'])]
        try:
            error_doc.build(error_story)
        except: pass # If even error PDF fails, return empty
        buffer.seek(0)
        
    return buffer