from io import BytesIO
from textwrap import wrap

import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.cm import ScalarMappable


def plot_security_report(table):

    data = pd.DataFrame(table)

    # Sorting by failureRate for a meaningful arrangement
    data_sorted = data.sort_values("failureRate", ascending=False)

    # Values for the plot
    angles = np.linspace(0, 2 * np.pi, len(data_sorted), endpoint=False)
    failure_rate = data_sorted["failureRate"]
    tokens = data_sorted["tokens"]

    # Styling parameters
    COLORS = ["#6C5B7B", "#C06C84", "#F67280", "#F8B195"]
    cmap = mpl.colors.LinearSegmentedColormap.from_list("custom", COLORS, N=256)
    norm = mpl.colors.Normalize(vmin=tokens.min(), vmax=tokens.max())

    # Polar plot setup
    fig, ax = plt.subplots(figsize=(10, 8), subplot_kw={"projection": "polar"})
    ax.set_theta_offset(np.pi / 2)
    ax.set_theta_direction(-1)
    ax.set_facecolor("white")
    # Bars for failureRate with colors based on 'tokens'
    bars = ax.bar(
        angles,
        failure_rate,
        width=0.3,
        color=[cmap(norm(t)) for t in tokens],
        alpha=0.75,
        label="Failure Rate %",
    )

    # Add labels for the modules
    module_labels = ["\n".join(wrap(m, 10)) for m in data_sorted["module"]]
    ax.set_xticks(angles)

    # Add dashed vertical lines. These are just references

    ax.set_xticklabels(module_labels, fontsize=7, color="#333")

    # Color bar for the tokens
    sm = ScalarMappable(cmap=cmap, norm=norm)
    sm.set_array([])
    cbar = plt.colorbar(sm, ax=ax, orientation="horizontal", pad=0.1)
    cbar.set_label("Token Count (k)", fontsize=12, color="#444")

    # Grid and legend
    ax.grid(True, color="gray", linestyle=":", linewidth=0.5)
    plt.legend(loc="upper right", bbox_to_anchor=(1.1, 1.1))
    ax.vlines(angles, 0, 100, color="#444", ls=(0, (4, 4)), zorder=11)

    # Title and subtitle
    title = "Security Report for Different Modules"
    # fig.suptitle(title, fontsize=18, weight="bold", ha="center", va="top")

    caption = "Report generated by https://github.com/msoedov/agentic_security"

    fig.text(0.5, 0.025, caption, fontsize=10, ha="center", va="baseline")

    buf = BytesIO()
    plt.savefig(buf, format="jpeg")
    plt.close(fig)
    buf.seek(0)
    return buf