# Copyright (c) Advanced Micro Devices, Inc.
# SPDX-License-Identifier:  MIT

from __future__ import annotations

import math
import sys
import traceback
from io import StringIO
from typing import Any, Optional

import pandas as pd
import plotext as plt
import plotly.express as px
import plotly.graph_objects as go
from textual.widgets import Static

from utils.mem_chart_gfx9 import plot_mem_chart as plot_mem_chart_gfx9
from utils.mem_chart_gfx11 import plot_mem_chart as plot_mem_chart_gfx11

# Constants
MIN_PLOT_WIDTH = 20
WIDTH_MULTIPLIER_SMALL = 3
WIDTH_MULTIPLIER_TINY = 100
HEIGHT_MULTIPLIER_SMALL = 10
HEIGHT_MULTIPLIER_TINY = 300
WIDTH_THRESHOLD_SMALL = 20
WIDTH_THRESHOLD_TINY = 1
HEIGHT_THRESHOLD_SMALL = 20
HEIGHT_THRESHOLD_TINY = 0.5
DEFAULT_WIDTH_OFFSET = 40


def simple_bar(df: pd.DataFrame, title: Optional[str] = None) -> Optional[str]:
    """
    Plot data with simple bar chart
    """

    # TODO: handle None properly

    if "Metric" in df.columns and "Avg" in df.columns:
        metric_dict = (
            pd
            .DataFrame([df["Metric"], df["Avg"]])
            .replace("", 0)
            .replace(float("inf"), -1)  # It should not happen
            .replace(float("-inf"), -1)
            .transpose()
            .set_index("Metric")
            .to_dict()["Avg"]
        )
    else:
        raise NameError(
            f"simple_bar: No Metric or Avg in df columns: {str(df.columns)}"
        )

    plt.clear_figure()

    # adjust plot size along x axis based on the max value
    w = max(list(metric_dict.values())) - DEFAULT_WIDTH_OFFSET
    if w < WIDTH_THRESHOLD_SMALL and w > WIDTH_THRESHOLD_TINY:
        w *= WIDTH_MULTIPLIER_SMALL
    elif w < WIDTH_THRESHOLD_TINY:
        w *= WIDTH_MULTIPLIER_TINY
    plt.simple_bar(list(metric_dict.keys()), list(metric_dict.values()), width=w)

    plot_content = plt.build()
    if not plot_content or plot_content.strip() == "":
        return None
    return f"\n{plot_content}\n"


def simple_multiple_bar(df: pd.DataFrame, title: Optional[str] = None) -> Optional[str]:
    """
    Plot data with simple multiple bar chart
    """

    # TODO: handle Nan and None properly

    plt.clear_figure()
    t_df = (
        df.fillna(0).replace("", 0).replace(float("inf"), -1).replace(float("-inf"), -1)
    )
    sub_labels = t_df.transpose().to_dict("split")["index"]
    sub_labels.pop(0)
    data = t_df.transpose().to_dict("split")["data"]
    labels = data.pop(0)

    plt.theme("pro")
    # adjust plot size along y axis based on the max value
    h = max(max(y) for y in data)

    if h < HEIGHT_THRESHOLD_SMALL and h > HEIGHT_THRESHOLD_TINY:
        h *= HEIGHT_MULTIPLIER_SMALL
    elif h < HEIGHT_THRESHOLD_TINY or math.isclose(h, HEIGHT_THRESHOLD_TINY):
        h *= HEIGHT_MULTIPLIER_TINY

    plt.plot_size(height=h)
    plt.multiple_bar(labels, data)

    plot_content = plt.build()
    if not plot_content or plot_content.strip() == "":
        return None
    return f"\n{plot_content}\n"


def simple_box(
    df: pd.DataFrame, orientation: str = "v", title: Optional[str] = None
) -> Optional[str]:
    """
    Plot data with simple box/whisker chart.
    Accept pre-calculated data only for now.
    """

    plt.clear_figure()
    labels: list[str] = []
    data: list[list[float]] = []

    # TODO:
    # handle Nan and None properly
    # error checking for labels
    # show unit if provided

    labels_length = 0
    t_df = (
        df.fillna(0).replace("", 0).replace(float("inf"), -1).replace(float("-inf"), -1)
    )
    for _, row in t_df.iterrows():
        column_name = row.get("Metric") or row.get("Channel")

        if column_name is None:
            raise KeyError("Neither 'Metric' nor 'Channel' column found")

        labels.append(column_name)
        # TODO: need better fix for horizontal overflow
        labels_length += len(str(column_name)) + 8
        data.append([row["Max"], row["Q3"], row["Median"], row["Q1"], row["Min"]])

    # TODO: need better fix for horizontal overflow
    if orientation == "v":
        # adjust plot size along x axis based on total labels length
        plt.plot_size(labels_length, 30)

    plt.box(
        labels,
        data,
        width=0.1,
        colors=["blue+", "orange+"],
        orientation=orientation,
    )
    plt.theme("pro")

    plot_content = plt.build()
    if not plot_content or plot_content.strip() == "":
        return None
    return f"\n{plot_content}\n"


