# 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
import matplotlib.pyplot as plt
import markdown
from xml.etree import ElementTree as ET

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)
        elements.append(Image(img_buffer, width=500, height=300, keepAspectRatio=True))
        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)
        
        # Parse HTML and convert to ReportLab paragraphs
        root = ET.fromstring(html)
        for elem in root.iter():
            if elem.tag == 'h3':
                elements.append(Paragraph(elem.text, self.styles['Heading3']))
            elif elem.tag == 'h4':
                elements.append(Paragraph(elem.text, self.styles['Heading4']))
            elif elem.tag == 'p':
                elements.append(Paragraph(ET.tostring(elem, encoding='unicode', method='text'), self.styles['Justify']))
            elif elem.tag == 'ul':
                for li in elem.findall('li'):
                    bullet_text = '• ' + ET.tostring(li, encoding='unicode', method='text').strip()
                    elements.append(Paragraph(bullet_text, self.styles['BodyText']))
        
        return elements