Copy

import { useEventListener } from "@vueuse/core";
import { storeToRefs } from "pinia";
import Plotly from "plotly.js-dist-min";
import { useChartStore, useUIStore } from "src/store";
import {
  defaultMarkerSize,
  selectedColor,
  selectedMarkerSize,
  unselectedMarkerSize,
  unselectedOpacity,
} from "utils/chart";
import { watch, type Ref } from "vue";

export function useChartSelection(
  chartContainer: Ref<HTMLDivElement | undefined>,
  allColors: Ref<number[]>,
  allRemoteAddrs: Ref<any[]>,
  allPorts?: Ref<number[]>,
  allSizes?: Ref<number[]>,
  allCounts?: Ref<number[]>,
  allCountries?: Ref<string[]>,
) {
  const chartStore = useChartStore().state;
  const uiStore = useUIStore();
  const {
    selectedIP,
    selectedPort,
    selectedCountry,
    selectedCount,
    selectedIPCount,
  } = storeToRefs(uiStore);

  function applyHighlight() {
    if (!chartContainer.value || !(chartContainer.value as any).data?.length)
      return;

    if (!selectedIP.value && !selectedPort.value && !selectedCountry.value) {
      resetSelectionVisuals();
      return;
    }

    if (selectedIP.value) {
      highlightIPVisuals(selectedIP.value);
    } else if (selectedPort.value) {
      highlightPortVisuals(selectedPort.value);
    } else if (selectedCountry.value) {
      highlightCountryVisuals(selectedCountry.value);
    }
  }

  function resetSelectionVisuals() {
    if (!chartContainer.value || !(chartContainer.value as any).data?.length)
      return;
    const resetData: any = {
      "marker.color": [allColors.value.slice()],
      "marker.opacity": [1],
    };
    if (allSizes?.value) {
      resetData["marker.size"] = [allSizes.value.slice()];
    } else {
      resetData["marker.size"] = [defaultMarkerSize];
    }
    Plotly.restyle(chartContainer.value, resetData);
    selectedCount.value = 0;
    selectedIPCount.value = 0;
  }

  function resetSelection() {
    uiStore.resetSelection();
    resetSelectionVisuals();
  }

  function toggleIPExclusion(ip: string) {
    const excluded = chartStore.remote_addr.includes("!" + ip);
    if (excluded) {
      chartStore.remote_addr = chartStore.remote_addr.filter(
        (a) => a !== "!" + ip,
      );
    } else {
      chartStore.remote_addr = [...chartStore.remote_addr, "!" + ip];
    }
  }

  function deleteSelectedIP() {
    if (selectedIP.value) {
      toggleIPExclusion(selectedIP.value);
      resetSelection();
    }
  }

  function highlightIP(ip: string) {
    if (selectedIP.value === ip) {
      resetSelection();
      return;
    }
    selectedIP.value = ip;
    selectedPort.value = undefined;
    selectedCountry.value = undefined;
    highlightIPVisuals(ip);
  }

  function highlightIPVisuals(ip: string) {
    if (!chartContainer.value || !(chartContainer.value as any).data?.length)
      return;

    const isMatch = (val: any) => {
      if (Array.isArray(val)) return val.includes(ip);
      return val === ip;
    };

    const uniqueIPs = new Set<string>();
    const matchingCount = allRemoteAddrs.value.reduce((acc, val, i) => {
      if (isMatch(val)) {
        if (Array.isArray(val)) val.forEach((v) => uniqueIPs.add(v));
        else uniqueIPs.add(val);
        return acc + (allCounts?.value?.[i] || 1);
      }
      return acc;
    }, 0);

    const newColorArr = allColors.value.map((origColor, i) =>
      isMatch(allRemoteAddrs.value[i]) ? selectedColor : origColor,
    );
    const newOpacityArr = allRemoteAddrs.value.map((val) =>
      isMatch(val) ? 1 : unselectedOpacity,
    );

    const restyleData: any = {
      "marker.color": [newColorArr],
      "marker.opacity": [newOpacityArr],
      "marker.line.width": [0],
    };

    if (allSizes?.value) {
      // Keep original sizes for map or other dynamic charts
      restyleData["marker.size"] = [allSizes.value];
    } else {
      restyleData["marker.size"] = [
        allRemoteAddrs.value.map((val) =>
          isMatch(val) ? selectedMarkerSize : unselectedMarkerSize,
        ),
      ];
    }

    Plotly.restyle(chartContainer.value, restyleData);
    selectedCount.value = matchingCount;
    selectedIPCount.value = uniqueIPs.size;
  }

  function highlightPort(portValue: number | string) {
    const port = Number(portValue);
    if (selectedPort.value === port) {
      resetSelection();
      return;
    }
    selectedPort.value = port;
    selectedIP.value = undefined;
    selectedCountry.value = undefined;
    highlightPortVisuals(port);
  }

  function highlightPortVisuals(port: number) {
    if (
      !chartContainer.value ||
      !(chartContainer.value as any).data?.length ||
      !allPorts?.value
    )
      return;

    const uniqueIPs = new Set<string>();
    const matchingCount = allPorts.value.reduce((acc, p, i) => {
      if (p === port) {
        const addr = allRemoteAddrs.value[i];
        if (Array.isArray(addr)) addr.forEach((a) => uniqueIPs.add(a));
        else uniqueIPs.add(addr);
        return acc + (allCounts?.value?.[i] || 1);
      }
      return acc;
    }, 0);

    const newColorArr = allColors.value.map((origColor, i) =>
      allPorts.value![i] === port ? selectedColor : origColor,
    );
    const newOpacityArr = allPorts.value.map((p) =>
      p === port ? 1 : unselectedOpacity,
    );

    const restyleData: any = {
      "marker.color": [newColorArr],
      "marker.opacity": [newOpacityArr],
      "marker.line.width": [0],
    };

    if (allSizes?.value) {
      restyleData["marker.size"] = [allSizes.value];
    } else {
      restyleData["marker.size"] = [
        allPorts.value.map((p) =>
          p === port ? selectedMarkerSize : unselectedMarkerSize,
        ),
      ];
    }

    Plotly.restyle(chartContainer.value, restyleData);
    selectedCount.value = matchingCount;
    selectedIPCount.value = uniqueIPs.size;
  }

  function highlightCountry(country: string) {
    if (selectedCountry.value === country) {
      resetSelection();
      return;
    }
    selectedCountry.value = country;
    selectedIP.value = undefined;
    selectedPort.value = undefined;
    highlightCountryVisuals(country);
  }

  function highlightCountryVisuals(country: string) {
    if (
      !chartContainer.value ||
      !(chartContainer.value as any).data?.length ||
      !allCountries?.value
    )
      return;

    const uniqueIPs = new Set<string>();
    const matchingCount = allCountries.value.reduce((acc, c, i) => {
      if (c === country) {
        const addr = allRemoteAddrs.value[i];
        if (Array.isArray(addr)) addr.forEach((a) => uniqueIPs.add(a));
        else uniqueIPs.add(addr);
        return acc + (allCounts?.value?.[i] || 1);
      }
      return acc;
    }, 0);

    const newColorArr = allColors.value.map((origColor, i) =>
      allCountries.value[i] === country ? selectedColor : origColor,
    );
    const newOpacityArr = allCountries.value.map((c) =>
      c === country ? 1 : unselectedOpacity,
    );

    const restyleData: any = {
      "marker.color": [newColorArr],
      "marker.opacity": [newOpacityArr],
      "marker.line.width": [0],
    };

    if (allSizes?.value) {
      restyleData["marker.size"] = [allSizes.value];
    } else {
      restyleData["marker.size"] = [
        allCountries.value.map((c) =>
          c === country ? selectedMarkerSize : unselectedMarkerSize,
        ),
      ];
    }

    Plotly.restyle(chartContainer.value, restyleData);
    selectedCount.value = matchingCount;
    selectedIPCount.value = uniqueIPs.size;
  }

  function handleKeyDown(e: KeyboardEvent) {
    const active = document.activeElement;
    if (
      active &&
      (active.tagName === "INPUT" ||
        active.tagName === "TEXTAREA" ||
        (active as HTMLElement).isContentEditable)
    ) {
      return;
    }

    if (e.key === "Escape" || e.key === "Esc") {
      resetSelection();
    } else if (e.key === "Delete" || e.key === "Del") {
      deleteSelectedIP();
    }
  }

  useEventListener(document, "keydown", handleKeyDown);

  watch(
    [
      chartContainer,
      allColors,
      allRemoteAddrs,
      allPorts,
      allSizes,
      allCounts,
      allCountries,
    ].filter((s) => s !== undefined),
    () => {
      applyHighlight();
    },
    { flush: "post" },
  );

  return {
    selectedIP,
    selectedPort,
    selectedCountry,
    selectedCount,
    selectedIPCount,
    selectedColor,
    highlightIP,
    highlightPort,
    highlightCountry,
    resetSelection,
    deleteSelectedIP,
    toggleIPExclusion,
  };
}