from pathlib import Path
from typing import Union
import pandas as pd
from PIL import Image, ImageDraw, ImageFont


def render(
    df: pd.DataFrame,
    output_path: Union[str, Path],
    font_size: int = 14,
    header_bg_color: tuple = (70, 130, 180),
    header_text_color: tuple = (255, 255, 255),
    cell_bg_color: tuple = (255, 255, 255),
    cell_text_color: tuple = (0, 0, 0),
    margin: int = 10,
):
    """Render a pandas DataFrame as a PNG image."""
    # Calculate dimensions
    num_rows, num_cols = df.shape

    # Get column widths based on content
    col_widths = []
    for col in df.columns:
        max_width = len(str(col))
        for val in df[col]:
            max_width = max(max_width, len(str(val)))
        col_widths.append(max_width)

    # Get row height
    row_height = font_size + 10

    # Calculate image dimensions
    cell_padding = 10
    header_height = row_height
    total_width = sum(col_widths) + (num_cols + 1) * cell_padding
    total_height = (num_rows + 1) * row_height + (num_rows + 2) * margin

    # Create image
    img = Image.new('RGB', (total_width, total_height), color=(255, 255, 255))
    draw = ImageDraw.Draw(img)

    # Try to load a font, fall back to default if not available
    try:
        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", font_size)
    except IOError:
        font = ImageFont.load_default()

    # Draw headers
    x_offset = margin + cell_padding
    y_offset = margin

    for i, col in enumerate(df.columns):
        # Draw header background
        header_left = x_offset + sum(col_widths[:i])
        header_top = y_offset
        header_right = header_left + col_widths[i]
        header_bottom = header_top + header_height
        draw.rectangle(
            [header_left, header_top, header_right, header_bottom],
            fill=header_bg_color
        )

        # Draw header text
        text_x = header_left + cell_padding // 2
        text_y = header_top + (header_height - font_size) // 2
        draw.text((text_x, text_y), str(col), fill=header_text_color, font=font)

    # Draw cells
    for row_idx, (_, row) in enumerate(df.iterrows()):
        y_offset += row_height
        for col_idx, value in enumerate(row):
            cell_left = margin + cell_padding + sum(col_widths[:col_idx])
            cell_top = y_offset
            cell_right = cell_left + col_widths[col_idx]
            cell_bottom = cell_top + row_height

            # Draw cell background (alternating colors for readability)
            if row_idx % 2 == 0:
                cell_color = cell_bg_color
            else:
                cell_color = (245, 245, 245)
            draw.rectangle(
                [cell_left, cell_top, cell_right, cell_bottom],
                fill=cell_color
            )

            # Draw cell border
            draw.rectangle(
                [cell_left, cell_top, cell_right, cell_bottom],
                outline=(200, 200, 200),
                width=1
            )

            # Draw cell text
            text_x = cell_left + cell_padding // 2
            text_y = cell_top + (row_height - font_size) // 2
            draw.text((text_x, text_y), str(value), fill=cell_text_color, font=font)

    # Save image
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    img.save(output_path, 'PNG')
