// Copyright (C) 2024 The Android Open Source Project
//
// 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 m from 'mithril';

import {extensions} from '../../components/extensions';
import type {time} from '../../base/time';
import {
  type QueryFlamegraphMetric,
  metricsFromTableOrSubquery,
} from '../../components/query_flamegraph';
import {FlamegraphPanel} from '../../components/flamegraph_panel';
import {FlamegraphProfile} from '../../components/flamegraph_profile';
import {convertTraceToPprofAndDownload} from '../../frontend/trace_converter';
import {Timestamp} from '../../components/widgets/timestamp';
import type {
  TrackEventDetailsPanel,
  TrackEventDetailsPanelSerializeArgs,
} from '../../public/details_panel';
import type {Trace} from '../../public/trace';
import {NUM} from '../../trace_processor/query_result';
import {Button, ButtonVariant} from '../../widgets/button';
import {Intent} from '../../widgets/common';
import {DetailsShell} from '../../widgets/details_shell';
import {Modal, showModal} from '../../widgets/modal';
import {
  Flamegraph,
  type FlamegraphState,
  FLAMEGRAPH_STATE_SCHEMA,
  type FlamegraphOptionalAction,
} from '../../widgets/flamegraph';
import type {SqlTableDefinition} from '../../components/widgets/sql/table/table_description';
import {PerfettoSqlTypes} from '../../trace_processor/perfetto_sql_type';
import {Stack} from '../../widgets/stack';
import {type ProfileDescriptor, ProfileType} from './common';

interface Props {
  ts: time;
  type: ProfileType;
}

export class HeapProfileFlamegraphDetailsPanel
  implements TrackEventDetailsPanel
{
  private readonly props: Props;
  private flamegraphModalDismissed = false;

  // TODO(lalitm): we should be able remove this around the 26Q2 timeframe
  // We moved serialization from being attached to selections to instead being
  // attached to the plugin that loaded the panel.
  readonly serialization: TrackEventDetailsPanelSerializeArgs<
    FlamegraphState | undefined
  > = {
    schema: FLAMEGRAPH_STATE_SCHEMA.optional(),
    state: undefined,
  };

  readonly metrics: ReadonlyArray<QueryFlamegraphMetric>;

  constructor(
    private readonly trace: Trace,
    private readonly heapGraphIncomplete: boolean,
    private readonly upid: number,
    private readonly profileDescriptor: ProfileDescriptor,
    private readonly ts: time,
    private readonly tsEnd: time,
    private state: FlamegraphState | undefined,
    private readonly onStateChange: (state: FlamegraphState) => void,
    onNodeSelected?: (args: {
      pathHashes: string;
      isDominator: boolean;
      upid: number;
      ts: time;
    }) => void,
  ) {
    this.props = {ts, type: profileDescriptor.type};
    this.metrics = flamegraphMetrics(
      this.trace,
      this.profileDescriptor,
      this.ts,
      this.tsEnd,
      this.upid,
      onNodeSelected
        ? (pathHashes, isDominator) =>
            onNodeSelected({pathHashes, isDominator, upid, ts})
        : undefined,
    );
    if (this.state === undefined) {
      this.state = Flamegraph.createDefaultState(this.metrics);
      onStateChange(this.state);
    }
  }

  async load() {
    // If the state in the serialization is not undefined, we should read from
    // it.
    // TODO(lalitm): remove this in 26Q2 - see comment on `serialization`.
    if (this.serialization.state !== undefined) {
      this.state = Flamegraph.updateState(
        this.serialization.state,
        this.metrics,
      );
      this.onStateChange(this.state);
      this.serialization.state = undefined;
    }
  }

  render() {
    const {type, ts} = this.props;
    return m(
      FlamegraphProfile,
      this.maybeShowModal(this.trace, type, this.heapGraphIncomplete),
      m(
        DetailsShell,
        {
          fillHeight: true,
          title: m('span', this.profileDescriptor.label),
          buttons: m(Stack, {orientation: 'horizontal', spacing: 'large'}, [
            m('span', `Snapshot time: `, m(Timestamp, {trace: this.trace, ts})),
            (type === ProfileType.NATIVE_HEAP_PROFILE ||
              type === ProfileType.JAVA_HEAP_SAMPLES) &&
              m(Button, {
                icon: 'file_download',
                intent: Intent.Primary,
                variant: ButtonVariant.Filled,
                onclick: () => {
                  downloadPprof(this.trace, this.upid, ts);
                },
              }),
          ]),
        },
        m(FlamegraphPanel, {
          trace: this.trace,
          metrics: this.metrics,
          state: this.state,
          onStateChange: (state) => {
            this.state = state;
            this.onStateChange(state);
          },
        }),
      ),
    );
  }

  private maybeShowModal(
    trace: Trace,
    type: ProfileType,
    heapGraphIncomplete: boolean,
  ) {
    if (type !== ProfileType.JAVA_HEAP_GRAPH || !heapGraphIncomplete) {
      return undefined;
    }
    if (this.flamegraphModalDismissed) {
      return undefined;
    }
    return m(Modal, {
      title: 'The flamegraph is incomplete',
      vAlign: 'TOP',
      content: m(
        'div',
        'The current trace does not have a fully formed flamegraph',
      ),
      buttons: [
        {
          text: 'Show the errors',
          primary: true,
          action: () => trace.navigate('#!/info'),
        },
        {
          text: 'Skip',
          action: () => {
            this.flamegraphModalDismissed = true;
          },
        },
      ],
    });
  }
}

