Copy

import { useResizeObserver, useThrottleFn } from "@vueuse/core";
import Plotly, {
  type PlotlyHTMLElement,
  type PlotMouseEvent,
  type PlotRelayoutEvent,
  type PlotSelectionEvent,
} from "plotly.js-dist-min";
import type { FilterActions } from "src/store";
import {
  categoricalChartMargin,
  colorscale,
  defaultChartMargin,
  defaultMarkerSize,
  gridcolor,
  hoverlabel,
  labelFontSize,
  modebar,
  portMaxAllowed,
  portMinAllowed,
  tickfontcolor,
} from "utils/chart";
import type { ChartDataPoint } from "utils/filters";
import {
  onBeforeUnmount,
  ref,
  toValue,
  type MaybeRefOrGetter,
  type Ref,
} from "vue";

export function usePlotlyChart(
  chartContainer: Ref<HTMLDivElement | undefined>,
  chartData: MaybeRefOrGetter<ChartDataPoint[]>,
  options?: {
    yField?: keyof ChartDataPoint;
    yTitle?: string;
    categoricalY?: boolean;
    onIPClick?: (ip: string, event: PlotMouseEvent) => void;
    onSelected?: (event: PlotSelectionEvent | undefined) => void;
    onRelayout?: (event: PlotRelayoutEvent) => void;
    xRange?: string[];
    yRange?: string[];
  },
) {
  const relayout = ref(false);
  const allColors = ref<number[]>([]);
  const allRemoteAddrs = ref<string[]>([]);
  const allPorts = ref<number[]>([]);

  const debouncedDrawPlot = useThrottleFn(() => drawPlot(), 250);

  onBeforeUnmount(() => {
    if (chartContainer.value) {
      Plotly.purge(chartContainer.value);
    }
  });

  function drawPlot(forceNewPlot: boolean = false) {
    if (!chartContainer.value) return;

    const data = toValue(chartData);

    const x: Date[] = [];
    const y: any[] = [];
    const colors: number[] = [];
    const customdata: Array<[string, string, number, string]> = [];
    const remoteAddrs: string[] = [];
    const ports: number[] = [];

    const yField = options?.yField || "dst_port";

    data.forEach((point: ChartDataPoint) => {
      x.push(point.time);
      y.push(point[yField]);
      colors.push(point.remote_addr_int);
      remoteAddrs.push(point.remote_addr);
      ports.push(point.dst_port);
      customdata.push([
        point.time_with_ms,
        point.remote_addr,
        point.dst_port,
        point.event ?? "",
      ]);
    });

    allColors.value = colors.slice();
    allRemoteAddrs.value = remoteAddrs.slice();
    allPorts.value = ports.slice();

    const traces: Plotly.Data[] = [
      {
        type: "scattergl" as const,
        mode: "markers" as const,
        x: x,
        y: y,
        marker: {
          size: defaultMarkerSize,
          color: colors.slice(),
          showscale: false,
          colorscale: colorscale,
        },
        hovertemplate:
          "Time: %{customdata[0]}<br>" +
          "IP: %{customdata[1]}<br>" +
          "Port: %{customdata[2]}<br>" +
          "Event: %{customdata[3]}<br>" +
          "<extra></extra>",
        customdata: customdata,
      },
    ];

    const layout: Partial<Plotly.Layout> = {
      xaxis: {
        title: {
          text: "Time",
          font: { color: tickfontcolor, size: labelFontSize },
        },
        type: "date" as const,
        gridcolor: gridcolor,
        range: options?.xRange?.length === 2 ? options.xRange : undefined,
        tickfont: { color: tickfontcolor, size: labelFontSize },
      },
      yaxis: {
        title: {
          text: options?.yTitle || "Port",
          font: { color: tickfontcolor, size: labelFontSize },
        },
        gridcolor: gridcolor,
        range: options?.yRange?.length === 2 ? options.yRange : undefined,
        tickfont: { color: tickfontcolor, size: labelFontSize },
      },
      hovermode: "closest" as const,
      margin: options?.categoricalY
        ? categoricalChartMargin
        : defaultChartMargin,
      plot_bgcolor: "transparent",
      paper_bgcolor: "transparent",
      hoverlabel: hoverlabel,
      modebar: modebar,
    };

    if (!layout.yaxis) {
      layout.yaxis = {};
    }

    if (options?.categoricalY) {
      layout.yaxis.type = "category";
      // Sort unique Y values to keep the axis consistent
      const uniqueY = [...new Set(y)].sort();
      layout.yaxis.categoryorder = "array";
      layout.yaxis.categoryarray = uniqueY;
    } else {
      layout.yaxis.tickformat = " ";
      layout.yaxis.maxallowed = portMaxAllowed;
      layout.yaxis.minallowed = portMinAllowed;
    }

    const config: Partial<Plotly.Config> = {
      responsive: true,
      scrollZoom: true,
      displaylogo: false,
    };

    if (relayout.value && !forceNewPlot) {
      const gd = chartContainer.value as any;
      const xRange = gd.layout.xaxis.range;
      const yRange = gd.layout.yaxis.range;

      const newLayout: Partial<Plotly.Layout> = {
        ...gd.layout,
      };

      if (gd.layout.xaxis.autorange === false) {
        newLayout.xaxis = { ...gd.layout.xaxis, range: xRange };
      }
      if (gd.layout.yaxis.autorange === false) {
        newLayout.yaxis = { ...gd.layout.yaxis, range: yRange };
      }

      Plotly.react(chartContainer.value, traces, newLayout, config);
    } else {
      relayout.value = true;
      const plotDiv = chartContainer.value as unknown as PlotlyHTMLElement;
      Plotly.newPlot(chartContainer.value, traces, layout, config);

      plotDiv.on("plotly_click", (eventData) => {
        const clickedIdx = eventData.points[0]?.pointIndex;
        const clickedIP = clickedIdx
          ? allRemoteAddrs.value[clickedIdx]
          : undefined;
        if (clickedIP) {
          options?.onIPClick?.(clickedIP, eventData);
        }
      });

      plotDiv.on("plotly_relayout", (eventData) => {
        options?.onRelayout?.(eventData);
      });

      plotDiv.on(
        "plotly_selected",
        (eventData: PlotSelectionEvent | undefined) => {
          options?.onSelected?.(eventData);
        },
      );
    }
  }

  useResizeObserver(chartContainer, () => {
    if (chartContainer.value && relayout.value) {
      Plotly.Plots.resize(chartContainer.value);
    }
  });

  return {
    drawPlot,
    debouncedDrawPlot,
    allColors,
    allRemoteAddrs,
    allPorts,
  };
}

