// 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 {Dispatch, SetStateAction, useCallback, useMemo, useRef, useState} from 'react';

import {RpcError} from '@protobuf-ts/runtime-rpc';
import {useQueryState as useNuqsQueryState} from 'nuqs';

import {ProfileTypesRequest, ProfileTypesResponse, QueryServiceClient} from '@parca/client';
import {DateTimeRange, IconButton, useGrpcMetadata, useParcaContext} from '@parca/components';
import {CloseIcon} from '@parca/icons';
import {Query} from '@parca/parser';
import {TEST_IDS, testId} from '@parca/test-utils';
import {millisToProtoTimestamp} from '@parca/utilities';

import {
  ProfileFilter,
  useProfileFilters,
} from '../ProfileView/components/ProfileFilters/useProfileFilters';
import {QueryControls} from '../QueryControls';
import {LabelsQueryProvider, useLabelsQueryProvider} from '../contexts/LabelsQueryProvider';
import {UnifiedLabelsProvider} from '../contexts/UnifiedLabelsContext';
import {stringParam} from '../hooks/urlParsers';
import {useLabelNames} from '../hooks/useLabels';
import {useQueryState} from '../hooks/useQueryState';
import useGrpcQuery from '../useGrpcQuery';
import {MetricsGraphSection} from './MetricsGraphSection';
import {useAutoQuerySelector} from './useAutoQuerySelector';

export interface QuerySelection {
  expression: string;
  from: number;
  to: number;
  timeSelection: string;
  sumBy?: string[];
  mergeFrom?: string;
  mergeTo?: string;
}

interface ProfileSelectorFeatures {
  showMetricsGraph: boolean;
  showSumBySelector?: boolean;
  showProfileTypeSelector?: boolean;
  disableProfileTypesDropdown?: boolean;
}

interface ProfileSelectorProps extends ProfileSelectorFeatures {
  queryClient: QueryServiceClient;
  closeProfile: () => void;
  enforcedProfileName: string;
  comparing: boolean;
  setDisplayHideMetricsGraphButton?: Dispatch<SetStateAction<boolean>>;
  suffix?: '_a' | '_b'; // For comparison mode
  onSearchHook?: () => void;
}

export interface IProfileTypesResult {
  loading: boolean;
  data?: ProfileTypesResponse;
  error?: RpcError;
}

export const useProfileTypes = (
  client: QueryServiceClient,
  start?: number,
  end?: number
): IProfileTypesResult => {
  const metadata = useGrpcMetadata();
  const metadataString = useMemo(() => JSON.stringify(metadata), [metadata]);
  const request: ProfileTypesRequest = {};

  if (start != null && end != null) {
    request.start = millisToProtoTimestamp(start);
    request.end = millisToProtoTimestamp(end);
  }

  const {isLoading, data, error} = useGrpcQuery({
    key: ['profileTypes', metadataString, start, end],
    queryFn: async abort => {
      const {response} = await client.profileTypes(request, {
        meta: metadata,
        abort,
      });
      return response;
    },
  });

  return {loading: isLoading, data, error: error as RpcError};
};