function flamegraphMetrics(
  trace: Trace,
  descriptor: ProfileDescriptor,
  ts: time,
  tsEnd: time,
  upid: number,
  onNodeSelected?: (pathHashes: string, isDominator: boolean) => void,
): ReadonlyArray<QueryFlamegraphMetric> {
  switch (descriptor.type) {
    case ProfileType.NATIVE_HEAP_PROFILE:
      return flamegraphMetricsForHeapProfile(
        ts,
        tsEnd,
        upid,
        descriptor.heapName!,
        [
          {
            name: 'Unreleased Malloc Size',
            unit: 'B',
            columnName: 'self_size',
          },
          {
            name: 'Unreleased Malloc Count',
            unit: '',
            columnName: 'self_count',
          },
          {
            name: 'Total Malloc Size',
            unit: 'B',
            columnName: 'self_alloc_size',
          },
          {
            name: 'Total Malloc Count',
            unit: '',
            columnName: 'self_alloc_count',
          },
        ],
      );
    case ProfileType.GENERIC_HEAP_PROFILE:
      return flamegraphMetricsForHeapProfile(
        ts,
        tsEnd,
        upid,
        descriptor.heapName!,
        [
          {
            name: 'Unreleased Size',
            unit: 'B',
            columnName: 'self_size',
          },
          {
            name: 'Unreleased Count',
            unit: '',
            columnName: 'self_count',
          },
          {
            name: 'Total Size',
            unit: 'B',
            columnName: 'self_alloc_size',
          },
          {
            name: 'Total Count',
            unit: '',
            columnName: 'self_alloc_count',
          },
        ],
      );
    case ProfileType.JAVA_HEAP_SAMPLES:
      return flamegraphMetricsForHeapProfile(
        ts,
        tsEnd,
        upid,
        descriptor.heapName!,
        [
          {
            name: 'Total Allocation Size',
            unit: 'B',
            columnName: 'self_size',
          },
          {
            name: 'Total Allocation Count',
            unit: '',
            columnName: 'self_count',
          },
        ],
      );
    case ProfileType.JAVA_HEAP_GRAPH:
      return [
        {
          name: 'Object Size',
          unit: 'B',
          dependencySql:
            'include perfetto module android.memory.heap_graph.class_tree;',
          statement: `
            select
              id,
              parent_id as parentId,
              ifnull(name, '[Unknown]') as name,
              root_type,
              heap_type,
              self_size as value,
              self_count,
              path_hash_stable
            from _heap_graph_class_tree
            where graph_sample_ts = ${ts} and upid = ${upid}
          `,
          unaggregatableProperties: [
            {name: 'root_type', displayName: 'Root Type'},
            {name: 'heap_type', displayName: 'Heap Type'},
          ],
          aggregatableProperties: [
            {
              name: 'self_count',
              displayName: 'Self Count',
              mergeAggregation: 'SUM',
            },
            {
              name: 'path_hash_stable',
              displayName: 'Path Hash',
              mergeAggregation: 'CONCAT_WITH_COMMA',
              isVisible: (_) => false,
            },
          ],
          optionalNodeActions: getHeapGraphNodeOptionalActions(
            trace,
            false,
            onNodeSelected,
          ),
          optionalRootActions: getHeapGraphRootOptionalActions(trace, false),
        },
        {
          name: 'Object Count',
          unit: '',
          dependencySql:
            'include perfetto module android.memory.heap_graph.class_tree;',
          statement: `
            select
              id,
              parent_id as parentId,
              ifnull(name, '[Unknown]') as name,
              root_type,
              heap_type,
              self_size,
              self_count as value,
              path_hash_stable
            from _heap_graph_class_tree
            where graph_sample_ts = ${ts} and upid = ${upid}
          `,
          unaggregatableProperties: [
            {name: 'root_type', displayName: 'Root Type'},
            {name: 'heap_type', displayName: 'Heap Type'},
          ],
          aggregatableProperties: [
            {
              name: 'path_hash_stable',
              displayName: 'Path Hash',
              mergeAggregation: 'CONCAT_WITH_COMMA',
              isVisible: (_) => false,
            },
          ],
          optionalNodeActions: getHeapGraphNodeOptionalActions(
            trace,
            false,
            onNodeSelected,
          ),
          optionalRootActions: getHeapGraphRootOptionalActions(trace, false),
        },
        {
          name: 'Dominated Object Size',
          unit: 'B',
          dependencySql:
            'include perfetto module android.memory.heap_graph.dominator_class_tree;',
          statement: `
            select
              id,
              parent_id as parentId,
              ifnull(name, '[Unknown]') as name,
              root_type,
              heap_type,
              self_size as value,
              self_count,
              path_hash_stable
            from _heap_graph_dominator_class_tree
            where graph_sample_ts = ${ts} and upid = ${upid}
          `,
          unaggregatableProperties: [
            {name: 'root_type', displayName: 'Root Type'},
            {name: 'heap_type', displayName: 'Heap Type'},
          ],
          aggregatableProperties: [
            {
              name: 'self_count',
              displayName: 'Self Count',
              mergeAggregation: 'SUM',
            },
            {
              name: 'path_hash_stable',
              displayName: 'Path Hash',
              mergeAggregation: 'CONCAT_WITH_COMMA',
              isVisible: (_) => false,
            },
          ],
          optionalNodeActions: getHeapGraphNodeOptionalActions(
            trace,
            true,
            onNodeSelected,
          ),
          optionalRootActions: getHeapGraphRootOptionalActions(trace, true),
        },
        {
          name: 'Dominated Object Count',
          unit: '',
          dependencySql:
            'include perfetto module android.memory.heap_graph.dominator_class_tree;',
          statement: `
            select
              id,
              parent_id as parentId,
              ifnull(name, '[Unknown]') as name,
              root_type,
              heap_type,
              self_size,
              self_count as value,
              path_hash_stable
            from _heap_graph_dominator_class_tree
            where graph_sample_ts = ${ts} and upid = ${upid}
          `,
          unaggregatableProperties: [
            {name: 'root_type', displayName: 'Root Type'},
            {name: 'heap_type', displayName: 'Heap Type'},
          ],
          aggregatableProperties: [
            {
              name: 'path_hash_stable',
              displayName: 'Path Hash',
              mergeAggregation: 'CONCAT_WITH_COMMA',
              isVisible: (_) => false,
            },
          ],
          optionalNodeActions: getHeapGraphNodeOptionalActions(
            trace,
            true,
            onNodeSelected,
          ),
          optionalRootActions: getHeapGraphRootOptionalActions(trace, true),
        },
      ];
  }
}

