// Copyright 2022 The Parca Authors
// Licensed 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 React, {
  memo,
  useCallback,
  useDeferredValue,
  useEffect,
  useMemo,
  useRef,
  useState,
} from 'react';

import {Table, tableFromIPC} from '@uwdata/flechette';
import {useContextMenu} from 'react-contexify';

import {FlamegraphArrow} from '@parca/client';
import {FlameGraphSkeleton, SandwichFlameGraphSkeleton, useParcaContext} from '@parca/components';
import {USER_PREFERENCES, useCurrentColorProfile, useUserPreference} from '@parca/hooks';
import {ProfileType} from '@parca/parser';
import {getColorForFeature} from '@parca/store';
import {type ColorConfig} from '@parca/utilities';

import {ProfileSource} from '../../ProfileSource';
import {useProfileFilters} from '../../ProfileView/components/ProfileFilters/useProfileFilters';
import {useProfileViewContext} from '../../ProfileView/context/ProfileViewContext';
import {TimelineGuide} from '../../TimelineGuide';
import {alignedUint8Array} from '../../utils';
import ContextMenuWrapper, {ContextMenuWrapperRef} from './ContextMenuWrapper';
import {FlameNode, RowHeight, colorByColors} from './FlameGraphNodes';
import {MemoizedTooltip} from './MemoizedTooltip';
import {MiniMap} from './MiniMap';
import {TooltipProvider} from './TooltipContext';
import {ZoomControls} from './ZoomControls';
import {useBatchedRendering} from './useBatchedRendering';
import {useScrollViewport} from './useScrollViewport';
import {useVisibleNodes} from './useVisibleNodes';
import {useZoom} from './useZoom';
import {
  CurrentPathFrame,
  boundsFromProfileSource,
  extractFeature,
  extractFilenameFeature,
  getCurrentPathFrameData,
  getMaxDepth,
  isCurrentPathFrameMatch,
} from './utils';

export const FIELD_LABELS_ONLY = 'labels_only';
export const FIELD_MAPPING_FILE = 'mapping_file';
export const FIELD_MAPPING_BUILD_ID = 'mapping_build_id';
export const FIELD_LOCATION_ADDRESS = 'location_address';
export const FIELD_LOCATION_LINE = 'location_line';
export const FIELD_INLINED = 'inlined';
export const FIELD_TIMESTAMP = 'timestamp';
export const FIELD_DURATION = 'duration';
export const FIELD_GROUPBY_METADATA = 'groupby_metadata';
export const FIELD_FUNCTION_NAME = 'function_name';
export const FIELD_FUNCTION_SYSTEM_NAME = 'function_system_name';
export const FIELD_FUNCTION_FILE_NAME = 'function_file_name';
export const FIELD_FUNCTION_START_LINE = 'function_startline';
export const FIELD_CHILDREN = 'children';
export const FIELD_LABELS = 'labels';
export const FIELD_CUMULATIVE = 'cumulative';
export const FIELD_FLAT = 'flat';
export const FIELD_DIFF = 'diff';
export const FIELD_PARENT = 'parent';
export const FIELD_DEPTH = 'depth';
export const FIELD_VALUE_OFFSET = 'value_offset';

interface FlameGraphArrowProps {
  arrow: FlamegraphArrow;
  total: bigint;
  filtered: bigint;
  profileType?: ProfileType;
  profileSource: ProfileSource;
  width?: number;
  curPath: CurrentPathFrame[];
  setCurPath: (path: CurrentPathFrame[]) => void;
  isHalfScreen: boolean;
  mappingsListFromMetadata: string[];
  filenamesListFromMetadata: string[];
  colorBy: string;
  compareAbsolute: boolean;
  isFlameChart?: boolean;
  isRenderedAsFlamegraph?: boolean;
  isInSandwichView?: boolean;
  tooltipId?: string;
  maxFrameCount?: number;
  isExpanded?: boolean;
  zoomControlsRef?: React.RefObject<HTMLDivElement | null>;
}

export const getMappingColors = (
  mappingsList: string[],
  isDarkMode: boolean,
  currentColorProfile: ColorConfig
): colorByColors => {
  const mappingFeatures = mappingsList.map(mapping => extractFeature(mapping));

  const colors: colorByColors = {};
  Object.entries(mappingFeatures).forEach(([_, feature]) => {
    colors[feature.name] = getColorForFeature(feature.name, isDarkMode, currentColorProfile.colors);
  });
  return colors;
};

