// Copyright (C) 2025 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, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

import {nodeRegistry} from './node_registry';
import {SlicesSourceNode} from './nodes/sources/slices_source';
import {
  modalForTableSelection,
  TableSourceNode,
  TableSourceState,
} from './nodes/sources/table_source';
import {SqlSourceNode, SqlSourceState} from './nodes/sources/sql_source';
import {AggregationNode, AggregationNodeState} from './nodes/aggregation_node';
import {
  ModifyColumnsNode,
  ModifyColumnsState,
} from './nodes/modify_columns_node';
import {AddColumnsNode, AddColumnsNodeState} from './nodes/add_columns_node';
import {
  IntervalIntersectNode,
  IntervalIntersectNodeState,
} from './nodes/interval_intersect_node';
import {MergeNode, MergeNodeState} from './nodes/merge_node';
import {SortNode, SortNodeState} from './nodes/sort_node';
import {FilterNode, FilterNodeState} from './nodes/filter_node';
import {UnionNode, UnionNodeState} from './nodes/union_node';
import {
  LimitAndOffsetNode,
  LimitAndOffsetNodeState,
} from './nodes/limit_and_offset_node';

export function registerCoreNodes() {
  nodeRegistry.register('slice', {
    name: 'Slices',
    description: 'Explore all the slices from your trace.',
    icon: 'bar_chart',
    hotkey: 's',
    type: 'source',
    factory: (state) => new SlicesSourceNode(state),
  });

  nodeRegistry.register('table', {
    name: 'Perfetto Table',
    description:
      'Query and explore data from any table in the Perfetto standard library.',
    icon: 'table_chart',
    hotkey: 't',
    type: 'source',
    preCreate: async ({sqlModules}) => {
      const selection = await modalForTableSelection(sqlModules);
      if (selection) {
        return {
          sqlTable: selection.sqlTable,
          sqlModules,
        };
      }
      return null;
    },
    factory: (state) => new TableSourceNode(state as TableSourceState),
  });

  nodeRegistry.register('sql', {
    name: 'Query Node',
    description:
      'Start with a custom SQL query to act as a source for further exploration.',
    icon: 'code',
    hotkey: 'q',
    type: 'source',
    factory: (state) => new SqlSourceNode(state as SqlSourceState),
  });

  nodeRegistry.register('aggregation', {
    name: 'Aggregation',
    description: 'Group and aggregate data from the source node.',
    icon: 'functions',
    type: 'modification',
    factory: (state) => new AggregationNode(state as AggregationNodeState),
  });

  nodeRegistry.register('modify_columns', {
    name: 'Modify Columns',
    description: 'Select, rename, and add new columns to the data.',
    icon: 'edit',
    type: 'modification',
    factory: (state) => new ModifyColumnsNode(state as ModifyColumnsState),
  });

  nodeRegistry.register('add_columns', {
    name: 'Add Columns',
    description:
      'Add columns from another node via LEFT JOIN. Connect a node to the left-side port.',
    icon: 'add_box',
    type: 'modification',
    factory: (state) => {
      const fullState: AddColumnsNodeState = {
        ...state,
        prevNode: state.prevNode!,
        selectedColumns: (state as AddColumnsNodeState).selectedColumns ?? [],
        leftColumn: (state as AddColumnsNodeState).leftColumn ?? 'id',
        rightColumn: (state as AddColumnsNodeState).rightColumn ?? 'id',
        autoExecute: false,
      };
      return new AddColumnsNode(fullState);
    },
  });

  nodeRegistry.register('interval_intersect', {
    name: 'Interval Intersect',
    description: 'Intersect the intervals with another table.',
    icon: 'timeline',
    type: 'multisource',
    factory: (state, context) => {
      if (!context) {
        throw new Error(
          'NodeFactoryContext is required for IntervalIntersectNode',
        );
      }
      const fullState: IntervalIntersectNodeState = {
        ...state,
        prevNodes: state.prevNodes ?? [],
      };
      return new IntervalIntersectNode(fullState);
    },
  });

  nodeRegistry.register('merge', {
    name: 'Merge',
    description:
      'Join two tables using equality columns or custom SQL condition.',
    icon: 'merge',
    type: 'multisource',
    factory: (state) => {
      const fullState: MergeNodeState = {
        ...state,
        prevNodes: state.prevNodes ?? [],
        leftQueryAlias: 'left',
        rightQueryAlias: 'right',
        conditionType: 'equality',
        leftColumn: '',
        rightColumn: '',
        sqlExpression: '',
      };
      return new MergeNode(fullState);
    },
  });

  nodeRegistry.register('sort_node', {
    name: 'Sort',
    description: 'Sort rows by one or more columns.',
    icon: 'sort',
    type: 'modification',
    factory: (state) => new SortNode(state as SortNodeState),
  });

  nodeRegistry.register('filter_node', {
    name: 'Filter',
    description: 'Filter rows based on column values.',
    icon: 'filter_alt',
    type: 'modification',
    factory: (state) => new FilterNode(state as FilterNodeState),
  });

  nodeRegistry.register('union_node', {
    name: 'Union',
    description: 'Combine rows from multiple sources.',
    icon: 'merge_type',
    type: 'multisource',
    factory: (state) => {
      const fullState: UnionNodeState = {
        ...state,
        prevNodes: state.prevNodes ?? [],
        selectedColumns: [],
      };
      const node = new UnionNode(fullState);
      node.onPrevNodesUpdated();
      return node;
    },
  });

  nodeRegistry.register('limit_and_offset_node', {
    name: 'Limit and Offset',
    description: 'Limit the number of rows returned and optionally skip rows.',
    icon: 'filter_list',
    type: 'modification',
    factory: (state) =>
      new LimitAndOffsetNode(state as LimitAndOffsetNodeState),
  });
}
