from __future__ import annotations

import io
from typing import Any, Dict, Iterable, Optional, Sequence

import pandas as pd

from mnemosynecore.mattermost import send_file_bytes, send_file_bytes_test


def _require_matplotlib():
    try:
        import matplotlib

        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
    except ImportError as exc:
        raise ImportError(
            "Для функций визуализации установите extra зависимость: "
            "mnemosynecore[viz]"
        ) from exc
    return plt


def _fig_to_png(fig) -> bytes:
    buffer = io.BytesIO()
    fig.savefig(buffer, format="png", bbox_inches="tight", dpi=150)
    buffer.seek(0)
    png = buffer.read()
    buffer.close()
    return png


def plot_line_png(
    *,
    df: pd.DataFrame,
    x: str,
    y: str | Sequence[str],
    title: str = "",
    figsize: tuple[int, int] = (10, 4),
    marker: str = "o",
    grid: bool = True,
    sort_by_x: bool = False,
) -> bytes:
    if not isinstance(df, pd.DataFrame):
        raise TypeError("df должен быть pandas.DataFrame")
    if x not in df.columns:
        raise ValueError(f"Колонка '{x}' не найдена")

    y_cols = [y] if isinstance(y, str) else list(y)
    if not y_cols:
        raise ValueError("Нужно передать минимум одну y-колонку")
    for col in y_cols:
        if col not in df.columns:
            raise ValueError(f"Колонка '{col}' не найдена")

    frame = df[[x] + y_cols].copy()
    if sort_by_x:
        frame = frame.sort_values(x)

    plt = _require_matplotlib()
    fig, ax = plt.subplots(figsize=figsize)

    for col in y_cols:
        ax.plot(frame[x], frame[col], marker=marker, label=col)

    if len(y_cols) > 1:
        ax.legend()
    if title:
        ax.set_title(title)
    ax.set_xlabel(x)
    ax.grid(grid, alpha=0.3)
    fig.autofmt_xdate()
    fig.tight_layout()

    png = _fig_to_png(fig)
    plt.close(fig)
    return png


def plot_bar_png(
    *,
    df: pd.DataFrame,
    x: str,
    y: str | Sequence[str],
    title: str = "",
    figsize: tuple[int, int] = (10, 4),
    stacked: bool = False,
    sort_by_x: bool = False,
) -> bytes:
    if not isinstance(df, pd.DataFrame):
        raise TypeError("df должен быть pandas.DataFrame")
    if x not in df.columns:
        raise ValueError(f"Колонка '{x}' не найдена")

    y_cols = [y] if isinstance(y, str) else list(y)
    if not y_cols:
        raise ValueError("Нужно передать минимум одну y-колонку")
    for col in y_cols:
        if col not in df.columns:
            raise ValueError(f"Колонка '{col}' не найдена")

    frame = df[[x] + y_cols].copy()
    if sort_by_x:
        frame = frame.sort_values(x)

    plt = _require_matplotlib()
    fig, ax = plt.subplots(figsize=figsize)
    frame.set_index(x)[y_cols].plot(kind="bar", ax=ax, stacked=stacked)
    if title:
        ax.set_title(title)
    ax.set_xlabel(x)
    ax.grid(axis="y", alpha=0.3)
    fig.tight_layout()

    png = _fig_to_png(fig)
    plt.close(fig)
    return png


def plot_heatmap_png(
    *,
    df: pd.DataFrame,
    value_cols: Optional[Sequence[str]] = None,
    index_col: Optional[str] = None,
    title: str = "",
    figsize: tuple[int, int] = (8, 6),
    cmap: str = "Blues",
    annotate: bool = False,
) -> bytes:
    if not isinstance(df, pd.DataFrame):
        raise TypeError("df должен быть pandas.DataFrame")

    if value_cols is None:
        value_cols = list(df.select_dtypes(include=["number"]).columns)
    value_cols = list(value_cols)
    if not value_cols:
        raise ValueError("Нужно передать value_cols или иметь numeric-колонки в df")

    for col in value_cols:
        if col not in df.columns:
            raise ValueError(f"Колонка '{col}' не найдена")

    matrix = df[value_cols].apply(pd.to_numeric, errors="coerce")
    matrix = matrix.fillna(0.0)
    y_labels = (
        df[index_col].astype(str).tolist()
        if index_col and index_col in df.columns
        else [str(i) for i in range(len(matrix))]
    )

    plt = _require_matplotlib()
    fig, ax = plt.subplots(figsize=figsize)
    image = ax.imshow(matrix.to_numpy(), aspect="auto", cmap=cmap)
    ax.set_xticks(range(len(value_cols)))
    ax.set_xticklabels(value_cols, rotation=45, ha="right")
    ax.set_yticks(range(len(y_labels)))
    ax.set_yticklabels(y_labels)
    if title:
        ax.set_title(title)

    if annotate:
        values = matrix.to_numpy()
        for i in range(values.shape[0]):
            for j in range(values.shape[1]):
                ax.text(j, i, f"{values[i, j]:.2f}", ha="center", va="center")

    fig.colorbar(image, ax=ax)
    fig.tight_layout()

    png = _fig_to_png(fig)
    plt.close(fig)
    return png


def send_plot(
    *,
    channel_id: str,
    bot_id: str,
    png_bytes: bytes,
    file_name: str = "plot.png",
    text: str = "",
    root_id: Optional[str] = None,
    silent: bool = False,
) -> Dict[str, Optional[str]]:
    return send_file_bytes(
        channel_id=channel_id,
        bot_id=bot_id,
        file_name=file_name,
        file_bytes=png_bytes,
        text=text,
        mime_type="image/png",
        root_id=root_id,
        silent=silent,
    )


def send_plot_test(
    *,
    channel_id: str,
    bot_id: str,
    png_bytes: bytes,
    file_name: str = "plot.png",
    text: str = "",
    root_id: Optional[str] = None,
    dir_path: Optional[str] = None,
    creds_path: Optional[str] = None,
    silent: bool = False,
) -> Dict[str, Optional[str]]:
    return send_file_bytes_test(
        channel_id=channel_id,
        bot_id=bot_id,
        file_name=file_name,
        file_bytes=png_bytes,
        text=text,
        mime_type="image/png",
        root_id=root_id,
        dir_path=dir_path,
        creds_path=creds_path,
        silent=silent,
    )
