app/DocumentGeneration/chart_generator.py

# Chart Generator Module
# Contains parse_embedded_chart_data() and create_chart_from_data() functions

import re
import io
import os
import sys
import base64
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from typing import List, Dict, Any, Optional

# Add the parent directory to sys.path to import from models.py
sys.path.append(os.path.dirname(os.path.dirname(__file__)))
from models import ExportSection

def parse_embedded_chart_data(sections: List[ExportSection]) -> List[Dict[str, Any]]:
    """Parse chart data embedded in section content using CHARTDATA markers"""
    charts = []
    
    for section in sections:
        if section.content:
            # Look for chart data markers - handle both formats
            chart_matches = re.findall(r'<!--CHART_DATA_START-->(.*?)<!--CHART_DATA_END-->', section.content, re.DOTALL)
            
            # Also check for alternative markers as fallback
            alt_chart_matches = re.findall(r'<!--CHARTDATASTART-->(.*?)<!--CHARTDATAEND-->', section.content, re.DOTALL)
            chart_matches.extend(alt_chart_matches)
            
            # logging.info(f"Section '{section.title}': Found {len(chart_matches)} chart data markers")
            
            for chart_data in chart_matches:
                try:
                    # Parse the chart data array
                    chart_array = eval(chart_data.strip())
                    # logging.info(f"Parsed chart array with {len(chart_array)} items")
                    # logging.info(f"Chart array type: {type(chart_array)}")
                    
                    # Each chart is an array with: [Section, ChartType, Title, Period, Data]
                    for chart_item in chart_array:
                        if len(chart_item) >= 5:
                            section_name = chart_item[0]
                            chart_type = chart_item[1]  # DC=Doughnut, BG=Bar Graph, etc.
                            chart_title = chart_item[2]
                            period = chart_item[3]
                            data = chart_item[4]
                            
                            chart_info = {
                                "section": section_name,
                                "type": chart_type,
                                "title": chart_title,
                                "period": period,
                                "data": data,
                                "source_section": section.title
                            }
                            
                            charts.append(chart_info)
                            # logging.info(f"Added chart: {chart_type} - {chart_title} for section {section_name}")
                        else:
                            # logging.warning(f"Chart item has insufficient data: {len(chart_item)} items, expected 5")
                            pass
                
                except Exception as e:
                    # logging.error(f"Failed to parse chart data for section '{section.title}': {str(e)}")
                    continue
            
            # Additional debugging: Check for any chart-related content
            if not chart_matches:
                # Look for any chart-related markers or text
                chart_indicators = re.findall(r'<!--.*?CHART.*?-->', section.content, re.DOTALL)
                if chart_indicators:
                    # logging.info(f"Section '{section.title}': Found chart indicators: {chart_indicators}")
                    pass
                
                # Check for general chart text
                if 'CHART' in section.content.upper() or 'chart' in section.content.lower():
                    # logging.info(f"Section '{section.title}': Contains chart-related text")
                    pass
    
    # logging.info(f"Total charts parsed: {len(charts)}")
    return charts

