# import io
# from reportlab.lib.pagesizes import letter
# from reportlab.platypus import SimpleDocTemplate, Image, Paragraph, Spacer
# from reportlab.lib.styles import getSampleStyleSheet
# import matplotlib.pyplot as plt

# class ReportGenerator:
#     def __init__(self):
#         self.styles = getSampleStyleSheet()

#     def create_combined_pdf(self, intervention_fig, student_metrics_fig, recommendations):
#         buffer = io.BytesIO()
#         doc = SimpleDocTemplate(buffer, pagesize=letter)
        
#         elements = []
        
#         # Add the intervention statistics chart
#         elements.extend(self._add_chart(intervention_fig, "Intervention Statistics"))
        
#         # Add the student metrics chart
#         elements.extend(self._add_chart(student_metrics_fig, "Student Metrics"))
        
#         # Add the AI recommendations
#         elements.extend(self._add_recommendations(recommendations))
        
#         # Build the PDF
#         doc.build(elements)
#         buffer.seek(0)
#         return buffer

#     def _add_chart(self, fig, title):
#         elements = []
#         elements.append(Paragraph(title, self.styles['Heading2']))
#         img_buffer = io.BytesIO()
        
#         if hasattr(fig, 'write_image'):  # Plotly figure
#             fig.write_image(img_buffer, format="png")
#         elif isinstance(fig, plt.Figure):  # Matplotlib figure
#             fig.savefig(img_buffer, format='png')
#             plt.close(fig)  # Close the figure to free up memory
#         else:
#             raise ValueError(f"Unsupported figure type: {type(fig)}")
        
#         img_buffer.seek(0)
#         elements.append(Image(img_buffer, width=500, height=300))
#         elements.append(Spacer(1, 12))
#         return elements

#     def _add_recommendations(self, recommendations):
#         elements = []
#         elements.append(Paragraph("AI Recommendations", self.styles['Heading1']))
#         elements.append(Paragraph(recommendations, self.styles['BodyText']))
#         return elements


import io
from reportlab.lib.pagesizes import letter
from reportlab.platypus import SimpleDocTemplate, Image, Paragraph, Spacer
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib.enums import TA_JUSTIFY
from reportlab.lib.units import inch
import matplotlib.pyplot as plt
import markdown
from xml.etree import ElementTree as ET
from PIL import Image as PILImage
from xml.parsers.expat import ExpatError
from html import escape

class ReportGenerator:
    def __init__(self):
        self.styles = getSampleStyleSheet()
        self.styles.add(ParagraphStyle(name='Justify', alignment=TA_JUSTIFY))

    def create_combined_pdf(self, intervention_fig, student_metrics_fig, recommendations):
        buffer = io.BytesIO()
        doc = SimpleDocTemplate(buffer, pagesize=letter)
        
        elements = []
        
        elements.extend(self._add_chart(intervention_fig, "Intervention Statistics"))
        elements.extend(self._add_chart(student_metrics_fig, "Student Metrics"))
        elements.extend(self._add_recommendations(recommendations))
        
        doc.build(elements)
        buffer.seek(0)
        return buffer

    def _add_chart(self, fig, title):
        elements = []
        elements.append(Paragraph(title, self.styles['Heading2']))
        img_buffer = io.BytesIO()
        
        if hasattr(fig, 'write_image'):  # Plotly figure
            fig.write_image(img_buffer, format="png", width=700, height=400)
        elif isinstance(fig, plt.Figure):  # Matplotlib figure
            fig.set_size_inches(10, 6)  # Set a consistent size
            fig.savefig(img_buffer, format='png', dpi=100, bbox_inches='tight')
            plt.close(fig)
        else:
            raise ValueError(f"Unsupported figure type: {type(fig)}")
        
        img_buffer.seek(0)
        
        # Use PIL to get image dimensions
        with PILImage.open(img_buffer) as img:
            img_width, img_height = img.size
        
        # Calculate width and height to maintain aspect ratio
        max_width = 6.5 * inch  # Maximum width (letter width is 8.5 inches, leaving margins)
        max_height = 4 * inch   # Maximum height
        
        aspect = img_width / float(img_height)
        
        if img_width > max_width:
            img_width = max_width
            img_height = img_width / aspect
        
        if img_height > max_height:
            img_height = max_height
            img_width = img_height * aspect
        
        # Reset buffer position
        img_buffer.seek(0)
        
        # Create ReportLab Image with calculated dimensions
        img = Image(img_buffer, width=img_width, height=img_height)
        
        elements.append(img)
        elements.append(Spacer(1, 12))
        return elements

    def _add_recommendations(self, recommendations):
        elements = []
        elements.append(Paragraph("AI Recommendations", self.styles['Heading1']))
        
        # Convert markdown to HTML
        html = markdown.markdown(recommendations)
        
        # Wrap the HTML in a root element to ensure valid XML
        wrapped_html = f"<root>{html}</root>"
        
        try:
            root = ET.fromstring(wrapped_html)
        except ExpatError:
            # If parsing fails, fallback to treating the entire content as plain text
            elements.append(Paragraph(escape(recommendations), self.styles['BodyText']))
            return elements
        
        for elem in root:
            if elem.tag == 'h3':
                elements.append(Paragraph(elem.text or "", self.styles['Heading3']))
            elif elem.tag == 'h4':
                elements.append(Paragraph(elem.text or "", self.styles['Heading4']))
            elif elem.tag == 'p':
                text = ''.join(elem.itertext())
                elements.append(Paragraph(text, self.styles['Justify']))
            elif elem.tag == 'ul':
                for li in elem.findall('li'):
                    bullet_text = '• ' + ''.join(li.itertext()).strip()
                    elements.append(Paragraph(bullet_text, self.styles['BodyText']))
            else:
                # For any other tags, just extract the text
                text = ''.join(elem.itertext())
                if text.strip():
                    elements.append(Paragraph(text, self.styles['BodyText']))
        
        return elements