function flamegraphMetricsForHeapProfile(
  ts: time,
  tsEnd: bigint,
  upid: number,
  heapName: string,
  metrics: {name: string; unit: string; columnName: string}[],
) {
  return metricsFromTableOrSubquery({
    tableOrSubquery: `
      (
        -- Any selection overlap with an allocation slice includes the
        -- slice in the result. Practically this means that we might need to
        -- extend the right-side boundary.
        with alloc_bound as (
          select ts
          from heap_profile_allocation
          where ts >= ${tsEnd}
            and upid = ${upid} and heap_name = '${heapName}'
          order by ts asc
          limit 1
        ),
        -- The native heap profiler data model is delta-encoded.
        -- Unreleased allocations will be recorded across continuous dumps
        -- and trace processor is responsible for deduplicating them.
        -- If an allocation at ts1 is released in ts2, this will
        -- be represented as an unmatched memory released in ts2 (negative size).
        -- For the purposes of looking at ts2+ slices, we need to ignore
        -- the negative sized data points.
        alloc_class as (
          select callsite_id, if(sum(count) > 0, 1, 0) as positive_alloc
          from heap_profile_allocation
          where ts >= ${ts} and ts <= ifnull((SELECT ts FROM alloc_bound), ${tsEnd})
            and upid = ${upid} and heap_name = '${heapName}'
          group by callsite_id
        )
        select
          id,
          parent_id as parentId,
          name,
          mapping_name,
          source_file || ':' || line_number as source_location,
          self_size,
          self_count,
          self_alloc_size,
          self_alloc_count
        from _android_heap_profile_callstacks_for_allocations!((
          select
            callsite_id,
            iif(positive_alloc, size, 0) as size,
            iif(positive_alloc, count, 0) as count,
            max(size, 0) as alloc_size,
            max(count, 0) as alloc_count
          from heap_profile_allocation a
          join alloc_class using (callsite_id)
          where a.ts >= ${ts} and a.ts <= ifnull((SELECT ts FROM alloc_bound), ${tsEnd})
            and a.upid = ${upid} and a.heap_name = '${heapName}'
        ))
      )
    `,
    tableMetrics: metrics,
    dependencySql:
      'include perfetto module android.memory.heap_profile.callstacks',
    unaggregatableProperties: [{name: 'mapping_name', displayName: 'Mapping'}],
    aggregatableProperties: [
      {
        name: 'source_location',
        displayName: 'Source Location',
        mergeAggregation: 'ONE_OR_SUMMARY',
      },
    ],
    nameColumnLabel: 'Symbol',
  });
}