def create_enhanced_chart_style(chart_type, data, title, business_idea):
    """Create charts with enhanced professional styling"""
    # Set modern style
    plt.style.use('default')
    
    # Design palette from image
    colors = ['#6A2EAB', '#5A2E9B', '#C0C0C0', '#FF8C00', '#FFD700', '#696969', '#FF4500']
    
    if chart_type == "DC":  # Doughnut Chart
        fig, ax = plt.subplots(figsize=(14, 10))
        
        # Extract labels and values from data
        labels = [item[0] for item in data]
        values = [item[1] for item in data]
        
        # Create doughnut chart with enhanced styling
        wedges, texts, autotexts = ax.pie(values, labels=labels, colors=colors[:len(values)],
                                         autopct='%1.1f%%', startangle=90,
                                         wedgeprops=dict(width=0.6, edgecolor='white', linewidth=2),
                                         textprops={'fontsize': 18})
        
        # Enhanced title styling
        ax.set_title(title, fontsize=24, fontweight='bold', pad=20, color='#374151')
        
        # Add legend with enhanced styling
        legend = ax.legend(wedges, labels, title="Categories", 
                          loc="center left", bbox_to_anchor=(1, 0, 0.5, 1),
                          fontsize=16, title_fontsize=18)
        legend.get_title().set_color('#FF6B6B')
        
        # Add business name watermark
        fig.text(0.5, 0.02, f"Generated for: {business_idea}", 
                ha='center', va='bottom', fontsize=8, color='#9CA3AF', style='italic')
        
        plt.tight_layout()
        return fig
        
    elif chart_type == "BG":  # Bar Graph
        fig, ax = plt.subplots(figsize=(16, 10))
        
        # Extract labels and values
        labels = [item[0] for item in data]
        values = [item[1] for item in data]
        
        # Create bar chart with enhanced styling
        bars = ax.bar(labels, values, color=colors[:len(values)], 
                     edgecolor='white', linewidth=1, alpha=0.8)
        
        # Add value labels on bars
        for bar, value in zip(bars, values):
            height = bar.get_height()
            ax.text(bar.get_x() + bar.get_width()/2., height + max(values)*0.01,
                   f'{value}', ha='center', va='bottom', fontweight='bold', fontsize=18)
        
        # Enhanced styling
        ax.set_title(title, fontsize=24, fontweight='bold', pad=20, color='#374151')
        ax.set_xlabel('Categories', fontsize=18, color='#6B7280')
        ax.set_ylabel('Values', fontsize=18, color='#6B7280')
        
        # Rotate x-axis labels for better readability
        plt.setp(ax.get_xticklabels(), rotation=45, ha='right', fontsize=16)
        plt.setp(ax.get_yticklabels(), fontsize=16)
        
        # Add grid for better readability
        ax.yaxis.grid(True, alpha=0.3)
        ax.set_axisbelow(True)
        
        # Add business name watermark
        fig.text(0.5, 0.02, f"Generated for: {business_idea}", 
                ha='center', va='bottom', fontsize=8, color='#9CA3AF', style='italic')
        
        plt.tight_layout()
        return fig
        
    elif chart_type == "LG":  # Line Graph
        fig, ax = plt.subplots(figsize=(16, 10))
        
        # Extract data
        x_values = [item[0] for item in data]
        y_values = [item[1] for item in data]
        
        # Create line chart with enhanced styling
        ax.plot(x_values, y_values, color='#FF6B6B', linewidth=3, marker='o', 
               markersize=8, markerfacecolor='white', markeredgecolor='#FF6B6B')
        
        # Fill area under line
        ax.fill_between(x_values, y_values, alpha=0.3, color='#4ECDC4')
        
        # Enhanced styling
        ax.set_title(title, fontsize=24, fontweight='bold', pad=20, color='#374151')
        ax.set_xlabel('Time Period', fontsize=18, color='#6B7280')
        ax.set_ylabel('Values', fontsize=18, color='#6B7280')
        
        # Add grid
        ax.grid(True, alpha=0.3)
        ax.set_axisbelow(True)
        
        # Increase tick label font sizes
        plt.setp(ax.get_xticklabels(), fontsize=16)
        plt.setp(ax.get_yticklabels(), fontsize=16)
        
        # Add business name watermark
        fig.text(0.5, 0.02, f"Generated for: {business_idea}", 
                ha='center', va='bottom', fontsize=8, color='#9CA3AF', style='italic')
        
        plt.tight_layout()
        return fig
        
    elif chart_type == "PG":  # Pie Graph
        fig, ax = plt.subplots(figsize=(14, 10))
        
        # Extract labels and values
        labels = [item[0] for item in data]
        values = [item[1] for item in data]
        
        # Create pie chart with enhanced styling
        wedges, texts, autotexts = ax.pie(values, labels=labels, colors=colors[:len(values)],
                                         autopct='%1.1f%%', startangle=90,
                                         wedgeprops=dict(edgecolor='white', linewidth=2),
                                         textprops={'fontsize': 18})
        
        # Enhanced title styling
        ax.set_title(title, fontsize=24, fontweight='bold', pad=20, color='#374151')
        
        # Add business name watermark
        fig.text(0.5, 0.02, f"Generated for: {business_idea}", 
                ha='center', va='bottom', fontsize=8, color='#9CA3AF', style='italic')
        
        plt.tight_layout()
        return fig
    
    return None