def px_simple_bar(
    df: pd.DataFrame,
    title: Optional[str] = None,
    id: Optional[int] = None,
    style: Optional[dict[str, Any]] = None,
    orientation: str = "h",
) -> go.Figure:
    """
    Plot data with simple bar chart
    """

    # TODO: handle None properly
    if "Metric" in df.columns and ("Count" in df.columns or "Value" in df.columns):
        detected_label = "Count" if "Count" in df.columns else "Value"
        df[detected_label] = [
            x.astype(int) if x != "" else 0 for x in df[detected_label]
        ]
    else:
        raise NameError("simple_bar: No Metric or Count in df columns!")

    # Assign figure characteristics
    range_color = style.get("range_color", None) if style else None
    label_txt = style.get("label_txt", None) if style else None
    xrange = style.get("xrange", None) if style else None
    if label_txt is not None:
        label_txt = label_txt.strip("()")
        try:
            label_txt = label_txt.replace("+ $normUnit", df["Unit"][0])
        except KeyError:
            print("No units found in df. Auto labeling.")

    # Overrides for figure chatacteristics
    if id == 1701.1:
        label_txt = "%"
        range_color = [0, 100]
        xrange = [0, 110]
    if id == 1701.2:
        label_txt = "Gb/s"
        range_color = [0, 1638]
        xrange = [0, 1638]

    fig = px.bar(
        df,
        title=title,
        x=detected_label,
        y="Metric",
        color=detected_label,
        range_color=range_color,
        labels={detected_label: label_txt},
        orientation=orientation,
    ).update_xaxes(range=xrange)

    return fig


def px_simple_multi_bar(
    df: pd.DataFrame, title: Optional[str] = None, id: Optional[int] = None
) -> list[go.Figure]:
    """
    Plot data with simple multiple bar chart
    """

    # TODO: handle Nan and None properly
    if "Metric" in df.columns and "Avg" in df.columns:
        df["Avg"] = [x.astype(int) if x != "" else 0 for x in df["Avg"]]
    else:
        raise NameError("simple_multi_bar: No Metric or Count in df columns!")

    dfigs: list[go.Figure] = []
    nested_bar: dict[str, dict[str, Any]] = {}
    df_unit = df["Unit"][0]
    if id == 1604:
        nested_bar = {"NC": {}, "UC": {}, "RW": {}, "CC": {}}
        for _, row in df.iterrows():
            nested_bar[row["Coherency"]][row["Xfer"]] = row["Avg"]
    if id == 1704:
        nested_bar = {"Read": {}, "Write": {}}
        for _, row in df.iterrows():
            nested_bar[row["Transaction"]][row["Type"]] = row["Avg"]

    for group, metric in nested_bar.items():
        dfigs.append(
            px
            .bar(
                title=group,
                x=metric.values(),
                y=metric.keys(),
                labels={"x": df_unit, "y": ""},
                text=metric.values(),
            )
            .update_xaxes(showgrid=False, rangemode="nonnegative")
            .update_yaxes(showgrid=False)
        )
    return dfigs


class RooflinePlot(Static):
    """Roofline Plot visualization widget."""

    DEFAULT_CSS = """
    RooflinePlot {
        border: solid $accent;
        padding: 0;
        width: auto;
        height: auto;
        overflow-y: auto;
        overflow-x: auto;
        background: $surface;
        color: $text;
    }
    """

    def __init__(self, df: pd.DataFrame, **kwargs: Any) -> None:
        """Initialize the roofline plot"""
        super().__init__("", classes="roofline", **kwargs)
        self.df = df

        try:
            plot_str = str(self.df.get("4. Roofline", "No roofline data generated"))
            self.update(plot_str)
        except Exception as e:
            error_message = f"Roofline plot error: {str(e)}\n{traceback.format_exc()}"
            self.update(error_message)