export const getFilenameColors = (
  filenamesList: string[],
  isDarkMode: boolean,
  currentColorProfile: ColorConfig
): colorByColors => {
  const filenameFeatures = filenamesList.map(filename => extractFilenameFeature(filename));

  const colors: colorByColors = {};
  Object.entries(filenameFeatures).forEach(([_, feature]) => {
    colors[feature.name] = getColorForFeature(feature.name, isDarkMode, currentColorProfile.colors);
  });
  return colors;
};

const noop = (): void => {};

export const FlameGraphArrow = memo(function FlameGraphArrow({
  arrow,
  total,
  filtered,
  width,
  setCurPath,
  curPath,
  profileType,
  profileSource,
  compareAbsolute,
  isFlameChart = false,
  isRenderedAsFlamegraph = false,
  isInSandwichView = false,
  isHalfScreen,
  tooltipId = 'default',
  maxFrameCount,
  isExpanded = false,
  mappingsListFromMetadata,
  filenamesListFromMetadata,
  colorBy,
  zoomControlsRef,
}: FlameGraphArrowProps): React.JSX.Element {
  const [highlightSimilarStacksPreference] = useUserPreference<boolean>(
    USER_PREFERENCES.HIGHLIGHT_SIMILAR_STACKS.key
  );
  const [hoveringRow, setHoveringRow] = useState<number | undefined>(undefined);
  const [dockedMetainfo] = useUserPreference<boolean>(USER_PREFERENCES.GRAPH_METAINFO_DOCKED.key);
  const {perf, isDarkMode} = useParcaContext();

  const table: Table = useMemo(() => {
    const result = tableFromIPC(alignedUint8Array(arrow.record), {useBigInt: true});

    if (perf?.setMeasurement != null) {
      perf.setMeasurement('flamegraph.node_count', result.numRows);
    }

    return result;
  }, [arrow, perf]);
  const svg = useRef(null);
  const containerRef = useRef<HTMLDivElement>(null);
  const renderStartTime = useRef<number>(0);
  const hasInitialRenderCompleted = useRef(false);

  const [svgElement, setSvgElement] = useState<SVGSVGElement | null>(null);

  const {excludeBinary} = useProfileFilters();

  const {compareMode} = useProfileViewContext();
  const currentColorProfile = useCurrentColorProfile();
  const colorForSimilarNodes = currentColorProfile.colorForSimilarNodes;

  const colorByValue = colorBy === undefined || colorBy === '' ? 'binary' : colorBy;

  const filenameColors = useMemo(() => {
    const colors = getFilenameColors(filenamesListFromMetadata, isDarkMode, currentColorProfile);
    return colors;
  }, [isDarkMode, filenamesListFromMetadata, currentColorProfile]);

  const mappingColors = useMemo(() => {
    const colors = getMappingColors(mappingsListFromMetadata, isDarkMode, currentColorProfile);
    return colors;
  }, [isDarkMode, mappingsListFromMetadata, currentColorProfile]);

  const colorByList = {
    filename: filenameColors,
    binary: mappingColors,
  };

  type ColorByKey = keyof typeof colorByList;

  const colorByColors: colorByColors = colorByList[colorByValue as ColorByKey];

  const MENU_ID = 'flame-graph-context-menu';
  const contextMenuRef = useRef<ContextMenuWrapperRef>(null);
  const {show, hideAll} = useContextMenu({
    id: MENU_ID,
  });
  const displayMenu = useCallback(
    (e: React.MouseEvent, row: number): void => {
      e.preventDefault();
      // Race condition fix: Use callback to ensure context menu shows only after
      // row state has been updated and propagated through the hook chain.
      // This prevents empty function names on first click.
      contextMenuRef.current?.setRow(row, () => {
        show({
          event: e,
        });
      });
    },
    [show]
  );

  const hideBinary = (binaryToRemove: string): void => {
    // Add a new frame filter to hide this binary using the new ProfileFilters system
    excludeBinary(binaryToRemove);
  };

  const handleRowClick = useCallback(
    (row: number): void => {
      if (isFlameChart) {
        // In flame charts, we don't want to expand the node, so we return early.
        return;
      }
      // Walk down the stack starting at row until we reach the root (row 0).
      const path: CurrentPathFrame[] = [];
      let currentRow = row;
      while (currentRow > 0) {
        const frame = getCurrentPathFrameData(table, currentRow);
        path.push(frame);
        currentRow = table.getChild(FIELD_PARENT)?.get(currentRow) ?? 0;
      }

      // Reverse the path so that the root is first.
      path.reverse();
      setCurPath(path);
    },
    [table, setCurPath, isFlameChart]
  );

  const depthColumn = table.getChild(FIELD_DEPTH);
  const maxDepth = getMaxDepth(depthColumn);

  const rootCumulative = table.getChild(FIELD_CUMULATIVE)?.get(0);
  const isEmptySandwichView =
    isInSandwichView && table.numRows > 0 && (rootCumulative === null || rootCumulative === 0n);

  // Apply frame limit if maxFrameCount is provided and not expanded
  const effectiveDepth =
    maxFrameCount !== undefined && !isExpanded ? Math.min(maxDepth, maxFrameCount) : maxDepth;

  // Use deferred value to prevent UI blocking when expanding frames
  const deferredEffectiveDepth = useDeferredValue(effectiveDepth);

  const totalHeight = isInSandwichView
    ? deferredEffectiveDepth * RowHeight
    : (deferredEffectiveDepth + 1) * RowHeight;

  // Get the viewport of the container, this is used to determine which rows are visible.
  const viewport = useScrollViewport(containerRef);

  const isZoomEnabled = isFlameChart;

  const {zoomLevel, zoomIn, zoomOut, resetZoom, zoomToPosition, setZoomWithScroll, scrollLeftRef} =
    useZoom(isZoomEnabled ? containerRef : {current: null});
  const zoomedWidth = isZoomEnabled ? Math.round((width ?? 1) * zoomLevel) : width ?? 0;

  // Reset zoom when the data changes (e.g. new query, different time range)
  useEffect(() => {
    resetZoom();
  }, [table, resetZoom]);

  // To find the selected row, we must walk the current path and look at which
  // children of the current frame matches the path element exactly. Until the
  // end, the row we find at the end is our selected row.
  let currentRow = 0;
  for (const frame of curPath) {
    let childRows: number[] = Array.from(table.getChild(FIELD_CHILDREN)?.get(currentRow) ?? []);
    if (childRows.length === 0) {
      // If there are no children, we can stop here.
      break;
    }
    childRows = childRows.filter(c => isCurrentPathFrameMatch(table, c, frame));
    if (childRows.length === 0) {
      // If there are no children that match the current path frame, we can stop here.
      break;
    }
    if (childRows.length > 1) {
      // If there are multiple children that match the current path frame, we can stop here.
      // This is a case where the path is ambiguous and we cannot determine a single row.
      break;
    }
    // If there is exactly one child that matches the current path frame, we can continue.
    currentRow = childRows[0];
  }
  const selectedRow = currentRow;

  const visibleNodes = useVisibleNodes({
    table,
    viewport,
    total,
    width: zoomedWidth,
    selectedRow,
    effectiveDepth: deferredEffectiveDepth,
  });

  // Add nodes in incremental batches to avoid blocking the UI
  const {items: batchedNodes, isComplete: isBatchingComplete} = useBatchedRendering(visibleNodes, {
    batchSize: 500,
  });
  if (isBatchingComplete) {
    hasInitialRenderCompleted.current = true;
  }

  // Show skeleton only during initial load, not during scroll updates
  const showSkeleton =
    !hasInitialRenderCompleted.current && batchedNodes.length !== visibleNodes.length;

  useEffect(() => {
    if (perf?.markInteraction != null) {
      renderStartTime.current = performance.now();
    }
  }, [table, width, curPath, perf]);

  useEffect(() => {
    setSvgElement(svg.current);
  }, [tooltipId]);

  if (isEmptySandwichView) {
    return <div className="mx-auto text-center">No matching samples found</div>;
  }

  return (
    <TooltipProvider
      table={table}
      total={total}
      totalUnfiltered={total + filtered}
      profileType={profileType}
      unit={arrow.unit}
      compareAbsolute={compareAbsolute}
      tooltipId={tooltipId}
    >
      <div className="relative">
        {isZoomEnabled && (
          <ZoomControls
            zoomLevel={zoomLevel}
            zoomIn={zoomIn}
            zoomOut={zoomOut}
            resetZoom={resetZoom}
            portalRef={zoomControlsRef}
          />
        )}
        <ContextMenuWrapper
          ref={contextMenuRef}
          menuId={MENU_ID}
          table={table}
          total={total}
          totalUnfiltered={total + filtered}
          compareAbsolute={compareAbsolute}
          resetPath={() => setCurPath([])}
          hideMenu={hideAll}
          hideBinary={hideBinary}
          unit={arrow.unit}
          profileType={profileType}
          isInSandwichView={isInSandwichView}
        />
        <MemoizedTooltip contextElement={svgElement} dockedMetainfo={dockedMetainfo} />
        {showSkeleton && (
          <div className="absolute inset-0 z-10">
            {isRenderedAsFlamegraph ? (
              <SandwichFlameGraphSkeleton isHalfScreen={isHalfScreen} isDarkMode={isDarkMode} />
            ) : (
              <FlameGraphSkeleton isHalfScreen={isHalfScreen} isDarkMode={isDarkMode} />
            )}
          </div>
        )}
        {isZoomEnabled && (
          <MiniMap
            containerRef={containerRef}
            table={table}
            width={width ?? 0}
            zoomedWidth={zoomedWidth}
            totalHeight={totalHeight}
            maxDepth={deferredEffectiveDepth}
            colorByColors={colorByColors}
            colorBy={colorByValue}
            profileSource={profileSource}
            isDarkMode={isDarkMode}
            scrollLeft={viewport.scrollLeft}
            scrollLeftRef={scrollLeftRef}
            onZoomToPosition={zoomToPosition}
            onSetZoomWithScroll={setZoomWithScroll}
          />
        )}
        <div
          ref={containerRef}
          className={`${
            isZoomEnabled ? '[scrollbar-width:none] [&::-webkit-scrollbar]:hidden' : ''
          } will-change-transform webkit-overflow-scrolling-touch contain ${
            !isZoomEnabled ? 'overflow-auto' : ''
          }`}
          style={{
            width: width ?? '100%',
            ...(isZoomEnabled ? {overflowX: 'scroll' as const, overflowY: 'auto' as const} : {}),
            contain: 'layout style paint',
            visibility: !showSkeleton ? 'visible' : 'hidden',
          }}
        >
          <div>
            {isFlameChart && (
              <TimelineGuide
                bounds={boundsFromProfileSource(profileSource)}
                width={zoomedWidth}
                height={totalHeight}
                margin={0}
                ticks={12}
                timeUnit="nanoseconds"
              />
            )}
            <svg
              className="relative font-robotoMono"
              width={zoomedWidth}
              height={totalHeight}
              preserveAspectRatio="xMinYMid"
              ref={svg}
            >
              {batchedNodes.map(row => (
                <FlameNode
                  key={row}
                  table={table}
                  row={row}
                  colors={colorByColors}
                  colorBy={colorByValue}
                  totalWidth={zoomedWidth}
                  height={RowHeight}
                  darkMode={isDarkMode}
                  compareMode={compareMode}
                  colorForSimilarNodes={colorForSimilarNodes}
                  selectedRow={selectedRow}
                  onClick={() => handleRowClick(row)}
                  onContextMenu={displayMenu}
                  hoveringRow={highlightSimilarStacksPreference ? hoveringRow : undefined}
                  setHoveringRow={highlightSimilarStacksPreference ? setHoveringRow : noop}
                  isFlameChart={isFlameChart}
                  profileSource={profileSource}
                  isRenderedAsFlamegraph={isRenderedAsFlamegraph}
                  isInSandwichView={isInSandwichView}
                  maxDepth={maxDepth}
                  effectiveDepth={deferredEffectiveDepth}
                  tooltipId={tooltipId}
                />
              ))}
            </svg>
          </div>
        </div>
      </div>
    </TooltipProvider>
  );
});

export default FlameGraphArrow;