def create_chart_from_data(chart_info: Dict[str, Any], business_idea: str) -> bytes:
    """Create a matplotlib chart from parsed chart data with enhanced styling"""
    try:
        chart_type = chart_info["type"]
        chart_title = chart_info["title"]
        data = chart_info["data"]
        
        # logging.info(f"Creating chart: {chart_type} - {chart_title} with {len(data)} data points")
        
        # Use enhanced styling function
        fig = create_enhanced_chart_style(chart_type, data, chart_title, business_idea)
        
        if fig is None:
            # Fallback to original method if enhanced styling fails
            if chart_type == "DC":  # Doughnut Chart
                fig, ax = plt.subplots(figsize=(14, 10))  # Larger size for better visibility
                
                # Extract labels and values from data
                labels = [item[0] for item in data]
                values = [item[1] for item in data]
                
                # High contrast color scheme
                colors = ['#6A2EAB', '#5A2E9B', '#C0C0C0', '#FF8C00', '#FFD700', '#696969', '#FF4500']
                
                # Create doughnut chart
                wedges, texts, autotexts = ax.pie(values, labels=labels, colors=colors[:len(values)],
                                                 autopct='%1.1f%%', startangle=90,
                                                 wedgeprops=dict(width=0.6, edgecolor='white', linewidth=2),
                                                 textprops={'fontsize': 18})
                
                # Enhanced title
                ax.set_title(chart_title, fontsize=22, fontweight='bold', pad=15, color='#374151')
                
                # Add legend
                ax.legend(wedges, labels, title="Categories", loc="center left", bbox_to_anchor=(1, 0, 0.5, 1), fontsize=16)
                
                plt.tight_layout()
            elif chart_type == "BG":  # Bar Graph
                fig, ax = plt.subplots(figsize=(16, 10))
                
                # Extract labels and values
                labels = [item[0] for item in data]
                values = [item[1] for item in data]
                
                # High contrast color scheme
                colors = ['#6A2EAB', '#5A2E9B', '#C0C0C0', '#FF8C00', '#FFD700', '#696969', '#FF4500']
                
                # Create bar chart
                bars = ax.bar(labels, values, color=colors[:len(values)], 
                             edgecolor='white', linewidth=1, alpha=0.8)
                
                # Add value labels on bars
                for bar, value in zip(bars, values):
                    height = bar.get_height()
                    ax.text(bar.get_x() + bar.get_width()/2., height + max(values)*0.01,
                           f'{value}', ha='center', va='bottom', fontweight='bold', fontsize=18)
                
                ax.set_title(chart_title, fontsize=22, fontweight='bold', pad=15, color='#374151')
                ax.set_xlabel('Categories', fontsize=18, color='#6B7280')
                ax.set_ylabel('Values', fontsize=18, color='#6B7280')
                
                # Rotate x-axis labels
                plt.setp(ax.get_xticklabels(), rotation=45, ha='right', fontsize=16)
                plt.setp(ax.get_yticklabels(), fontsize=16)
                
                # Add grid
                ax.yaxis.grid(True, alpha=0.3)
                ax.set_axisbelow(True)
                
                plt.tight_layout()
            elif chart_type == "LG":  # Line Graph
                fig, ax = plt.subplots(figsize=(16, 10))
                
                # Extract data
                x_values = [item[0] for item in data]
                y_values = [item[1] for item in data]
                
                # Create line chart
                ax.plot(x_values, y_values, color='#FF6B6B', linewidth=3, marker='o', 
                       markersize=8, markerfacecolor='white', markeredgecolor='#FF6B6B')
                
                # Fill area under line
                ax.fill_between(x_values, y_values, alpha=0.3, color='#4ECDC4')
                
                ax.set_title(chart_title, fontsize=22, fontweight='bold', pad=15, color='#374151')
                ax.set_xlabel('Time Period', fontsize=18, color='#6B7280')
                ax.set_ylabel('Values', fontsize=18, color='#6B7280')
                
                # Add grid
                ax.grid(True, alpha=0.3)
                ax.set_axisbelow(True)
                
                # Increase tick label font sizes
                plt.setp(ax.get_xticklabels(), fontsize=16)
                plt.setp(ax.get_yticklabels(), fontsize=16)
                
                plt.tight_layout()
            elif chart_type == "PG":  # Pie Graph
                fig, ax = plt.subplots(figsize=(14, 10))
                
                # Extract labels and values
                labels = [item[0] for item in data]
                values = [item[1] for item in data]
                
                # High contrast color scheme
                colors = ['#6A2EAB', '#5A2E9B', '#C0C0C0', '#FF8C00', '#FFD700', '#696969', '#FF4500']
                
                # Create pie chart
                wedges, texts, autotexts = ax.pie(values, labels=labels, colors=colors[:len(values)],
                                                 autopct='%1.1f%%', startangle=90,
                                                 wedgeprops=dict(edgecolor='white', linewidth=2),
                                                 textprops={'fontsize': 18})
                
                ax.set_title(chart_title, fontsize=22, fontweight='bold', pad=15, color='#374151')
                
                plt.tight_layout()
            else:
                # logging.warning(f"Unknown chart type: {chart_type}")
                return None
        
        # Save to bytes
        buffer = io.BytesIO()
        fig.savefig(buffer, format='png', dpi=300, bbox_inches='tight', 
                   facecolor='white', edgecolor='none')
        buffer.seek(0)
        chart_bytes = buffer.getvalue()
        buffer.close()
        
        # Close the figure to free memory
        plt.close(fig)
        
        return chart_bytes
        
    except Exception as e:
        # logging.error(f"Failed to create chart: {str(e)}")
        return None

def fix_base64_padding(base64_string: str) -> str:
    """Fix common base64 padding issues"""
    # Remove any whitespace or newlines
    base64_string = base64_string.strip()
    
    # Add padding if needed
    missing_padding = len(base64_string) % 4
    if missing_padding:
        base64_string += '=' * (4 - missing_padding)
    
    return base64_string

def robust_base64_decode(base64_string: str) -> Optional[bytes]:
    """Robustly decode base64 string with error handling and padding correction"""
    try:
        # First try direct decoding
        return base64.b64decode(base64_string)
    except Exception as e1:
        try:
            # Try with padding correction
            fixed_string = fix_base64_padding(base64_string)
            return base64.b64decode(fixed_string)
        except Exception as e2:
            try:
                # Try with URL-safe base64
                return base64.urlsafe_b64decode(fix_base64_padding(base64_string))
            except Exception as e3:
                try:
                    # Try with standard base64 ignoring padding
                    return base64.b64decode(base64_string + '==', validate=False)
                except Exception as e4:
                    return None