const ProfileSelector = ({
  queryClient,
  closeProfile,
  enforcedProfileName,
  comparing,
  showMetricsGraph = true,
  showSumBySelector = true,
  showProfileTypeSelector = true,
  setDisplayHideMetricsGraphButton,
  suffix,
  onSearchHook,
}: ProfileSelectorProps): JSX.Element => {
  const {externalProfilerComponent, additionalMetricsGraph} = useParcaContext();
  const [queryBrowserMode, setRawQueryBrowserMode] = useNuqsQueryState(
    'query_browser_mode',
    stringParam.withOptions({history: 'replace'})
  );
  const setQueryBrowserMode = useCallback(
    (mode: string | null) => {
      void setRawQueryBrowserMode(mode);
    },
    [setRawQueryBrowserMode]
  );

  const profileFilterDefaults = externalProfilerComponent?.profileFilterDefaults as
    | ProfileFilter[]
    | undefined;
  const {forceApplyFilters} = useProfileFilters();

  const handleProfileTypeChange = useCallback(() => {
    if (profileFilterDefaults != null && profileFilterDefaults.length > 0) {
      forceApplyFilters(profileFilterDefaults);
    }
  }, [forceApplyFilters, profileFilterDefaults]);

  // Use the new useQueryState hook - reads directly from URL params
  const {
    querySelection,
    draftSelection,
    setDraftExpression,
    setDraftTimeRange,
    setDraftSumBy,
    setDraftProfileName,
    setDraftMatchers,
    commitDraft,
    profileSelection,
    setProfileSelection,
    sumByLoading,
    draftParsedQuery,
  } = useQueryState({suffix, onProfileTypeChange: handleProfileTypeChange});

  // Use draft state for local state instead of committed state
  const [timeRangeSelection, setTimeRangeSelection] = useState(
    DateTimeRange.fromRangeKey(draftSelection.timeSelection, draftSelection.from, draftSelection.to)
  );

  // Sync local timeRangeSelection when URL state changes externally (e.g., "Switch to 1 minute" button)
  const [prevQueryTimeSelection, setPrevQueryTimeSelection] = useState(
    querySelection.timeSelection
  );
  const [prevQueryFrom, setPrevQueryFrom] = useState(querySelection.from);
  const [prevQueryTo, setPrevQueryTo] = useState(querySelection.to);

  if (
    prevQueryTimeSelection !== querySelection.timeSelection ||
    prevQueryFrom !== querySelection.from ||
    prevQueryTo !== querySelection.to
  ) {
    setPrevQueryTimeSelection(querySelection.timeSelection);
    setPrevQueryFrom(querySelection.from);
    setPrevQueryTo(querySelection.to);
    setTimeRangeSelection(
      DateTimeRange.fromRangeKey(
        querySelection.timeSelection,
        querySelection.from,
        querySelection.to
      )
    );
  }

  const [queryExpressionString, setQueryExpressionString] = useState(draftSelection.expression);

  const [advancedModeForQueryBrowser, setAdvancedModeForQueryBrowser] = useState(
    queryBrowserMode === 'advanced'
  );

  // Handler to update draft when time range changes
  const handleTimeRangeChange = useCallback(
    (range: DateTimeRange) => {
      setTimeRangeSelection(range);
      setDraftTimeRange(range.getFromMs(), range.getToMs(), range.getRangeKey());
    },
    [setDraftTimeRange]
  );

  const profileType = useMemo(() => {
    return Query.parse(queryExpressionString).profileType();
  }, [queryExpressionString]);

  const from = timeRangeSelection.getFromMs();
  const to = timeRangeSelection.getToMs();

  const {
    loading: profileTypesLoading,
    data: profileTypesData,
    error,
  } = useProfileTypes(queryClient, from, to);

  const {result} = useLabelNames(queryClient, profileType.toString(), from, to);

  const labels = useMemo(() => {
    return result.response?.labelNames === undefined ? [] : result.response.labelNames;
  }, [result]);

  const [prevEnforcedProfileName, setPrevEnforcedProfileName] = useState(enforcedProfileName);
  const [prevQueryExpression, setPrevQueryExpression] = useState(querySelection.expression);

  if (
    prevEnforcedProfileName !== enforcedProfileName ||
    prevQueryExpression !== querySelection.expression
  ) {
    setPrevEnforcedProfileName(enforcedProfileName);
    setPrevQueryExpression(querySelection.expression);
    if (enforcedProfileName !== '') {
      const [q, changed] = Query.parse(querySelection.expression).setProfileName(
        enforcedProfileName
      );
      if (changed) {
        setQueryExpressionString(q.toString());
      } else {
        setQueryExpressionString(querySelection.expression);
      }
    } else {
      setQueryExpressionString(querySelection.expression);
    }
  }

  const enforcedProfileNameQuery = (): Query => {
    const pq = Query.parse(queryExpressionString);
    const [q] = pq.setProfileName(enforcedProfileName);
    return q;
  };

  const query =
    enforcedProfileName !== '' ? enforcedProfileNameQuery() : Query.parse(queryExpressionString);
  const selectedProfileName = query.profileName();

  const setQueryExpression = (updateTs = false): void => {
    if (onSearchHook != null) {
      onSearchHook();
    }
    // When updateTs is true, re-evaluate the time range to current values
    if (updateTs) {
      // Force re-evaluation of time range (important for relative ranges like "last 15 minutes")
      const currentFrom = timeRangeSelection.getFromMs(true);
      const currentTo = timeRangeSelection.getToMs(true);
      const currentRangeKey = timeRangeSelection.getRangeKey();
      // Commit with refreshed time range
      commitDraft({
        from: currentFrom,
        to: currentTo,
        timeSelection: currentRangeKey,
      });
    } else {
      // Commit the draft with existing values
      commitDraft();
    }
  };

  const setMatchersString = (matchers: string): void => {
    // Update draft state only
    setDraftMatchers(matchers);
    setQueryExpressionString(`${selectedProfileName}{${matchers}}`);
  };

  const setProfileName = (profileName: string | undefined): void => {
    if (profileName === undefined) {
      return;
    }
    // Update draft state only
    setDraftProfileName(profileName);
    const [newQuery, changed] = query.setProfileName(profileName);
    if (changed) {
      const q = newQuery.toString();
      setQueryExpressionString(q);
    }
  };

  useAutoQuerySelector({
    selectedProfileName,
    profileTypesData,
    setProfileName,
    setQueryExpression,
    querySelection,
    loading: sumByLoading,
    defaultProfileType: externalProfilerComponent?.defaultProfileType,
  });

  const searchDisabled =
    queryExpressionString === undefined ||
    queryExpressionString === '' ||
    queryExpressionString === '{}';

  const queryBrowserRef = useRef<HTMLDivElement>(null);
  const sumByRef = useRef(null);

  return (
    <>
      <div className="mb-2 flex flex-col">
        {additionalMetricsGraph?.({
          querySelection,
          queryClient,
          suffix,
          timeRange: timeRangeSelection,
          onTimeRangeChange: handleTimeRangeChange,
          commitTimeRange: () => setQueryExpression(true),
          selectTimeRange: (range: DateTimeRange) => {
            commitDraft({
              from: range.getFromMs(),
              to: range.getToMs(),
              timeSelection: range.getRangeKey(),
            });
          },
        })}
        <LabelsQueryProvider
          setMatchersString={setMatchersString}
          runQuery={setQueryExpression}
          currentQuery={query}
          profileType={selectedProfileName ?? profileType.toString()}
          queryClient={queryClient}
          start={timeRangeSelection.getFromMs()}
          end={timeRangeSelection.getToMs()}
          suffix={suffix}
        >
          <LabelsSource>
            <QueryControls
              showProfileTypeSelector={showProfileTypeSelector}
              showSumBySelector={showSumBySelector}
              profileTypesData={profileTypesData}
              profileTypesLoading={profileTypesLoading}
              selectedProfileName={selectedProfileName}
              setProfileName={setProfileName}
              setMatchersString={setMatchersString}
              setQueryExpression={setQueryExpression}
              query={query}
              queryBrowserRef={queryBrowserRef}
              timeRangeSelection={timeRangeSelection}
              setTimeRangeSelection={handleTimeRangeChange}
              searchDisabled={searchDisabled}
              setQueryBrowserMode={setQueryBrowserMode}
              advancedModeForQueryBrowser={advancedModeForQueryBrowser}
              setAdvancedModeForQueryBrowser={setAdvancedModeForQueryBrowser}
              queryClient={queryClient}
              sumByRef={sumByRef}
              labels={labels}
              sumBySelection={draftSelection.sumBy ?? []}
              sumBySelectionLoading={sumByLoading}
              setUserSumBySelection={setDraftSumBy}
              profileType={profileType}
              profileTypesError={error}
              externalProfilerComponent={externalProfilerComponent}
              draftSelection={draftSelection}
              setDraftMatchers={setDraftMatchers}
              draftParsedQuery={draftParsedQuery}
              commitDraft={commitDraft}
            />
          </LabelsSource>
        </LabelsQueryProvider>
        {comparing && (
          <div>
            <IconButton
              onClick={() => closeProfile()}
              icon={<CloseIcon />}
              {...testId(TEST_IDS.COMPARE_CLOSE_BUTTON)}
            />
          </div>
        )}
      </div>
      <MetricsGraphSection
        showMetricsGraph={showMetricsGraph}
        setDisplayHideMetricsGraphButton={setDisplayHideMetricsGraphButton}
        querySelection={querySelection}
        profileSelection={profileSelection}
        comparing={comparing}
        sumBy={querySelection.sumBy}
        defaultSumByLoading={sumByLoading}
        queryClient={queryClient}
        queryExpressionString={queryExpressionString}
        setTimeRangeSelection={handleTimeRangeChange}
        selectQuery={commitDraft}
        setProfileSelection={setProfileSelection}
        query={query}
        setQueryExpression={setQueryExpression}
        setNewQueryExpression={setDraftExpression}
        commitDraft={commitDraft}
        profileTypesLoading={profileTypesLoading}
        hasNoProfileTypes={
          !profileTypesLoading &&
          (profileTypesData?.types == null || profileTypesData.types.length === 0)
        }
      />
    </>
  );
};

export default ProfileSelector;

const LabelsSource = ({children}: {children: React.ReactNode}): JSX.Element => {
  const {
    labelNames,
    labelValues,
    isLabelNamesLoading,
    isLabelValuesLoading,
    refetchLabelValues,
    refetchLabelNames,
    currentLabelName,
    setCurrentLabelName,
    suffix,
  } = useLabelsQueryProvider();

  return (
    <UnifiedLabelsProvider
      labelNames={labelNames}
      labelValues={labelValues}
      isLabelNamesLoading={isLabelNamesLoading}
      isLabelValuesLoading={isLabelValuesLoading}
      refetchLabelValues={refetchLabelValues}
      refetchLabelNames={refetchLabelNames}
      currentLabelName={currentLabelName}
      setCurrentLabelName={setCurrentLabelName}
      suffix={suffix}
    >
      {children}
    </UnifiedLabelsProvider>
  );
};