# 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 Dosage"))
#         elements.extend(self._add_chart(student_metrics_fig, "Student Attendance and Engagement"))
#         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("MTSS.ai Analysis", 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


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, TA_LEFT
from reportlab.lib.units import inch
import matplotlib.pyplot as plt
import markdown
from bs4 import BeautifulSoup
from PIL import Image as PILImage
from reportlab.lib.colors import black

class ReportGenerator:
    def __init__(self):
        self.styles = getSampleStyleSheet()
        self.styles['BodyText'].alignment = TA_JUSTIFY
        
        self.styles['Bullet'].leftIndent = 20
        self.styles['Bullet'].firstLineIndent = 0
        self.styles['Bullet'].alignment = TA_LEFT
        
        self.styles['Heading1'].fontSize = 18
        self.styles['Heading2'].fontSize = 16
        self.styles['Heading3'].fontSize = 14
        self.styles['Heading4'].fontSize = 12

        # Add a new style for bold text
        self.styles.add(ParagraphStyle(name='Bold', parent=self.styles['BodyText'], fontName='Helvetica-Bold'))

    def create_combined_pdf(self, intervention_fig, student_metrics_fig, recommendations):
        buffer = io.BytesIO()
        doc = SimpleDocTemplate(buffer, pagesize=letter, topMargin=0.5*inch, bottomMargin=0.5*inch)
        
        elements = []
        
        elements.extend(self._add_chart(intervention_fig, "Intervention Dosage"))
        elements.extend(self._add_chart(student_metrics_fig, "Student Attendance and Engagement"))
        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("MTSS.ai Analysis", self.styles['Heading1']))
        
        html = markdown.markdown(recommendations)
        soup = BeautifulSoup(html, 'html.parser')
        
        for element in soup.find_all(['h1', 'h2', 'h3', 'h4', 'h5', 'h6', 'p', 'ul']):
            if element.name.startswith('h'):
                level = int(element.name[1])
                style = f'Heading{min(level, 4)}'
                elements.append(Paragraph(element.text, self.styles[style]))
            elif element.name == 'p':
                elements.append(self._create_paragraph_with_inline_styles(element))
            elif element.name == 'ul':
                for li in element.find_all('li'):
                    bullet_text = '• ' + self._get_text_with_inline_styles(li)
                    elements.append(Paragraph(bullet_text, self.styles['Bullet']))
            elements.append(Spacer(1, 6))
        
        return elements

    def _create_paragraph_with_inline_styles(self, element):
        text = self._get_text_with_inline_styles(element)
        return Paragraph(text, self.styles['BodyText'])

    def _get_text_with_inline_styles(self, element):
        text = ""
        for child in element.children:
            if isinstance(child, str):
                text += child
            elif child.name in ['strong', 'b']:
                text += f'<b>{child.text}</b>'
            elif child.name in ['em', 'i']:
                text += f'<i>{child.text}</i>'
            elif child.name == 'u':
                text += f'<u>{child.text}</u>'
        return text