class MemoryChart(Static):
    """Memory chart visualization widget."""

    DEFAULT_CSS = """
    MemoryChart {
        border: solid $accent;
        padding: 0;
        width: auto;
        height: auto;
        overflow-y: auto;
        overflow-x: auto;
        background: $surface;
        color: $text;
    }
    """

    def __init__(self, df: pd.DataFrame, **kwargs: Any) -> None:
        super().__init__("", classes="mem-chart", **kwargs)
        self.df = df

    def on_mount(self) -> None:
        try:
            if self.df is None or self.df.empty:
                self.update("No chart data generated")
                return

            if not {"Metric", "Value"}.issubset(self.df.columns):
                self.update("Error: Missing required columns")
                return

            metric_dict = dict(zip(self.df["Metric"], self.df["Value"]))

            # Route to arch-specific chart renderer
            mspec = getattr(self.app, "mspec", None)
            gpu_arch = mspec.gpu_arch if mspec else ""
            if gpu_arch.startswith("gfx115"):
                plot_func = plot_mem_chart_gfx11
            else:
                plot_func = plot_mem_chart_gfx9

            original_stdout = sys.stdout
            try:
                with StringIO() as string_buffer:
                    sys.stdout = string_buffer
                    result = plot_func("", "per_kernel", metric_dict)
                    stdout_output = string_buffer.getvalue()
            finally:
                sys.stdout = original_stdout

            plot_str = next(
                (x for x in [stdout_output, str(result) if result else None] if x),
                "No chart data generated",
            )
            self.update(plot_str)

        except Exception as e:
            self.update(f"Memory chart error: {str(e)}")


class SimpleBar(Static):
    """Simple Bar visualization widget."""

    DEFAULT_CSS = """
    SimpleBar {
        padding: 0;
        width: auto;
        height: auto;
        overflow-y: auto;
        overflow-x: auto;
        background: $surface;
        color: $text;
    }
    """

    def __init__(self, df: pd.DataFrame, **kwargs: Any) -> None:
        super().__init__("", classes="simple-bar", **kwargs)
        self.df = df

        try:
            result = simple_bar(self.df)

            if result:
                plot_str = str(result)
                escaped_content = plot_str.replace("[", r"\[").replace("]", r"\]")
                self.update(escaped_content)
            else:
                self.update("No simple bar data generated")

        except Exception as e:
            error_message = f"Simple Bar error: {str(e)}\n{traceback.format_exc()}"
            escaped_error = error_message.replace("[", r"\[").replace("]", r"\]")
            self.update(f"Error: {escaped_error}")


class SimpleBox(Static):
    DEFAULT_CSS = """
    SimpleBox {
        padding: 0;
        width: auto;
        height: auto;
        overflow-y: auto;
        overflow-x: auto;
        background: $surface;
        color: $text;
    }
    """

    def __init__(self, df: pd.DataFrame, **kwargs: Any) -> None:
        super().__init__("", classes="simple-box", **kwargs)
        self.df = df

        try:
            result = simple_box(self.df)

            if result:
                plot_str = str(result)
                escaped_content = plot_str.replace("[", r"\[").replace("]", r"\]")
                self.update(escaped_content)
            else:
                self.update("No simple box data generated")

        except Exception as e:
            error_message = f"Simple Box error: {str(e)}\n{traceback.format_exc()}"
            escaped_error = error_message.replace("[", r"\[").replace("]", r"\]")
            self.update(f"Error: {escaped_error}")


class SimpleMultiBar(Static):
    """Simple Multiple Bar visualization widget."""

    DEFAULT_CSS = """
    SimpleMultiBar {
        padding: 0;
        width: auto;
        height: auto;
        overflow-y: auto;
        overflow-x: auto;
        background: $surface;
        color: $text;
    }
    """

    def __init__(self, df: pd.DataFrame, **kwargs: Any) -> None:
        super().__init__("", classes="simple-multi-bar", **kwargs)
        self.df = df

        try:
            result = simple_multiple_bar(self.df)

            if result:
                plot_str = str(result)
                escaped_content = plot_str.replace("[", r"\[").replace("]", r"\]")
                self.update(escaped_content)
            else:
                self.update("No simple multi bar data generated")

        except Exception as e:
            error_message = (
                f"Simple Multiple Box error: {str(e)}\n{traceback.format_exc()}"
            )
            escaped_error = error_message.replace("[", r"\[").replace("]", r"\]")
            self.update(f"Error: {escaped_error}")
