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

# Layout parameters
COLUMNS = 3

# Derived constants
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

# Figure dimensions
FIGURE_WIDTH = 20  # Wider to accommodate columns
MAX_HEIGHT = 12  # Maximum height in inches
MIN_HEIGHT_PER_ROW = 2.2
FIGURE_PADDING = 2

# Bar styling
BAR_HEIGHT_RATIO = 0.22  # Bar height as ratio of vertical spacing
VERTICAL_SPACING_RATIO = 0.2  # Base vertical position ratio
AMD_BAR_OFFSET = 0.25  # AMD bar offset ratio
NVIDIA_BAR_OFFSET = 0.54  # NVIDIA bar offset ratio

# Colors
COLORS = {
    'passed': '#4CAF50',
    'failed': '#E53E3E',
    'skipped': '#FFD54F',
    'error': '#8B0000',
    'empty': "#5B5B5B"
}

# Font styling
MODEL_NAME_FONT_SIZE = 16
LABEL_FONT_SIZE = 14
LABEL_OFFSET = 1  # Distance of label from bar


def draw_text_and_bar(
    label: str,
    stats: dict[str, int],
    y_bar: float,
    column_left_position: float,
    bar_height: float,
    ax: plt.Axes,
) -> None:
    """Draw a horizontal bar chart for given stats and its label on the left."""
    # Text
    label_x = column_left_position - LABEL_OFFSET
    ax.text(
        label_x, y_bar, label, ha='right', va='center', color='#CCCCCC', fontsize=LABEL_FONT_SIZE,
        fontfamily='monospace', fontweight='normal'
    )
    # Bar
    total = sum(stats.values())
    if total > 0:
        left = column_left_position
        for category in ['passed', 'failed', 'skipped', 'error']:
            if stats[category] > 0:
                width = stats[category] / total * BAR_WIDTH
                ax.barh(y_bar, width, left=left, height=bar_height, color=COLORS[category], alpha=0.9)
                left += width
    else:
        ax.barh(y_bar, BAR_WIDTH, left=column_left_position, height=bar_height, color=COLORS['empty'], alpha=0.9)

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)
    rows = (model_count + COLUMNS - 1) // COLUMNS  # Ceiling division

    # Figure dimensions - wider for columns, height based on rows
    height_per_row = min(MIN_HEIGHT_PER_ROW, MAX_HEIGHT / max(rows, 1))
    figure_height = min(MAX_HEIGHT, rows * height_per_row + FIGURE_PADDING)

    fig, ax = plt.subplots(figsize=(FIGURE_WIDTH, figure_height), facecolor='#000000')
    ax.set_facecolor('#000000')

    visible_model_count = 0
    max_y = 0

    for i, model_name in enumerate(available_models):
        if model_name not in df.index:
            continue

        row = df.loc[model_name]

        # Extract and process model data
        amd_stats, nvidia_stats = extract_model_data(row)[:2]

        # 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 = (VERTICAL_SPACING_RATIO + row) * vertical_spacing
        y_model_name = y_base    # Model name above AMD bar
        y_amd_bar = y_base + vertical_spacing * AMD_BAR_OFFSET       # AMD bar
        y_nvidia_bar = y_base + vertical_spacing * NVIDIA_BAR_OFFSET    # 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=MODEL_NAME_FONT_SIZE, fontfamily='monospace', fontweight='bold')

        # AMD label and bar in this column
        bar_height = min(0.4, vertical_spacing * BAR_HEIGHT_RATIO)
        # Draw AMD bar
        draw_text_and_bar("amd", amd_stats, y_amd_bar, col_left, bar_height, ax)
        # Draw NVIDIA bar
        draw_text_and_bar("nvidia", nvidia_stats, y_nvidia_bar, col_left, bar_height, ax)

        # 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()
    return fig