async function downloadPprof(trace: Trace, upid: number, ts: time) {
  const pid = await trace.engine.query(
    `select pid from process where upid = ${upid}`,
  );
  if (!trace.traceInfo.downloadable) {
    showModal({
      title: 'Download not supported',
      content: m('div', 'This trace file does not support downloads'),
    });
    return;
  }
  const blob = await trace.getTraceFile();
  convertTraceToPprofAndDownload(blob, pid.firstRow({pid: NUM}).pid, ts);
}

function getHeapGraphDuplicateObjectsView(
  isDominator: boolean,
): SqlTableDefinition {
  return {
    name: `_heap_graph${tableModifier(isDominator)}duplicate_objects`,
    columns: [
      {column: 'class_name', type: PerfettoSqlTypes.STRING},
      {column: 'path_count', type: PerfettoSqlTypes.INT},
      {column: 'object_count', type: PerfettoSqlTypes.INT},
      {column: 'total_size', type: PerfettoSqlTypes.INT},
      {column: 'total_native_size', type: PerfettoSqlTypes.INT},
    ],
  };
}

function getHeapGraphNodeOptionalActions(
  trace: Trace,
  isDominator: boolean,
  onNodeSelected?: (pathHashes: string, isDominator: boolean) => void,
): ReadonlyArray<FlamegraphOptionalAction> {
  if (!trace.plugins.isPluginEnabled('com.android.HeapDumpExplorer')) {
    return [];
  }
  return [
    {
      name: 'Open in Heapdump Explorer',
      execute: async ({properties, node}) => {
        const pathHashes = properties.get('path_hash_stable');
        if (pathHashes === undefined) return;

        onNodeSelected?.(pathHashes, isDominator);

        const name = node?.name;
        const nameSuffix =
          name !== undefined ? `_${encodeURIComponent(name)}` : '';
        trace.navigate(`#!/heapdump/flamegraph_objects${nameSuffix}`);
      },
    },
  ];
}

function getHeapGraphRootOptionalActions(
  trace: Trace,
  isDominator: boolean,
): ReadonlyArray<FlamegraphOptionalAction> {
  return [
    {
      name: 'Reference paths by class',
      execute: async () => {
        const viewName = `_heap_graph${tableModifier(isDominator)}duplicate_objects`;
        const macroArgs = `_heap_graph${tableModifier(isDominator)}path_hashes`;
        const macroExpr = `_heap_graph_duplicate_objects_agg!(${macroArgs})`;
        const statement = `CREATE OR REPLACE PERFETTO VIEW ${viewName} AS SELECT * FROM ${macroExpr};`;

        // Create view to be returned
        await trace.engine.query(statement);
        extensions.addLegacySqlTableTab(trace, {
          table: getHeapGraphDuplicateObjectsView(isDominator),
        });
      },
    },
  ];
}

function tableModifier(isDominator: boolean): string {
  return isDominator ? '_dominator_' : '_';
}
