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