src/viz_report.py

import streamlit as st
import pandas as pd
from fpdf import FPDF
import matplotlib.pyplot as plt
import seaborn as sns
import io
from datetime import datetime
import tempfile # To handle temporary files
import os       # To interact with the operating system

class ComprehensivePDF(FPDF):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.user_name = ""

    def header(self):
        self.set_font('Arial', 'B', 12)
        self.cell(0, 10, 'Visio AI - Comprehensive Analysis Report', 0, 1, 'C')
        self.ln(5)

    def footer(self):
        self.set_y(-15)
        self.set_font('Arial', 'I', 8)
        self.cell(0, 10, f'Page {self.page_no()}', 0, 0, 'C')

    def add_title_page(self, user_name=""):
        self.user_name = user_name
        self.add_page()
        self.set_font('Arial', 'B', 24)
        self.cell(0, 20, 'Comprehensive Analysis Report', 0, 1, 'C')
        self.ln(20)
        
        self.set_font('Arial', '', 12)
        if self.user_name:
            self.cell(0, 10, f"Prepared for: {self.user_name}", 0, 1, 'C')
        self.cell(0, 10, f"Date Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}", 0, 1, 'C')
        self.ln(20)
        
        try:
            self.image("images/favicon.png", x=85, y=100, w=40)
        except FileNotFoundError:
            self.set_font('Arial', 'I', 10)
            self.cell(0, 10, "[Logo Not Found]", 0, 1, 'C')

        self.set_y(-40)
        self.set_font('Arial', 'I', 10)
        self.cell(0, 10, "Generated by Visio AI", 0, 1, 'C')
        
    def add_section_title(self, title):
        self.add_page()
        self.set_font('Arial', 'B', 16)
        self.cell(0, 10, title, 0, 1, 'L')
        self.ln(5)
        
    def add_text_content(self, title, content):
        self.set_font('Arial', 'B', 12)
        self.cell(0, 10, title, 0, 1, 'L')
        self.set_font('Courier', '', 10)
        self.multi_cell(0, 5, content)
        self.ln(5)

    # --- NEW BULLETPROOF HELPER FUNCTION ---
    def add_image_from_object(self, image_object, width):
        """
        Saves a matplotlib figure or PIL image to a temporary file
        and adds it to the PDF, then deletes the file.
        """
        fp = None
        try:
            # Create a named temporary file
            with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as fp:
                # Save the image object to the temporary file
                image_object.save(fp, format="PNG")
                temp_path = fp.name
            
            # Add the image to the PDF from the temporary file path
            self.image(temp_path, w=width)
        finally:
            # Ensure the temporary file is deleted
            if fp and os.path.exists(fp.name):
                os.remove(fp.name)

def render_report_page():
    st.markdown("<h2 style='text-align: center; color: #4A90E2;'>📄 Comprehensive Report Generator</h2>", unsafe_allow_html=True)
    st.markdown("<p style='text-align: center;'>Generate a complete PDF report of your session's activities.</p>", unsafe_allow_html=True)
    
    user_name = st.text_input("Enter your name (optional, will be shown on the report cover)")

    if st.button("Generate Full Report 🚀"):
        if 'updated_df' not in st.session_state and 'viz_ai_img_result' not in st.session_state and 'word_cloud_result' not in st.session_state:
            st.warning("There is no activity to report. Please train a model or use the AI tools first.")
            return

        with st.spinner("Assembling your comprehensive report..."):
            pdf = ComprehensivePDF()
            pdf.add_title_page(user_name)

            # --- Section 1: Data Analysis & ML ---
            if 'updated_df' in st.session_state and st.session_state.updated_df is not None:
                df = st.session_state.updated_df
                pdf.add_section_title("1. Dataset & Machine Learning Analysis")
                
                buffer = io.StringIO()
                df.info(buf=buffer)
                pdf.add_text_content("Dataset Information", buffer.getvalue())
                
                pdf.add_text_content("Numerical Summary", df.describe(include='number').to_string())
                if not df.select_dtypes(include='object').empty:
                    pdf.add_text_content("Categorical Summary", df.describe(include='object').to_string())

                if 'trained_model' in st.session_state and st.session_state.trained_model is not None:
                    metrics = st.session_state.model_metrics
                    algo_name = st.session_state.selected_algo_name
                    
                    pdf.set_font('Arial', 'B', 12)
                    pdf.cell(0, 10, f"Machine Learning Model: {algo_name}", 0, 1, 'L')
                    
                    metrics_str = ""
                    for key, val in metrics.items():
                        if key != 'Confusion Matrix':
                            metrics_str += f"{key}: {val:.4f}\n" if isinstance(val, float) else f"{key}: {val}\n"
                    pdf.add_text_content("Performance Metrics", metrics_str)
                    
                    if 'Confusion Matrix' in metrics:
                        fig_cm, ax_cm = plt.subplots()
                        sns.heatmap(metrics['Confusion Matrix'], annot=True, fmt='d', cmap='Blues', ax=ax_cm)
                        ax_cm.set_title('Confusion Matrix')
                        pdf.set_font('Arial', 'B', 12)
                        pdf.cell(0, 10, "Confusion Matrix", 0, 1, 'L')
                        pdf.add_image_from_object(fig_cm, width=170)
                        plt.close(fig_cm) # Close the figure to free memory

            # --- Section 2: Viz AI Image Analysis ---
            if 'viz_ai_img_result' in st.session_state and st.session_state.viz_ai_img_result is not None:
                img_result = st.session_state.viz_ai_img_result
                pdf.add_section_title("2. Viz AI Image Analysis")
                
                pdf.add_image_from_object(img_result['image'], width=150)
                pdf.ln(5)
                
                pdf.add_text_content("Model Used", img_result['model'])
                pdf.add_text_content("User Prompt", img_result['prompt'])
                pdf.add_text_content("AI Analysis", img_result['analysis'])
            
            # --- Section 3: Word Cloud ---
            if 'word_cloud_result' in st.session_state and st.session_state.word_cloud_result is not None:
                wc_result = st.session_state.word_cloud_result
                pdf.add_section_title("3. Word Cloud Analysis")
                pdf.add_text_content("Source File", wc_result['source'])
                pdf.add_text_content("Settings", wc_result['settings'])
                pdf.set_font('Arial', 'B', 12)
                pdf.cell(0, 10, "Generated Word Cloud", 0, 1, 'L')
                pdf.add_image_from_object(wc_result['figure'], width=170)
                plt.close(wc_result['figure']) # Close the figure to free memory

            # --- Generate Download ---
            pdf_output = pdf.output()
            st.success("Report Generated!")
            st.download_button(
                label="📥 Download Full Report",
                data=pdf_output,
                file_name=f"VisioAI_Comprehensive_Report_{datetime.now().strftime('%Y%m%d')}.pdf",
                mime="application/pdf"
            )