import matplotlib.pyplot as plt
import pandas as pd
from utils import generate_underlined_line
from data import extract_model_data

# Figure dimensions
FIGURE_WIDTH_DUAL = 18
FIGURE_HEIGHT_DUAL = 9

# Colors
COLORS = {
    'passed': '#4CAF50',    # Medium green
    'failed': '#E53E3E',    # More red
    'skipped': '#FFD54F',   # Medium yellow
    'error': '#8B0000'      # Dark red
}

# Styling constants
BLACK = '#000000'
LABEL_COLOR = '#AAAAAA'
TITLE_COLOR = '#FFFFFF'

# Font sizes
DEVICE_TITLE_FONT_SIZE = 28

# Layout constants
SEPARATOR_LINE_Y_END = 0.85
SUBPLOT_TOP = 0.85
SUBPLOT_WSPACE = 0.4
PIE_START_ANGLE = 90
BORDER_LINE_WIDTH = 0.5
SEPARATOR_ALPHA = 0.5
SEPARATOR_LINE_WIDTH = 1
DEVICE_TITLE_PAD = 2
MODEL_TITLE_Y = 1

# Processing constants
MAX_FAILURE_ITEMS = 10


def _process_failure_category(failures_obj: dict, category: str, info_lines: list) -> None:
    """Process a single failure category (multi or single) and add to info_lines."""
    if category in failures_obj and failures_obj[category]:
        info_lines.append(generate_underlined_line(f"{category.title()} GPU failure details:"))
        if isinstance(failures_obj[category], list):
            # Handle list of failures (could be strings or dicts)
            for i, failure in enumerate(failures_obj[category][:MAX_FAILURE_ITEMS]):
                if isinstance(failure, dict):
                    # Extract meaningful info from dict (e.g., test name, line, etc.)
                    failure_str = failure.get('line', failure.get('test', 
                                            failure.get('name', str(failure))))
                    info_lines.append(f"  {i+1}. {failure_str}")
                else:
                    info_lines.append(f"  {i+1}. {str(failure)}")
            if len(failures_obj[category]) > MAX_FAILURE_ITEMS:
                remaining = len(failures_obj[category]) - MAX_FAILURE_ITEMS
                info_lines.append(f"... and {remaining} more")
        else:
            info_lines.append(str(failures_obj[category]))
        info_lines.append("")


def extract_failure_info(failures_obj, device: str, multi_count: int, single_count: int) -> str:
    """Extract failure information from failures object."""
    if (not failures_obj or pd.isna(failures_obj)) and multi_count == 0 and single_count == 0:
        return f"No failures on {device}"
    
    info_lines = []
    
    # Add counts summary
    if multi_count > 0 or single_count > 0:
        info_lines.append(generate_underlined_line(f"Failure Summary for {device}:"))
        if multi_count > 0:
            info_lines.append(f"Multi GPU failures: {multi_count}")
        if single_count > 0:
            info_lines.append(f"Single GPU failures: {single_count}")
        info_lines.append("")
    
    # Try to extract detailed failure information
    try:
        if isinstance(failures_obj, dict):
            _process_failure_category(failures_obj, 'multi', info_lines)
            _process_failure_category(failures_obj, 'single', info_lines)
        
        return "\n".join(info_lines) if info_lines else f"No detailed failure info for {device}"
    
    except Exception as e:
        if multi_count > 0 or single_count > 0:
            error_msg = (f"Failures detected on {device} (Multi: {multi_count}, Single: {single_count})\n"
                        f"Details unavailable: {str(e)}")
            return error_msg
        return f"Error processing failure info for {device}: {str(e)}"


