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

import {IntervalIntersectNode} from './interval_intersect_node';
import {ModifyColumnsNode} from './modify_columns_node';
import type {QueryNode} from '../../query_node';
import {notifyNextNodes} from '../graph_utils';
import {
  type PerfettoSqlType,
  PerfettoSqlTypes,
} from '../../../../trace_processor/perfetto_sql_type';
import {
  createMockNode,
  createMockIntervalNode,
  createColumnInfo,
  createColumnInfoWithType,
  expectValidationSuccess,
} from '../testing/test_utils';

const intCols = (...names: string[]) =>
  names.map((n) => createColumnInfo(n, 'int'));

function makeIIN(
  inputNodes: QueryNode[],
  partitionColumns?: string[],
  tsDurSource?: number,
): IntervalIntersectNode {
  return new IntervalIntersectNode(
    {inputNodes, partitionColumns, tsDurSource},
    {},
  );
}

describe('IntervalIntersectNode', () => {
  describe('constructor', () => {
    it('should set autoExecute to false by default', () => {
      const node = makeIIN([]);

      expect(node.context.autoExecute).toBeUndefined();
    });
  });

  describe('finalCols', () => {
    it('should return empty array when no prev nodes', () => {
      const node = makeIIN([]);

      expect(node.finalCols).toEqual([]);
    });

    it('should include ts and dur from intersection without suffix', () => {
      const node1 = createMockIntervalNode('node1');
      const node2 = createMockIntervalNode('node2');

      const node = makeIIN([node1, node2]);

      const cols = node.finalCols;
      expect(cols[0].name).toBe('ts');
      expect(cols[1].name).toBe('dur');
    });

    it('should include partition columns without suffix', () => {
      const node1 = createMockNode({
        nodeId: 'node1',
        columns: intCols('id', 'ts', 'dur', 'utid'),
      });
      const node2 = createMockNode({
        nodeId: 'node2',
        columns: intCols('id', 'ts', 'dur', 'utid'),
      });

      const node = makeIIN([node1, node2], ['utid']);

      const cols = node.finalCols;
      expect(cols[0].name).toBe('ts');
      expect(cols[1].name).toBe('dur');
      expect(cols[2].name).toBe('utid');
    });

    it('should include id_N, ts_N, dur_N for each input with correct types', () => {
      const node1 = createMockIntervalNode('node1');
      const node2 = createMockIntervalNode('node2');
      const node3 = createMockIntervalNode('node3');

      const node = makeIIN([node1, node2, node3]);

      const cols = node.finalCols;

      // After ts, dur (indexes 0, 1), we should have id_0, ts_0, dur_0
      expect(cols[2].name).toBe('id_0');
      expect(cols[2].type).toEqual({kind: 'int'});
      expect(cols[3].name).toBe('ts_0');
      expect(cols[3].type).toEqual({kind: 'timestamp'}); // ts columns are TIMESTAMP type
      expect(cols[4].name).toBe('dur_0');
      expect(cols[4].type).toEqual({kind: 'duration'}); // dur columns are DURATION type

      // Then id_1, ts_1, dur_1
      expect(cols[5].name).toBe('id_1');
      expect(cols[5].type).toEqual({kind: 'int'});
      expect(cols[6].name).toBe('ts_1');
      expect(cols[6].type).toEqual({kind: 'timestamp'});
      expect(cols[7].name).toBe('dur_1');
      expect(cols[7].type).toEqual({kind: 'duration'});

      // Then id_2, ts_2, dur_2
      expect(cols[8].name).toBe('id_2');
      expect(cols[8].type).toEqual({kind: 'int'});
      expect(cols[9].name).toBe('ts_2');
      expect(cols[9].type).toEqual({kind: 'timestamp'});
      expect(cols[10].name).toBe('dur_2');
      expect(cols[10].type).toEqual({kind: 'duration'});
    });

    it('should include non-duplicated columns from all inputs', () => {
      const node1 = createMockNode({
        nodeId: 'node1',
        columns: [
          createColumnInfo('id', 'int'),
          createColumnInfo('ts', 'int'),
          createColumnInfo('dur', 'int'),
          createColumnInfo('name', 'string'),
          createColumnInfo('value', 'int'),
        ],
      });
      const node2 = createMockNode({
        nodeId: 'node2',
        columns: [
          createColumnInfo('id', 'int'),
          createColumnInfo('ts', 'int'),
          createColumnInfo('dur', 'int'),
          createColumnInfo('status', 'string'),
        ],
      });

      const node = makeIIN([node1, node2]);

      const cols = node.finalCols;
      const colNames = cols.map((c) => c.name);

      // Should have ts, dur, id_0, ts_0, dur_0, id_1, ts_1, dur_1
      expect(colNames).toContain('ts');
      expect(colNames).toContain('dur');
      expect(colNames).toContain('id_0');
      expect(colNames).toContain('ts_0');
      expect(colNames).toContain('dur_0');
      expect(colNames).toContain('id_1');
      expect(colNames).toContain('ts_1');
      expect(colNames).toContain('dur_1');

      // Should have non-duplicated columns
      expect(colNames).toContain('name');
      expect(colNames).toContain('value');
      expect(colNames).toContain('status');

      // Verify types are preserved
      const nameCol = cols.find((c) => c.name === 'name');
      expect(nameCol?.type).toEqual({kind: 'string'});
      const valueCol = cols.find((c) => c.name === 'value');
      expect(valueCol?.type).toEqual({kind: 'int'});
      const statusCol = cols.find((c) => c.name === 'status');
      expect(statusCol?.type).toEqual({kind: 'string'});
    });

    it('should exclude duplicated columns entirely', () => {
      const node1 = createMockNode({
        nodeId: 'node1',
        columns: [
          createColumnInfo('id', 'int'),
          createColumnInfo('ts', 'int'),
          createColumnInfo('dur', 'int'),
          createColumnInfo('name', 'string'),
          createColumnInfo('value', 'int'),
        ],
      });
      const node2 = createMockNode({
        nodeId: 'node2',
        columns: [
          createColumnInfo('id', 'int'),
          createColumnInfo('ts', 'int'),
          createColumnInfo('dur', 'int'),
          createColumnInfo('name', 'string'), // Duplicate
          createColumnInfo('status', 'string'),
        ],
      });

      const node = makeIIN([node1, node2]);

      const cols = node.finalCols;
      const colNames = cols.map((c) => c.name);

      // Should NOT include 'name' since it's duplicated
      expect(colNames).not.toContain('name');
      // Should include unique columns
      expect(colNames).toContain('value');
      expect(colNames).toContain('status');
    });

    it('should exclude partition columns from duplicated columns', () => {
      const node1 = createMockNode({
        nodeId: 'node1',
        columns: [
          createColumnInfo('id', 'int'),
          createColumnInfo('ts', 'int'),
          createColumnInfo('dur', 'int'),
          createColumnInfo('utid', 'int'),
          createColumnInfo('name', 'string'),
        ],
      });
      const node2 = createMockNode({
        nodeId: 'node2',
        columns: [
          createColumnInfo('id', 'int'),
          createColumnInfo('ts', 'int'),
          createColumnInfo('dur', 'int'),
          createColumnInfo('utid', 'int'),
          createColumnInfo('status', 'string'),
        ],
      });

      const node = makeIIN([node1, node2], ['utid']);

      const cols = node.finalCols;
      const colNames = cols.map((c) => c.name);

      // utid should appear once without suffix (as partition column)
      const utidOccurrences = colNames.filter((n) => n === 'utid');
      expect(utidOccurrences.length).toBe(1);

      // Should have other non-duplicated columns
      expect(colNames).toContain('name');
      expect(colNames).toContain('status');
    });

    it('should handle multiple partition columns', () => {
      const node1 = createMockNode({
        nodeId: 'node1',
        columns: intCols('id', 'ts', 'dur', 'utid', 'upid'),
      });
      const node2 = createMockNode({
        nodeId: 'node2',
        columns: intCols('id', 'ts', 'dur', 'utid', 'upid'),
      });

      const node = makeIIN([node1, node2], ['utid', 'upid']);

      const cols = node.finalCols;
      expect(cols[0].name).toBe('ts');
      expect(cols[1].name).toBe('dur');
      expect(cols[2].name).toBe('utid');
      expect(cols[3].name).toBe('upid');
    });

    it('should set all columns as checked', () => {
      const node1 = createMockNode({
        nodeId: 'node1',
        columns: [
          createColumnInfo('id', 'int', {checked: false}),
          createColumnInfo('ts', 'int', {checked: false}),
          createColumnInfo('dur', 'int', {checked: false}),
          createColumnInfo('name', 'string', {checked: false}),
        ],
      });
      const node2 = createMockNode({
        nodeId: 'node2',
        columns: [
          createColumnInfo('id', 'int', {checked: false}),
          createColumnInfo('ts', 'int', {checked: false}),
          createColumnInfo('dur', 'int', {checked: false}),
        ],
      });

      const node = makeIIN([node1, node2]);

      const cols = node.finalCols;
      // All columns should be checked
      for (const col of cols) {
        expect(col.checked).toBe(true);
      }
    });

    it('should preserve types for all columns including non-duplicated columns', () => {
      const node1 = createMockNode({
        nodeId: 'node1',
        columns: [
          createColumnInfo('id', 'int'),
          createColumnInfo('ts', 'int'),
          createColumnInfo('dur', 'int'),
          createColumnInfo('name', 'string'),
          createColumnInfo('value', 'double'),
          createColumnInfo('count', 'int'),
        ],
      });
      const node2 = createMockNode({
        nodeId: 'node2',
        columns: [
          createColumnInfo('id', 'int'),
          createColumnInfo('ts', 'int'),
          createColumnInfo('dur', 'int'),
          createColumnInfo('status', 'string'),
          createColumnInfo('priority', 'int'),
        ],
      });
      const node3 = createMockNode({
        nodeId: 'node3',
        columns: [
          createColumnInfo('id', 'int'),
          createColumnInfo('ts', 'int'),
          createColumnInfo('dur', 'int'),
          createColumnInfo('category', 'string'),
        ],
      });

      const node = makeIIN([node1, node2, node3]);

      const cols = node.finalCols;

      // Check id_N, ts_N, dur_N types for each input
      // ts columns are TIMESTAMP type, dur columns are DURATION type
      const id0 = cols.find((c) => c.name === 'id_0');
      expect(id0?.type).toEqual({kind: 'int'});
      const ts0 = cols.find((c) => c.name === 'ts_0');
      expect(ts0?.type).toEqual({kind: 'timestamp'});
      const dur0 = cols.find((c) => c.name === 'dur_0');
      expect(dur0?.type).toEqual({kind: 'duration'});

      const id1 = cols.find((c) => c.name === 'id_1');
      expect(id1?.type).toEqual({kind: 'int'});
      const ts1 = cols.find((c) => c.name === 'ts_1');
      expect(ts1?.type).toEqual({kind: 'timestamp'});
      const dur1 = cols.find((c) => c.name === 'dur_1');
      expect(dur1?.type).toEqual({kind: 'duration'});

      const id2 = cols.find((c) => c.name === 'id_2');
      expect(id2?.type).toEqual({kind: 'int'});
      const ts2 = cols.find((c) => c.name === 'ts_2');
      expect(ts2?.type).toEqual({kind: 'timestamp'});
      const dur2 = cols.find((c) => c.name === 'dur_2');
      expect(dur2?.type).toEqual({kind: 'duration'});

      // Check non-duplicated columns preserve their types
      const name = cols.find((c) => c.name === 'name');
      expect(name?.type).toEqual({kind: 'string'});
      const value = cols.find((c) => c.name === 'value');
      expect(value?.type).toEqual({kind: 'double'});
      const count = cols.find((c) => c.name === 'count');
      expect(count?.type).toEqual({kind: 'int'});
      const status = cols.find((c) => c.name === 'status');
      expect(status?.type).toEqual({kind: 'string'});
      const priority = cols.find((c) => c.name === 'priority');
      expect(priority?.type).toEqual({kind: 'int'});
      const category = cols.find((c) => c.name === 'category');
      expect(category?.type).toEqual({kind: 'string'});
    });

    it('should exclude columns with conflicting types', () => {
      const node1 = createMockNode({
        nodeId: 'node1',
        columns: [
          createColumnInfo('id', 'int'),
          createColumnInfo('ts', 'int'),
          createColumnInfo('dur', 'int'),
          createColumnInfo('value', 'double'),
          createColumnInfo('unique1', 'string'),
        ],
      });
      const node2 = createMockNode({
        nodeId: 'node2',
        columns: intCols('id', 'ts', 'dur', 'value', 'unique2'),
      });

      const node = makeIIN([node1, node2]);

      const cols = node.finalCols;
      const colNames = cols.map((c) => c.name);

      // Should NOT include 'value' since it's duplicated (even with different types)
      expect(colNames).not.toContain('value');
      // Should include unique columns
      expect(colNames).toContain('unique1');
      expect(colNames).toContain('unique2');
    });
  });

  describe('validation', () => {
    it('should fail validation with less than two inputs', () => {
      const node1 = createMockIntervalNode('node1');

      const node = makeIIN([node1]);

      expect(node.validate()).toBe(false);
      expect(node.context.issues?.queryError?.message).toContain(
        'requires at least two inputs',
      );
    });

    it('should pass validation with two valid inputs', () => {
      const node1 = createMockIntervalNode('node1');
      const node2 = createMockIntervalNode('node2');

      const node = makeIIN([node1, node2]);

      expectValidationSuccess(node);
    });

    it('should fail validation when input is missing required id column', () => {
      const node1 = createMockIntervalNode('node1');
      const node2 = createMockNode({
        nodeId: 'node2',
        columns: intCols('ts', 'dur'),
      });

      const node = makeIIN([node1, node2]);

      expect(node.validate()).toBe(false);
      expect(node.context.issues?.queryError?.message).toContain(
        'missing required columns',
      );
      expect(node.context.issues?.queryError?.message).toContain('id');
    });

    it('should fail validation when input is missing required ts column', () => {
      const node1 = createMockIntervalNode('node1');
      const node2 = createMockNode({
        nodeId: 'node2',
        columns: intCols('id', 'dur'),
      });

      const node = makeIIN([node1, node2]);

      expect(node.validate()).toBe(false);
      expect(node.context.issues?.queryError?.message).toContain(
        'missing required columns',
      );
      expect(node.context.issues?.queryError?.message).toContain('ts');
    });

    it('should fail validation when input is missing required dur column', () => {
      const node1 = createMockIntervalNode('node1');
      const node2 = createMockNode({
        nodeId: 'node2',
        columns: intCols('id', 'ts'),
      });

      const node = makeIIN([node1, node2]);

      expect(node.validate()).toBe(false);
      expect(node.context.issues?.queryError?.message).toContain(
        'missing required columns',
      );
      expect(node.context.issues?.queryError?.message).toContain('dur');
    });

    it('should fail validation when prev node validation fails', () => {
      const node1 = createMockIntervalNode('node1');
      const node2 = createMockIntervalNode('node2');
      node2.validate = () => false;

      const node = makeIIN([node1, node2]);

      expect(node.validate()).toBe(false);
      expect(node.context.issues?.queryError?.message).toContain('is invalid');
    });

    it('should clear previous errors when validation starts', () => {
      const node1 = createMockIntervalNode('node1');

      const node = makeIIN([node1]);

      // First validation should fail
      expect(node.validate()).toBe(false);
      expect(node.context.issues?.queryError).toBeDefined();

      // Add second node to make it valid
      const node2 = createMockIntervalNode('node2');
      node.secondaryInputs.connections.set(1, node2);

      // Second validation should pass and clear errors
      expectValidationSuccess(node);
      expect(node.context.issues?.queryError).toBeUndefined();
    });

    it('should fail validation when partition column is missing from an input', () => {
      const node1 = createMockNode({
        nodeId: 'node1',
        columns: intCols('id', 'ts', 'dur', 'utid'),
      });
      const node2 = createMockNode({
        nodeId: 'node2',
        columns: intCols('id', 'ts', 'dur', 'utid'),
      });
      const node3 = createMockNode({
        nodeId: 'node3',
        columns: intCols('id', 'ts', 'dur'),
      });

      const node = makeIIN([node1, node2, node3], ['utid']);

      expect(node.validate()).toBe(false);
      expect(node.context.issues?.queryError?.message).toContain(
        "Partition column 'utid' is missing from Input 2",
      );
      expect(node.context.issues?.queryError?.message).toContain(
        'remove the partitioning',
      );
    });

    it('should pass validation when all inputs have partition columns', () => {
      const node1 = createMockNode({
        nodeId: 'node1',
        columns: intCols('id', 'ts', 'dur', 'utid'),
      });
      const node2 = createMockNode({
        nodeId: 'node2',
        columns: intCols('id', 'ts', 'dur', 'utid'),
      });
      const node3 = createMockNode({
        nodeId: 'node3',
        columns: intCols('id', 'ts', 'dur', 'utid'),
      });

      const node = makeIIN([node1, node2, node3], ['utid']);

      expectValidationSuccess(node);
    });

    it('should fail validation when only some partition columns are missing', () => {
      const node1 = createMockNode({
        nodeId: 'node1',
        columns: intCols('id', 'ts', 'dur', 'utid', 'upid'),
      });
      const node2 = createMockNode({
        nodeId: 'node2',
        columns: intCols('id', 'ts', 'dur', 'utid'),
      });

      const node = makeIIN([node1, node2], ['utid', 'upid']);

      expect(node.validate()).toBe(false);
      expect(node.context.issues?.queryError?.message).toContain(
        "Partition column 'upid' is missing from Input 1",
      );
    });

    it('should pass validation with empty partition columns array', () => {
      const node1 = createMockNode({
        nodeId: 'node1',
        columns: intCols('id', 'ts', 'dur', 'utid'),
      });
      const node2 = createMockIntervalNode('node2');

      const node = makeIIN([node1, node2], []);

      expectValidationSuccess(node);
      expect(node.context.issues?.queryError).toBeUndefined();
    });
  });

  describe('getTitle', () => {
    it('should return "Interval Intersect"', () => {
      const node = makeIIN([]);

      expect(node.getTitle()).toBe('Interval Intersect');
    });
  });

  describe('clone', () => {
    it('should create a deep copy of the node', () => {
      const node1 = createMockIntervalNode('node1');
      const node2 = createMockIntervalNode('node2');

      const node = makeIIN([node1, node2], ['utid']);

      const cloned = node.clone() as IntervalIntersectNode;

      expect(cloned).toBeInstanceOf(IntervalIntersectNode);
      expect(cloned.nodeId).not.toBe(node.nodeId);
      expect((cloned as IntervalIntersectNode).attrs.partitionColumns).toEqual([
        'utid',
      ]);
    });

    it('should not share state arrays with original', () => {
      const node1 = createMockIntervalNode('node1');
      const node2 = createMockIntervalNode('node2');

      const node = makeIIN([node1, node2], ['utid']);

      const cloned = node.clone() as IntervalIntersectNode;

      // Modify cloned arrays
      (cloned as IntervalIntersectNode).attrs.partitionColumns!.push('upid');

      // Original should not be affected
      expect(node.attrs.partitionColumns).toEqual(['utid']);
    });
  });

  describe('serializeState', () => {
    it('should serialize state correctly', () => {
      const node1 = createMockIntervalNode('node1');
      const node2 = createMockIntervalNode('node2');
      const node3 = createMockIntervalNode('node3');

      const node = makeIIN([node1, node2, node3], ['utid']);

      const serialized = node.attrs;

      expect(serialized.partitionColumns).toEqual(['utid']);
    });

    it('should handle empty partition columns', () => {
      const node1 = createMockIntervalNode('node1');
      const node2 = createMockIntervalNode('node2');

      const node = makeIIN([node1, node2]);

      const serialized = node.attrs;

      expect(serialized.partitionColumns).toBeUndefined();
    });
  });

  describe('deserialize', () => {
    it('should restore partition columns via constructor', () => {
      const node = new IntervalIntersectNode({partitionColumns: ['utid']}, {});

      expect(node.attrs.partitionColumns).toEqual(['utid']);
    });

    it('should handle missing partition columns', () => {
      const node = new IntervalIntersectNode({}, {});

      expect(node.attrs.partitionColumns).toBeUndefined();
    });
  });

  describe('ModifyColumnsNode integration', () => {
    it('should pass columns with correct types to ModifyColumnsNode', () => {
      // Create input nodes with proper SQL types
      const node1 = createMockNode({
        nodeId: 'node1',
        columns: [
          createColumnInfoWithType('id', PerfettoSqlTypes.INT),
          createColumnInfoWithType('ts', PerfettoSqlTypes.TIMESTAMP),
          createColumnInfoWithType('dur', PerfettoSqlTypes.DURATION),
          createColumnInfoWithType('name', PerfettoSqlTypes.STRING),
        ],
      });
      const node2 = createMockNode({
        nodeId: 'node2',
        columns: [
          createColumnInfoWithType('id', PerfettoSqlTypes.INT),
          createColumnInfoWithType('ts', PerfettoSqlTypes.TIMESTAMP),
          createColumnInfoWithType('dur', PerfettoSqlTypes.DURATION),
          createColumnInfoWithType('status', PerfettoSqlTypes.STRING),
        ],
      });

      const intervalNode = makeIIN([node1, node2]);

      // Create ModifyColumnsNode with IntervalIntersectNode as input
      const modifyNode = new ModifyColumnsNode({selectedColumns: []}, {});
      modifyNode.primaryInput = intervalNode;
      modifyNode.onPrevNodesUpdated();

      const selectedCols = modifyNode.attrs.selectedColumns;

      // Verify ts column has TIMESTAMP type
      const tsCol = selectedCols.find((c) => c.name === 'ts');
      expect(tsCol).toBeDefined();
      expect(tsCol?.type).toEqual({kind: 'timestamp'});
      expect(tsCol?.type).toEqual(PerfettoSqlTypes.TIMESTAMP);

      // Verify dur column has DURATION type
      const durCol = selectedCols.find((c) => c.name === 'dur');
      expect(durCol).toBeDefined();
      expect(durCol?.type).toEqual({kind: 'duration'});
      expect(durCol?.type).toEqual(PerfettoSqlTypes.DURATION);

      // Verify ts_0 column has TIMESTAMP type
      const ts0Col = selectedCols.find((c) => c.name === 'ts_0');
      expect(ts0Col).toBeDefined();
      expect(ts0Col?.type).toEqual({kind: 'timestamp'});
      expect(ts0Col?.type).toEqual(PerfettoSqlTypes.TIMESTAMP);

      // Verify dur_0 column has DURATION type
      const dur0Col = selectedCols.find((c) => c.name === 'dur_0');
      expect(dur0Col).toBeDefined();
      expect(dur0Col?.type).toEqual({kind: 'duration'});
      expect(dur0Col?.type).toEqual(PerfettoSqlTypes.DURATION);

      // Verify ts_1 column has TIMESTAMP type
      const ts1Col = selectedCols.find((c) => c.name === 'ts_1');
      expect(ts1Col).toBeDefined();
      expect(ts1Col?.type).toEqual({kind: 'timestamp'});
      expect(ts1Col?.type).toEqual(PerfettoSqlTypes.TIMESTAMP);

      // Verify dur_1 column has DURATION type
      const dur1Col = selectedCols.find((c) => c.name === 'dur_1');
      expect(dur1Col).toBeDefined();
      expect(dur1Col?.type).toEqual({kind: 'duration'});
      expect(dur1Col?.type).toEqual(PerfettoSqlTypes.DURATION);
    });

    it('should pass partition columns with original types to ModifyColumnsNode', () => {
      const node1 = createMockNode({
        nodeId: 'node1',
        columns: [
          createColumnInfoWithType('id', PerfettoSqlTypes.INT),
          createColumnInfoWithType('ts', PerfettoSqlTypes.TIMESTAMP),
          createColumnInfoWithType('dur', PerfettoSqlTypes.DURATION),
          createColumnInfoWithType('utid', PerfettoSqlTypes.INT),
        ],
      });
      const node2 = createMockNode({
        nodeId: 'node2',
        columns: [
          createColumnInfoWithType('id', PerfettoSqlTypes.INT),
          createColumnInfoWithType('ts', PerfettoSqlTypes.TIMESTAMP),
          createColumnInfoWithType('dur', PerfettoSqlTypes.DURATION),
          createColumnInfoWithType('utid', PerfettoSqlTypes.INT),
        ],
      });

      const intervalNode = makeIIN([node1, node2], ['utid']);

      // Create ModifyColumnsNode with IntervalIntersectNode as input
      const modifyNode = new ModifyColumnsNode({selectedColumns: []}, {});
      modifyNode.primaryInput = intervalNode;
      modifyNode.onPrevNodesUpdated();

      const selectedCols = modifyNode.attrs.selectedColumns;

      // Verify utid column preserves its type
      const utidCol = selectedCols.find((c) => c.name === 'utid');
      expect(utidCol).toBeDefined();
      expect(utidCol?.type).toEqual({kind: 'int'});
      expect(utidCol?.type).toEqual(PerfettoSqlTypes.INT);
    });

    it('should pass all expected columns to ModifyColumnsNode', () => {
      const node1 = createMockNode({
        nodeId: 'node1',
        columns: [
          createColumnInfo('id', 'int'),
          createColumnInfo('ts', 'int'),
          createColumnInfo('dur', 'int'),
          createColumnInfo('name', 'string'),
        ],
      });
      const node2 = createMockNode({
        nodeId: 'node2',
        columns: [
          createColumnInfo('id', 'int'),
          createColumnInfo('ts', 'int'),
          createColumnInfo('dur', 'int'),
          createColumnInfo('status', 'string'),
        ],
      });

      const intervalNode = makeIIN([node1, node2]);

      // Create ModifyColumnsNode with IntervalIntersectNode as input
      const modifyNode = new ModifyColumnsNode({selectedColumns: []}, {});
      modifyNode.primaryInput = intervalNode;
      modifyNode.onPrevNodesUpdated();

      const colNames = modifyNode.attrs.selectedColumns.map((c) => c.name);

      // Should have all expected columns
      expect(colNames).toContain('ts');
      expect(colNames).toContain('dur');
      expect(colNames).toContain('id_0');
      expect(colNames).toContain('ts_0');
      expect(colNames).toContain('dur_0');
      expect(colNames).toContain('id_1');
      expect(colNames).toContain('ts_1');
      expect(colNames).toContain('dur_1');
      expect(colNames).toContain('name');
      expect(colNames).toContain('status');
    });

    it('should pass non-duplicated columns with original types to ModifyColumnsNode', () => {
      // Non-duplicated columns should preserve their types from input nodes
      const node1 = createMockNode({
        nodeId: 'node1',
        columns: [
          createColumnInfoWithType('id', PerfettoSqlTypes.INT),
          createColumnInfoWithType('ts', PerfettoSqlTypes.TIMESTAMP),
          createColumnInfoWithType('dur', PerfettoSqlTypes.DURATION),
          createColumnInfoWithType('name', PerfettoSqlTypes.STRING),
          createColumnInfoWithType('value', PerfettoSqlTypes.DOUBLE),
        ],
      });
      const node2 = createMockNode({
        nodeId: 'node2',
        columns: [
          createColumnInfoWithType('id', PerfettoSqlTypes.INT),
          createColumnInfoWithType('ts', PerfettoSqlTypes.TIMESTAMP),
          createColumnInfoWithType('dur', PerfettoSqlTypes.DURATION),
          createColumnInfoWithType('status', PerfettoSqlTypes.STRING),
        ],
      });

      const intervalNode = makeIIN([node1, node2]);

      // Create ModifyColumnsNode with IntervalIntersectNode as input
      const modifyNode = new ModifyColumnsNode({selectedColumns: []}, {});
      modifyNode.primaryInput = intervalNode;
      modifyNode.onPrevNodesUpdated();

      const selectedCols = modifyNode.attrs.selectedColumns;

      // Verify non-duplicated columns preserve their types
      const nameCol = selectedCols.find((c) => c.name === 'name');
      expect(nameCol).toBeDefined();
      expect(nameCol?.type).toEqual({kind: 'string'});
      expect(nameCol?.type).toEqual(PerfettoSqlTypes.STRING);

      const valueCol = selectedCols.find((c) => c.name === 'value');
      expect(valueCol).toBeDefined();
      expect(valueCol?.type).toEqual({kind: 'double'});
      expect(valueCol?.type).toEqual(PerfettoSqlTypes.DOUBLE);

      const statusCol = selectedCols.find((c) => c.name === 'status');
      expect(statusCol).toBeDefined();
      expect(statusCol?.type).toEqual({kind: 'string'});
      expect(statusCol?.type).toEqual(PerfettoSqlTypes.STRING);
    });

    it('should have column.type set correctly on finalCols', () => {
      // Test that IntervalIntersectNode.finalCols has column.type set
      const node1 = createMockNode({
        nodeId: 'node1',
        columns: [
          createColumnInfoWithType('id', PerfettoSqlTypes.INT),
          createColumnInfoWithType('ts', PerfettoSqlTypes.TIMESTAMP),
          createColumnInfoWithType('dur', PerfettoSqlTypes.DURATION),
        ],
      });
      const node2 = createMockNode({
        nodeId: 'node2',
        columns: [
          createColumnInfoWithType('id', PerfettoSqlTypes.INT),
          createColumnInfoWithType('ts', PerfettoSqlTypes.TIMESTAMP),
          createColumnInfoWithType('dur', PerfettoSqlTypes.DURATION),
        ],
      });

      const intervalNode = makeIIN([node1, node2]);

      const cols = intervalNode.finalCols;

      // Check ts column
      const tsCol = cols.find((c) => c.name === 'ts');
      expect(tsCol?.type).toEqual({kind: 'timestamp'});

      // Check dur column
      const durCol = cols.find((c) => c.name === 'dur');
      expect(durCol?.type).toEqual({kind: 'duration'});

      // Check ts_0 column
      const ts0Col = cols.find((c) => c.name === 'ts_0');
      expect(ts0Col?.type).toEqual({kind: 'timestamp'});

      // Check dur_0 column
      const dur0Col = cols.find((c) => c.name === 'dur_0');
      expect(dur0Col?.type).toEqual({kind: 'duration'});
    });

    it('should propagate partition columns to ModifyColumnsNode when added after creation', () => {
      // Create input nodes with partition columns available
      const node1 = createMockNode({
        nodeId: 'node1',
        columns: [
          createColumnInfoWithType('id', PerfettoSqlTypes.INT),
          createColumnInfoWithType('ts', PerfettoSqlTypes.TIMESTAMP),
          createColumnInfoWithType('dur', PerfettoSqlTypes.DURATION),
          createColumnInfoWithType('utid', PerfettoSqlTypes.INT),
          createColumnInfoWithType('track_id', PerfettoSqlTypes.INT),
        ],
      });
      const node2 = createMockNode({
        nodeId: 'node2',
        columns: [
          createColumnInfoWithType('id', PerfettoSqlTypes.INT),
          createColumnInfoWithType('ts', PerfettoSqlTypes.TIMESTAMP),
          createColumnInfoWithType('dur', PerfettoSqlTypes.DURATION),
          createColumnInfoWithType('utid', PerfettoSqlTypes.INT),
          createColumnInfoWithType('track_id', PerfettoSqlTypes.INT),
        ],
      });

      // Create IntervalIntersectNode WITHOUT partition columns initially
      const intervalNode = makeIIN([node1, node2]);

      // Create ModifyColumnsNode with IntervalIntersectNode as input
      const modifyNode = new ModifyColumnsNode({selectedColumns: []}, {});
      modifyNode.primaryInput = intervalNode;
      intervalNode.nextNodes.push(modifyNode);
      modifyNode.onPrevNodesUpdated();

      // Verify initial columns (no partition columns yet)
      let selectedCols = modifyNode.attrs.selectedColumns;
      expect(selectedCols.find((c) => c.name === 'utid')).toBeUndefined();
      expect(selectedCols.find((c) => c.name === 'track_id')).toBeUndefined();

      // Count initial columns (should be: ts, dur, id_0, ts_0, dur_0, id_1, ts_1, dur_1)
      const initialColCount = selectedCols.length;
      expect(initialColCount).toBe(8);

      // NOW add partition columns to the interval intersect node
      intervalNode.attrs.partitionColumns = ['utid', 'track_id'];

      // Notify downstream nodes about the column change
      notifyNextNodes(intervalNode);
      intervalNode.context.onchange?.();

      // Verify that ModifyColumnsNode received the updated columns including partitions
      selectedCols = modifyNode.attrs.selectedColumns;

      // Should now include utid and track_id partition columns
      const utidCol = selectedCols.find((c) => c.name === 'utid');
      expect(utidCol).toBeDefined();
      expect(utidCol?.type).toEqual({kind: 'int'});
      expect(utidCol?.type).toEqual(PerfettoSqlTypes.INT);

      const trackIdCol = selectedCols.find((c) => c.name === 'track_id');
      expect(trackIdCol).toBeDefined();
      expect(trackIdCol?.type).toEqual({kind: 'int'});
      expect(trackIdCol?.type).toEqual(PerfettoSqlTypes.INT);

      // Verify the column count increased by 2 (the 2 partition columns)
      expect(selectedCols.length).toBe(initialColCount + 2);
    });

    it('should handle partition columns that do not exist in input nodes', () => {
      // Create input nodes WITHOUT the partition columns
      const node1 = createMockNode({
        nodeId: 'node1',
        columns: [
          createColumnInfoWithType('id', PerfettoSqlTypes.INT),
          createColumnInfoWithType('ts', PerfettoSqlTypes.TIMESTAMP),
          createColumnInfoWithType('dur', PerfettoSqlTypes.DURATION),
        ],
      });
      const node2 = createMockNode({
        nodeId: 'node2',
        columns: [
          createColumnInfoWithType('id', PerfettoSqlTypes.INT),
          createColumnInfoWithType('ts', PerfettoSqlTypes.TIMESTAMP),
          createColumnInfoWithType('dur', PerfettoSqlTypes.DURATION),
        ],
      });

      // Create IntervalIntersectNode with partition columns that DON'T exist
      const intervalNode = makeIIN([node1, node2], ['utid', 'track_id']);

      const cols = intervalNode.finalCols;

      // Should still create partition columns, but with 'NA' type as fallback
      const utidCol = cols.find((c) => c.name === 'utid');
      expect(utidCol).toBeDefined();
      expect(utidCol?.type).toBeUndefined();
      expect(utidCol?.type).toBeUndefined();

      const trackIdCol = cols.find((c) => c.name === 'track_id');
      expect(trackIdCol).toBeDefined();
      expect(trackIdCol?.type).toBeUndefined();
      expect(trackIdCol?.type).toBeUndefined();
    });

    it('should use first input node type for partition columns when types differ', () => {
      // Create input nodes with DIFFERENT types for the same partition column
      const node1 = createMockNode({
        nodeId: 'node1',
        columns: [
          createColumnInfoWithType('id', PerfettoSqlTypes.INT),
          createColumnInfoWithType('ts', PerfettoSqlTypes.TIMESTAMP),
          createColumnInfoWithType('dur', PerfettoSqlTypes.DURATION),
          createColumnInfoWithType('utid', PerfettoSqlTypes.INT),
        ],
      });
      const node2 = createMockNode({
        nodeId: 'node2',
        columns: [
          createColumnInfoWithType('id', PerfettoSqlTypes.INT),
          createColumnInfoWithType('ts', PerfettoSqlTypes.TIMESTAMP),
          createColumnInfoWithType('dur', PerfettoSqlTypes.DURATION),
          createColumnInfoWithType('utid', PerfettoSqlTypes.STRING), // Different type!
        ],
      });

      const intervalNode = makeIIN([node1, node2], ['utid']);

      const cols = intervalNode.finalCols;

      // Should use the type from the FIRST input node
      const utidCol = cols.find((c) => c.name === 'utid');
      expect(utidCol).toBeDefined();
      expect(utidCol?.type).toEqual({kind: 'int'}); // From node1, not STRING from node2
      expect(utidCol?.type).toEqual(PerfettoSqlTypes.INT);
    });

    it('should handle empty input nodes gracefully', () => {
      // Edge case: What if an input node is undefined?
      const node1 = createMockNode({
        nodeId: 'node1',
        columns: [
          createColumnInfoWithType('id', PerfettoSqlTypes.INT),
          createColumnInfoWithType('ts', PerfettoSqlTypes.TIMESTAMP),
          createColumnInfoWithType('dur', PerfettoSqlTypes.DURATION),
        ],
      });

      // Test edge case: undefined input node (implementation handles this gracefully)
      const intervalNode = makeIIN([node1, undefined] as QueryNode[]);

      const cols = intervalNode.finalCols;

      // Should still have base columns
      expect(cols.find((c) => c.name === 'ts')).toBeDefined();
      expect(cols.find((c) => c.name === 'dur')).toBeDefined();
      expect(cols.find((c) => c.name === 'id_0')).toBeDefined();

      // Should NOT create columns for the undefined node
      expect(cols.find((c) => c.name === 'id_1')).toBeUndefined();
    });
  });

  describe('tsDurSource', () => {
    it('should have id column only when tsDurSource is a number', () => {
      const node1 = createMockIntervalNode('node1');
      const node2 = createMockIntervalNode('node2');

      // Test with default (intersection) - no id column
      const nodeDefault = makeIIN([node1, node2]);

      let cols = nodeDefault.finalCols;
      let colNames = cols.map((c) => c.name);

      expect(colNames).not.toContain('id'); // No id for intersection
      expect(colNames).toContain('ts');
      expect(colNames).toContain('dur');
      expect(colNames).toContain('id_0');
      expect(colNames).toContain('ts_0');
      expect(colNames).toContain('dur_0');
      expect(colNames).toContain('id_1');
      expect(colNames).toContain('ts_1');
      expect(colNames).toContain('dur_1');

      // Test with tsDurSource = 0 - has id column
      const nodeSource0 = makeIIN([node1, node2], undefined, 0);

      cols = nodeSource0.finalCols;
      colNames = cols.map((c) => c.name);

      expect(colNames).toContain('id'); // Has id when source is an input
      expect(colNames).toContain('ts');
      expect(colNames).toContain('dur');
      expect(colNames).toContain('id_0');
      expect(colNames).toContain('ts_0');
      expect(colNames).toContain('dur_0');
      expect(colNames).toContain('id_1');
      expect(colNames).toContain('ts_1');
      expect(colNames).toContain('dur_1');

      // Test with tsDurSource = 1 - has id column
      const nodeSource1 = makeIIN([node1, node2], undefined, 1);

      cols = nodeSource1.finalCols;
      colNames = cols.map((c) => c.name);

      expect(colNames).toContain('id'); // Has id when source is an input
      expect(colNames).toContain('ts');
      expect(colNames).toContain('dur');
      expect(colNames).toContain('id_0');
      expect(colNames).toContain('ts_0');
      expect(colNames).toContain('dur_0');
      expect(colNames).toContain('id_1');
      expect(colNames).toContain('ts_1');
      expect(colNames).toContain('dur_1');
    });

    it('should serialize and deserialize tsDurSource correctly', () => {
      const node1 = createMockIntervalNode('node1');
      const node2 = createMockIntervalNode('node2');

      const node = makeIIN([node1, node2], undefined, 1);

      const serialized = node.attrs;
      expect(serialized.tsDurSource).toBe(1);

      const restored = new IntervalIntersectNode(serialized, {});
      expect(restored.attrs.tsDurSource).toBe(1);
    });

    it('should include tsDurSource in clone', () => {
      const node1 = createMockIntervalNode('node1');
      const node2 = createMockIntervalNode('node2');

      const node = makeIIN([node1, node2], undefined, 0);

      const cloned = node.clone() as IntervalIntersectNode;

      expect((cloned as IntervalIntersectNode).attrs.tsDurSource).toBe(0);
    });

    it('should update id column type when tsDurSource changes to different input', () => {
      // Create two inputs with DIFFERENT id types
      const node1 = createMockNode({
        nodeId: 'node1',
        columns: intCols('id', 'ts', 'dur'),
      });
      const node2 = createMockNode({
        nodeId: 'node2',
        columns: [
          createColumnInfo('id', 'string'), // Input 1 has STRING id
          createColumnInfo('ts', 'int'),
          createColumnInfo('dur', 'int'),
        ],
      });

      const node = makeIIN([node1, node2], undefined, 0);

      // Check id type matches input 0
      let idCol = node.finalCols.find((c) => c.name === 'id');
      expect(idCol).toBeDefined();
      expect(idCol?.type).toEqual({kind: 'int'});

      // Change to input 1
      node.attrs.tsDurSource = 1;

      // Check id type should now match input 1
      idCol = node.finalCols.find((c) => c.name === 'id');
      expect(idCol).toBeDefined();
      expect(idCol?.type).toEqual({kind: 'string'});
    });

    it('should propagate id column type change to downstream ModifyColumnsNode when tsDurSource changes', () => {
      // Create two inputs with DIFFERENT id types
      const node1 = createMockNode({
        nodeId: 'node1',
        columns: [
          createColumnInfoWithType('id', PerfettoSqlTypes.INT),
          createColumnInfoWithType('ts', PerfettoSqlTypes.TIMESTAMP),
          createColumnInfoWithType('dur', PerfettoSqlTypes.DURATION),
        ],
      });
      const node2 = createMockNode({
        nodeId: 'node2',
        columns: [
          createColumnInfoWithType('id', PerfettoSqlTypes.STRING),
          createColumnInfoWithType('ts', PerfettoSqlTypes.TIMESTAMP),
          createColumnInfoWithType('dur', PerfettoSqlTypes.DURATION),
        ],
      });

      // Create IntervalIntersectNode with tsDurSource = 0
      const intervalNode = makeIIN([node1, node2], undefined, 0);

      // Create ModifyColumnsNode downstream
      const modifyNode = new ModifyColumnsNode({selectedColumns: []}, {});
      modifyNode.primaryInput = intervalNode;
      intervalNode.nextNodes.push(modifyNode);
      modifyNode.onPrevNodesUpdated();

      // Verify id column has INT type from input 0
      let idCol = modifyNode.attrs.selectedColumns.find((c) => c.name === 'id');
      expect(idCol).toBeDefined();
      expect(idCol?.type).toEqual({kind: 'int'});

      // Change tsDurSource to input 1
      intervalNode.attrs.tsDurSource = 1;
      notifyNextNodes(intervalNode);

      // Verify id column type updated to STRING from input 1
      idCol = modifyNode.attrs.selectedColumns.find((c) => c.name === 'id');
      expect(idCol).toBeDefined();
      expect(idCol?.type).toEqual({kind: 'string'});
    });

    it('should preserve user-modified type in ModifyColumnsNode when tsDurSource changes', () => {
      // Create two inputs with DIFFERENT id types
      const node1 = createMockNode({
        nodeId: 'node1',
        columns: [
          createColumnInfoWithType('id', PerfettoSqlTypes.INT),
          createColumnInfoWithType('ts', PerfettoSqlTypes.TIMESTAMP),
          createColumnInfoWithType('dur', PerfettoSqlTypes.DURATION),
        ],
      });
      const node2 = createMockNode({
        nodeId: 'node2',
        columns: [
          createColumnInfoWithType('id', PerfettoSqlTypes.STRING),
          createColumnInfoWithType('ts', PerfettoSqlTypes.TIMESTAMP),
          createColumnInfoWithType('dur', PerfettoSqlTypes.DURATION),
        ],
      });

      // Create IntervalIntersectNode with tsDurSource = 0
      const intervalNode = makeIIN([node1, node2], undefined, 0);

      // Create ModifyColumnsNode downstream
      const modifyNode = new ModifyColumnsNode({selectedColumns: []}, {});
      modifyNode.primaryInput = intervalNode;
      intervalNode.nextNodes.push(modifyNode);
      modifyNode.onPrevNodesUpdated();

      // Simulate user manually changing the id column type to a custom type
      const customType: PerfettoSqlType = {kind: 'double'};
      const idColIndex = modifyNode.attrs.selectedColumns.findIndex(
        (c) => c.name === 'id',
      );
      modifyNode.attrs.selectedColumns[idColIndex] = {
        ...modifyNode.attrs.selectedColumns[idColIndex],
        type: customType,
        typeUserModified: true,
      };

      // Change tsDurSource to input 1
      intervalNode.attrs.tsDurSource = 1;
      notifyNextNodes(intervalNode);

      // Verify user-modified type is PRESERVED (not overwritten by STRING)
      const idCol = modifyNode.attrs.selectedColumns.find(
        (c) => c.name === 'id',
      );
      expect(idCol).toBeDefined();
      expect(idCol?.type).toEqual(customType);
      expect(idCol?.typeUserModified).toBe(true);
    });
  });
});
