import matplotlib.pyplot as plt
import pandas as pd

def create_summary_page(df: pd.DataFrame, available_models: list[str]) -> plt.Figure:
    """Create a summary page with model names and both AMD/NVIDIA test stats bars."""
    if df.empty:
        fig, ax = plt.subplots(figsize=(16, 8), facecolor='#000000')
        ax.set_facecolor('#000000')
        ax.text(0.5, 0.5, 'No data available', 
                horizontalalignment='center', verticalalignment='center',
                transform=ax.transAxes, fontsize=20, color='#888888',
                fontfamily='monospace', weight='normal')
        ax.axis('off')
        return fig
    
    # Calculate dimensions for N-column layout
    model_count = len(available_models)
    columns = 3
    rows = (model_count + columns - 1) // columns  # Ceiling division
    
    # Figure dimensions - wider for 4 columns, height based on rows
    figure_width = 20  # Wider to accommodate 4 columns
    max_height = 12  # Maximum height in inches
    height_per_row = min(2.2, max_height / max(rows, 1))
    figure_height = min(max_height, rows * height_per_row + 2)
    
    fig, ax = plt.subplots(figsize=(figure_width, figure_height), facecolor='#000000')
    ax.set_facecolor('#000000')
    
    colors = {
        'passed': '#4CAF50',
        'failed': '#E53E3E', 
        'skipped': '#FFD54F',
        'error': '#8B0000',
        'empty': "#5B5B5B"
    }
    
    visible_model_count = 0
    max_y = 0
    
    # Column layout parameters
    column_width = 100 / columns  # Each column takes 25% of width
    bar_width = column_width * 0.8  # 80% of column width for bars
    bar_margin = column_width * 0.1  # 10% margin on each side
    
    for i, model_name in enumerate(available_models):
        if model_name not in df.index:
            continue
            
        row = df.loc[model_name]
        
        # Get values directly from dataframe
        success_amd = int(row.get('success_amd', 0)) if pd.notna(row.get('success_amd', 0)) else 0
        success_nvidia = int(row.get('success_nvidia', 0)) if pd.notna(row.get('success_nvidia', 0)) else 0
        failed_multi_amd = int(row.get('failed_multi_no_amd', 0)) if pd.notna(row.get('failed_multi_no_amd', 0)) else 0
        failed_multi_nvidia = int(row.get('failed_multi_no_nvidia', 0)) if pd.notna(row.get('failed_multi_no_nvidia', 0)) else 0
        failed_single_amd = int(row.get('failed_single_no_amd', 0)) if pd.notna(row.get('failed_single_no_amd', 0)) else 0
        failed_single_nvidia = int(row.get('failed_single_no_nvidia', 0)) if pd.notna(row.get('failed_single_no_nvidia', 0)) else 0
        
        # Calculate stats
        amd_stats = {
            'passed': success_amd,
            'failed': failed_multi_amd + failed_single_amd,
            'skipped': 0,
            'error': 0
        }
        
        nvidia_stats = {
            'passed': success_nvidia,
            'failed': failed_multi_nvidia + failed_single_nvidia,
            'skipped': 0,
            'error': 0
        }
        
        amd_total = sum(amd_stats.values())
        nvidia_total = sum(nvidia_stats.values())
        
        if amd_total == 0 and nvidia_total == 0:
            continue
            
        # Calculate position in 4-column grid
        col = visible_model_count % columns
        row = visible_model_count // columns
        
        # Calculate horizontal position for this column
        col_left = col * column_width + bar_margin
        col_center = col * column_width + column_width / 2
        
        # Calculate vertical position for this row - start from top
        vertical_spacing = height_per_row
        y_base = (0.2 + row) * vertical_spacing  # Start closer to top
        y_model_name = y_base    # Model name above AMD bar
        y_amd_bar = y_base + vertical_spacing * 0.25       # AMD bar
        y_nvidia_bar = y_base + vertical_spacing * 0.54    # NVIDIA bar
        max_y = max(max_y, y_nvidia_bar + vertical_spacing * 0.3)
        
        # Model name centered above the bars in this column
        ax.text(col_center, y_model_name, model_name.lower(), 
               ha='center', va='center', color='#FFFFFF', 
               fontsize=16, fontfamily='monospace', fontweight='bold')
        
        # AMD label and bar in this column
        bar_height = min(0.4, vertical_spacing * 0.22)  # Adjust bar height based on spacing
        label_x = col_left - 1  # Label position to the left of the bar
        ax.text(label_x, y_amd_bar, "amd", ha='right', va='center', color='#CCCCCC', fontsize=14, fontfamily='monospace', fontweight='normal')
        
        if amd_total > 0:           
            # AMD bar starts at column left position
            left = col_left
            for category in ['passed', 'failed', 'skipped', 'error']:
                if amd_stats[category] > 0:
                    width = amd_stats[category] / amd_total * bar_width
                    ax.barh(y_amd_bar, width, left=left, height=bar_height, 
                           color=colors[category], alpha=0.9)
                    # if width > 2:  # Smaller threshold for text display
                    #     ax.text(left + width/2, y_amd_bar, str(amd_stats[category]), 
                    #            ha='center', va='center', color='black', 
                    #            fontweight='bold', fontsize=10, fontfamily='monospace')
                    left += width
        else:
            ax.barh(y_amd_bar, bar_width, left=col_left, height=bar_height, color=colors['empty'], alpha=0.9)
            # ax.text(col_center, y_amd_bar, "No data", ha='center', va='center', color='black', fontweight='bold', fontsize=10, fontfamily='monospace')
        
        # NVIDIA label and bar in this column
        ax.text(label_x, y_nvidia_bar, "nvidia", ha='right', va='center', color='#CCCCCC', fontsize=14, fontfamily='monospace', fontweight='normal')

        if nvidia_total > 0:            
            # NVIDIA bar starts at column left position
            left = col_left
            for category in ['passed', 'failed', 'skipped', 'error']:
                if nvidia_stats[category] > 0:
                    width = nvidia_stats[category] / nvidia_total * bar_width
                    ax.barh(y_nvidia_bar, width, left=left, height=bar_height, 
                           color=colors[category], alpha=0.9)
                    # if width > 2:  # Smaller threshold for text display
                    #     ax.text(left + width/2, y_nvidia_bar, str(nvidia_stats[category]), 
                    #            ha='center', va='center', color='black', 
                    #            fontweight='bold', fontsize=10, fontfamily='monospace')
                    left += width
        else:
            ax.barh(y_nvidia_bar, bar_width, left=col_left, height=bar_height, color=colors['empty'], alpha=0.9)
            # ax.text(col_center, y_nvidia_bar, "No data", ha='center', va='center', color='black', fontweight='bold', fontsize=10, fontfamily='monospace')
        
        # Increment counter for next visible model
        visible_model_count += 1
    
    # Style the axes to be completely invisible and span full width
    ax.set_xlim(-5, 105)  # Slightly wider to accommodate labels
    ax.set_ylim(0, max_y)
    ax.set_xlabel('')
    ax.set_ylabel('')
    ax.spines['bottom'].set_visible(False)
    ax.spines['left'].set_visible(False)
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.set_xticks([])
    ax.set_yticks([])
    ax.yaxis.set_inverted(True)
    
    # Remove all margins to make figure stick to top
    plt.tight_layout()
    plt.subplots_adjust(left=0.02, right=0.98, top=1.0, bottom=0.02)
    
    return fig