def _create_pie_chart(ax: plt.Axes, device_label: str, filtered_stats: dict) -> None:
    """Create a pie chart for device statistics."""
    if not filtered_stats:
        ax.text(0.5, 0.5, 'No test results', 
               horizontalalignment='center', verticalalignment='center',
               transform=ax.transAxes, fontsize=14, color='#888888',
               fontfamily='monospace', weight='normal')
        ax.set_title(device_label, fontsize=DEVICE_TITLE_FONT_SIZE, weight='bold', 
                    pad=DEVICE_TITLE_PAD, color=TITLE_COLOR, fontfamily='monospace')
        ax.axis('off')
        return
        
    chart_colors = [COLORS[category] for category in filtered_stats.keys()]
    
    # Create minimal pie chart - full pie, no donut effect
    wedges, texts, autotexts = ax.pie(
        filtered_stats.values(), 
        labels=[label.lower() for label in filtered_stats.keys()],  # Lowercase for minimal look
        colors=chart_colors,
        autopct=lambda pct: f'{int(pct/100*sum(filtered_stats.values()))}',
        startangle=PIE_START_ANGLE,
        explode=None,  # No separation
        shadow=False,
        wedgeprops=dict(edgecolor='#1a1a1a', linewidth=BORDER_LINE_WIDTH),  # Minimal borders
        textprops={'fontsize': 12, 'weight': 'normal', 
                  'color': LABEL_COLOR, 'fontfamily': 'monospace'}
    )
    
    # Enhanced percentage text styling for better readability
    for autotext in autotexts:
        autotext.set_color(BLACK)  # Black text for better contrast
        autotext.set_weight('bold')
        autotext.set_fontsize(14)
        autotext.set_fontfamily('monospace')
    
    # Minimal category labels
    for text in texts:
        text.set_color(LABEL_COLOR)
        text.set_weight('normal')
        text.set_fontsize(13)
        text.set_fontfamily('monospace')
    
    # Device label closer to chart and bigger
    ax.set_title(device_label, fontsize=DEVICE_TITLE_FONT_SIZE, weight='normal', 
                pad=DEVICE_TITLE_PAD, color=TITLE_COLOR, fontfamily='monospace')


def plot_model_stats(df: pd.DataFrame, model_name: str) -> tuple[plt.Figure, str, str]:
    """Draws pie charts of model's passed, failed, skipped, and error stats for AMD and NVIDIA."""
    # Handle case where the dataframe is empty or the model name could not be found in it
    if df.empty or model_name not in df.index:
        # Create empty stats for both devices
        amd_filtered = {}
        nvidia_filtered = {}
        failed_multi_amd = failed_single_amd = failed_multi_nvidia = failed_single_nvidia = 0
        failures_amd = failures_nvidia = {}
    else:
        row = df.loc[model_name]
        
        # Extract and process model data
        amd_stats, nvidia_stats, failed_multi_amd, failed_single_amd, failed_multi_nvidia, failed_single_nvidia = \
            extract_model_data(row)
        
        # Filter out categories with 0 values for cleaner visualization
        amd_filtered = {k: v for k, v in amd_stats.items() if v > 0}
        nvidia_filtered = {k: v for k, v in nvidia_stats.items() if v > 0}
        
        # Generate failure info directly from dataframe
        failures_amd = row.get('failures_amd', {})
        failures_nvidia = row.get('failures_nvidia', {})
    
    # Always create figure with two subplots side by side with padding
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(FIGURE_WIDTH_DUAL, FIGURE_HEIGHT_DUAL), facecolor=BLACK)
    ax1.set_facecolor(BLACK)
    ax2.set_facecolor(BLACK)
    
    # Create both pie charts with device labels
    _create_pie_chart(ax1, "amd", amd_filtered)
    _create_pie_chart(ax2, "nvidia", nvidia_filtered)
    
    # Add subtle separation line between charts - stops at device labels level
    line_x = 0.5
    fig.add_artist(plt.Line2D([line_x, line_x], [0.0, SEPARATOR_LINE_Y_END], 
                              color='#333333', linewidth=SEPARATOR_LINE_WIDTH, 
                              alpha=SEPARATOR_ALPHA, transform=fig.transFigure))
    
    # Add central shared title for model name
    fig.suptitle(f'{model_name.lower()}', fontsize=32, weight='bold', 
                color='#CCCCCC', fontfamily='monospace', y=MODEL_TITLE_Y)
    
    # Clean layout with padding and space for central title
    plt.tight_layout()
    plt.subplots_adjust(top=SUBPLOT_TOP, wspace=SUBPLOT_WSPACE)
    
    amd_failed_info = extract_failure_info(failures_amd, 'AMD', failed_multi_amd, failed_single_amd)
    nvidia_failed_info = extract_failure_info(failures_nvidia, 'NVIDIA', failed_multi_nvidia, failed_single_nvidia)
    
    return fig, amd_failed_info, nvidia_failed_info