/*!
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
import { Box, useToken } from "@chakra-ui/react";
import {
  Chart as ChartJS,
  CategoryScale,
  LinearScale,
  PointElement,
  LineElement,
  BarElement,
  Filler,
  Title,
  Tooltip,
  Legend,
  TimeScale,
} from "chart.js";
import "chart.js/auto";
import "chartjs-adapter-dayjs-4/dist/chartjs-adapter-dayjs-4.esm";
import annotationPlugin from "chartjs-plugin-annotation";
import { useDeferredValue } from "react";
import { Bar } from "react-chartjs-2";
import { useTranslation } from "react-i18next";
import { useParams, useNavigate, useLocation, useSearchParams } from "react-router-dom";

import { useGanttServiceGetGanttData } from "openapi/queries";
import type { DagRunState, DagRunType } from "openapi/requests/types.gen";
import { useColorMode } from "src/context/colorMode";
import { useHover } from "src/context/hover";
import { useOpenGroups } from "src/context/openGroups";
import { useTimezone } from "src/context/timezone";
import { GRID_BODY_OFFSET_PX } from "src/layouts/Details/Grid/constants";
import { flattenNodes } from "src/layouts/Details/Grid/utils";
import { useGridRuns } from "src/queries/useGridRuns";
import { useGridStructure } from "src/queries/useGridStructure";
import { useGridTiSummariesStream } from "src/queries/useGridTISummaries";
import { getComputedCSSVariableValue } from "src/theme";
import { isStatePending, useAutoRefresh } from "src/utils";

import { createHandleBarClick, createHandleBarHover, createChartOptions, transformGanttData } from "./utils";

ChartJS.register(
  CategoryScale,
  LinearScale,
  PointElement,
  BarElement,
  LineElement,
  Filler,
  Title,
  Tooltip,
  Legend,
  annotationPlugin,
  TimeScale,
);

type Props = {
  readonly dagRunState?: DagRunState | undefined;
  readonly limit: number;
  readonly runAfterGte?: string | undefined;
  readonly runAfterLte?: string | undefined;
  readonly runType?: DagRunType | undefined;
  readonly triggeringUser?: string | undefined;
};

const CHART_PADDING = 36;
const CHART_ROW_HEIGHT = 20;
const MIN_BAR_WIDTH = 10;

export const Gantt = ({ dagRunState, limit, runAfterGte, runAfterLte, runType, triggeringUser }: Props) => {
  const { dagId = "", groupId: selectedGroupId, runId = "", taskId: selectedTaskId } = useParams();
  const [searchParams] = useSearchParams();
  const { openGroupIds } = useOpenGroups();
  const deferredOpenGroupIds = useDeferredValue(openGroupIds);
  const { t: translate } = useTranslation("common");
  const { selectedTimezone } = useTimezone();
  const { colorMode } = useColorMode();
  const { hoveredTaskId, setHoveredTaskId } = useHover();
  const navigate = useNavigate();
  const location = useLocation();

  const filterRoot = searchParams.get("root") ?? undefined;
  const includeUpstream = searchParams.get("upstream") === "true";
  const includeDownstream = searchParams.get("downstream") === "true";
  const depthParam = searchParams.get("depth");
  const depth = depthParam !== null && depthParam !== "" ? parseInt(depthParam, 10) : undefined;

  // Corresponds to border, brand.emphasized, and brand.muted
  const [
    lightGridColor,
    darkGridColor,
    lightSelectedColor,
    darkSelectedColor,
    lightHoverColor,
    darkHoverColor,
  ] = useToken("colors", ["gray.200", "gray.800", "brand.300", "brand.700", "brand.200", "brand.800"]);

  const gridColor = colorMode === "light" ? lightGridColor : darkGridColor;
  const selectedItemColor = colorMode === "light" ? lightSelectedColor : darkSelectedColor;
  const hoveredItemColor = colorMode === "light" ? lightHoverColor : darkHoverColor;

  const { data: gridRuns, isLoading: runsLoading } = useGridRuns({
    dagRunState,
    limit,
    runAfterGte,
    runAfterLte,
    runType,
    triggeringUser,
  });
  const { data: dagStructure, isLoading: structureLoading } = useGridStructure({
    dagRunState,
    depth,
    includeDownstream,
    includeUpstream,
    limit,
    root: filterRoot,
    runType,
    triggeringUser,
  });
  const selectedRun = gridRuns?.find((run) => run.run_id === runId);
  const refetchInterval = useAutoRefresh({ dagId });

  // Get grid summaries for groups and mapped tasks (which have min/max times)
  const { summariesByRunId } = useGridTiSummariesStream({
    dagId,
    runIds: runId && selectedRun ? [runId] : [],
    states: selectedRun ? [selectedRun.state] : [],
  });
  const gridTiSummaries = summariesByRunId.get(runId);
  const summariesLoading = Boolean(runId && selectedRun && !summariesByRunId.has(runId));

  // Single fetch for all Gantt data (individual task tries)
  const { data: ganttData, isLoading: ganttLoading } = useGanttServiceGetGanttData(
    { dagId, runId },
    undefined,
    {
      enabled: Boolean(dagId) && Boolean(runId) && Boolean(selectedRun),
      refetchInterval: (query) =>
        query.state.data?.task_instances.some((ti) => isStatePending(ti.state)) ? refetchInterval : false,
    },
  );

  const { flatNodes } = flattenNodes(dagStructure, deferredOpenGroupIds);

  const isLoading = runsLoading || structureLoading || summariesLoading || ganttLoading;

  const allTries = ganttData?.task_instances ?? [];
  const gridSummaries = gridTiSummaries?.task_instances ?? [];

  const data = isLoading || runId === "" ? [] : transformGanttData({ allTries, flatNodes, gridSummaries });

  const labels = flatNodes.map((node) => node.id);

  // Get all unique states and their colors
  const states = [...new Set(data.map((item) => item.state ?? "none"))];
  const stateColorTokens = useToken(
    "colors",
    states.map((state) => `${state}.solid`),
  );
  const stateColorMap = Object.fromEntries(
    states.map((state, index) => [
      state,
      getComputedCSSVariableValue(stateColorTokens[index] ?? "oklch(0.5 0 0)"),
    ]),
  );

  const chartData = {
    datasets: [
      {
        backgroundColor: data.map((dataItem) => stateColorMap[dataItem.state ?? "none"]),
        data: Boolean(selectedRun) ? data : [],
        maxBarThickness: CHART_ROW_HEIGHT,
        minBarLength: MIN_BAR_WIDTH,
      },
    ],
    labels,
  };

  const fixedHeight = flatNodes.length * CHART_ROW_HEIGHT + CHART_PADDING;
  const selectedId = selectedTaskId ?? selectedGroupId;

  const handleBarClick = createHandleBarClick({ dagId, data, location, navigate, runId });

  const handleBarHover = createHandleBarHover(data, setHoveredTaskId);

  const chartOptions = createChartOptions({
    data,
    gridColor,
    handleBarClick,
    handleBarHover,
    hoveredId: hoveredTaskId,
    hoveredItemColor,
    labels,
    selectedId,
    selectedItemColor,
    selectedRun,
    selectedTimezone,
    translate,
  });

  if (runId === "" || (!isLoading && !selectedRun)) {
    return undefined;
  }

  const handleChartMouseLeave = () => {
    setHoveredTaskId(undefined);

    // Clear all hover styles when mouse leaves the chart area
    const allTasks = document.querySelectorAll<HTMLDivElement>('[id*="-"]');

    allTasks.forEach((task) => {
      task.style.backgroundColor = "";
    });
  };

  return (
    <Box
      height={`${fixedHeight}px`}
      minW="250px"
      ml={-2}
      mt={`${GRID_BODY_OFFSET_PX}px`}
      onMouseLeave={handleChartMouseLeave}
      w="100%"
    >
      <Bar
        data={chartData}
        options={chartOptions}
        style={{
          paddingTop: flatNodes.length === 1 ? 15 : 1.5,
        }}
      />
    </Box>
  );
};