export function setRangeToFilter(
  event: PlotRelayoutEvent,
  filterActions: FilterActions,
) {
  let yRange: number[] | undefined;

  // Plotly relayout events can have a nested structure or flat keys
  if (event.yaxis?.range) {
    yRange = event.yaxis.range;
  } else if (
    event["yaxis.range[0]"] !== undefined &&
    event["yaxis.range[1]"] !== undefined
  ) {
    yRange = [event["yaxis.range[0]"], event["yaxis.range[1]"]];
  }

  if (yRange?.[0] && yRange?.[1]) {
    const start = Math.max(Math.round(yRange[0]), 0);
    const end = Math.min(Math.round(yRange[1]), 65535);
    // Use setDstPort to replace the filter with the current zoom range
    filterActions.setDstPort(`${start}-${end}`);
  } else if (event["yaxis.autorange"] || event.yaxis?.autorange) {
    // If the user resets zoom (e.g. double click), clear the port filter
    filterActions.state.dst_port = [];
  }

  let xRange: number[] | undefined;

  if (event.xaxis?.range) {
    xRange = event.xaxis.range;
  } else if (
    event["xaxis.range[0]"] !== undefined &&
    event["xaxis.range[1]"] !== undefined
  ) {
    xRange = [event["xaxis.range[0]"], event["xaxis.range[1]"]];
  }

  if (xRange?.[0] && xRange?.[1]) {
    // Plotly uses ISO 8601 in local timezone, so we need to convert to UTC and then slice off the timezone part
    let start = new Date(xRange[0]);
    let end = new Date(xRange[1]);

    // add utc offset to start and end
    start.setUTCMinutes(start.getUTCMinutes() - start.getTimezoneOffset());
    end.setUTCMinutes(end.getUTCMinutes() - end.getTimezoneOffset());

    // round start down to last second and round end up to next second
    start.setMilliseconds(0);
    end.setMilliseconds(0);
    end.setSeconds(end.getSeconds() + 1);

    filterActions.setTimeStart(start.toISOString().slice(0, 19));
    filterActions.setTimeEnd(end.toISOString().slice(0, 19));
  }
}