// 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 {UnionNode} from './union_node';
import type {QueryNode} from '../../query_node';
import {
  createMockNodeWithStructuredQuery,
  createColumnInfo,
  expectValidationError,
  expectValidationSuccess,
} from '../testing/test_utils';

describe('UnionNode', () => {
  describe('constructor', () => {
    it('should initialize with default values', () => {
      const node1 = createMockNodeWithStructuredQuery('node1', [
        createColumnInfo('id', 'int'),
        createColumnInfo('name', 'string'),
      ]);
      const node2 = createMockNodeWithStructuredQuery('node2', [
        createColumnInfo('id', 'int'),
        createColumnInfo('name', 'string'),
      ]);

      const unionNode = new UnionNode(
        {
          inputNodes: [node1, node2],
          selectedColumns: [
            createColumnInfo('id', 'int'),
            createColumnInfo('name', 'string'),
          ],
        },
        {},
      );

      expect(unionNode.context.autoExecute).toBeUndefined();
      expect(unionNode.attrs.selectedColumns.length).toBe(2);
      expect(unionNode.secondaryInputs.min).toBe(2);
      expect(unionNode.secondaryInputs.max).toBe('unbounded');
    });

    it('should initialize connections from inputNodes', () => {
      const node1 = createMockNodeWithStructuredQuery('node1', []);
      const node2 = createMockNodeWithStructuredQuery('node2', []);
      const node3 = createMockNodeWithStructuredQuery('node3', []);

      const unionNode = new UnionNode(
        {
          inputNodes: [node1, node2, node3],
          selectedColumns: [],
        },
        {},
      );

      expect(unionNode.secondaryInputs.connections.size).toBe(3);
      expect(unionNode.secondaryInputs.connections.get(0)).toBe(node1);
      expect(unionNode.secondaryInputs.connections.get(1)).toBe(node2);
      expect(unionNode.secondaryInputs.connections.get(2)).toBe(node3);
    });
  });

  describe('getCommonColumns', () => {
    it('should return empty array when there are no input nodes', () => {
      const unionNode = new UnionNode(
        {
          inputNodes: [],
          selectedColumns: [],
        },
        {},
      );

      expect(unionNode['getCommonColumns']()).toEqual([]);
    });

    it('should return all columns from single input', () => {
      const node1 = createMockNodeWithStructuredQuery('node1', [
        createColumnInfo('id', 'int'),
        createColumnInfo('name', 'string'),
      ]);

      const unionNode = new UnionNode(
        {
          inputNodes: [node1],
          selectedColumns: [],
        },
        {},
      );

      const commonCols = unionNode['getCommonColumns']();
      expect(commonCols.length).toBe(2);
      expect(commonCols.map((c) => c.name)).toEqual(['id', 'name']);
    });

    it('should return only common columns between two inputs', () => {
      const node1 = createMockNodeWithStructuredQuery('node1', [
        createColumnInfo('id', 'int'),
        createColumnInfo('name', 'string'),
        createColumnInfo('ts', 'int'),
      ]);
      const node2 = createMockNodeWithStructuredQuery('node2', [
        createColumnInfo('id', 'int'),
        createColumnInfo('name', 'string'),
        createColumnInfo('value', 'int'),
      ]);

      const unionNode = new UnionNode(
        {
          inputNodes: [node1, node2],
          selectedColumns: [],
        },
        {},
      );

      const commonCols = unionNode['getCommonColumns']();
      expect(commonCols.length).toBe(2);
      expect(commonCols.map((c) => c.name)).toEqual(['id', 'name']);
    });

    it('should return only columns present in all inputs', () => {
      const node1 = createMockNodeWithStructuredQuery('node1', [
        createColumnInfo('id', 'int'),
        createColumnInfo('name', 'string'),
        createColumnInfo('ts', 'int'),
      ]);
      const node2 = createMockNodeWithStructuredQuery('node2', [
        createColumnInfo('id', 'int'),
        createColumnInfo('name', 'string'),
        createColumnInfo('value', 'int'),
      ]);
      const node3 = createMockNodeWithStructuredQuery('node3', [
        createColumnInfo('id', 'int'),
        createColumnInfo('other', 'string'),
      ]);

      const unionNode = new UnionNode(
        {
          inputNodes: [node1, node2, node3],
          selectedColumns: [],
        },
        {},
      );

      const commonCols = unionNode['getCommonColumns']();
      expect(commonCols.length).toBe(1);
      expect(commonCols.map((c) => c.name)).toEqual(['id']);
    });

    it('should return empty array when there are no common columns', () => {
      const node1 = createMockNodeWithStructuredQuery('node1', [
        createColumnInfo('id', 'int'),
        createColumnInfo('name', 'string'),
      ]);
      const node2 = createMockNodeWithStructuredQuery('node2', [
        createColumnInfo('value', 'int'),
        createColumnInfo('ts', 'int'),
      ]);

      const unionNode = new UnionNode(
        {
          inputNodes: [node1, node2],
          selectedColumns: [],
        },
        {},
      );

      const commonCols = unionNode['getCommonColumns']();
      expect(commonCols.length).toBe(0);
    });

    it('should set all columns as checked by default', () => {
      const node1 = createMockNodeWithStructuredQuery('node1', [
        createColumnInfo('id', 'int'),
        createColumnInfo('name', 'string'),
      ]);
      const node2 = createMockNodeWithStructuredQuery('node2', [
        createColumnInfo('id', 'int'),
        createColumnInfo('name', 'string'),
      ]);

      const unionNode = new UnionNode(
        {
          inputNodes: [node1, node2],
          selectedColumns: [],
        },
        {},
      );

      const commonCols = unionNode['getCommonColumns']();
      expect(commonCols.every((c) => c.checked === true)).toBe(true);
    });
  });

  describe('finalCols', () => {
    it('should return only checked columns', () => {
      const node1 = createMockNodeWithStructuredQuery('node1', [
        createColumnInfo('id', 'int'),
        createColumnInfo('name', 'string'),
      ]);
      const node2 = createMockNodeWithStructuredQuery('node2', [
        createColumnInfo('id', 'int'),
        createColumnInfo('name', 'string'),
      ]);

      const unionNode = new UnionNode(
        {
          inputNodes: [node1, node2],
          selectedColumns: [
            createColumnInfo('id', 'int'),
            {...createColumnInfo('name', 'string'), checked: false},
          ],
        },
        {},
      );

      const finalCols = unionNode.finalCols;
      expect(finalCols.length).toBe(1);
      expect(finalCols[0].name).toBe('id');
    });

    it('should return empty array when no columns are checked', () => {
      const node1 = createMockNodeWithStructuredQuery('node1', [
        createColumnInfo('id', 'int'),
      ]);
      const node2 = createMockNodeWithStructuredQuery('node2', [
        createColumnInfo('id', 'int'),
      ]);

      const unionNode = new UnionNode(
        {
          inputNodes: [node1, node2],
          selectedColumns: [{...createColumnInfo('id', 'int'), checked: false}],
        },
        {},
      );

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

  describe('onPrevNodesUpdated', () => {
    it('should update selectedColumns based on new common columns', () => {
      const node1 = createMockNodeWithStructuredQuery('node1', [
        createColumnInfo('id', 'int'),
        createColumnInfo('name', 'string'),
      ]);
      const node2 = createMockNodeWithStructuredQuery('node2', [
        createColumnInfo('id', 'int'),
        createColumnInfo('name', 'string'),
      ]);

      const unionNode = new UnionNode(
        {
          inputNodes: [node1, node2],
          selectedColumns: [
            createColumnInfo('id', 'int'),
            {...createColumnInfo('name', 'string'), checked: false},
          ],
        },
        {},
      );

      // Update node2 to remove 'name'
      const updatedNode2 = createMockNodeWithStructuredQuery('node2', [
        createColumnInfo('id', 'int'),
        createColumnInfo('value', 'int'),
      ]);
      unionNode.secondaryInputs.connections.set(1, updatedNode2);

      unionNode.onPrevNodesUpdated();

      // Only 'id' should remain as it's the only common column
      expect(unionNode.attrs.selectedColumns.length).toBe(1);
      expect(unionNode.attrs.selectedColumns[0].name).toBe('id');
      // 'id' should preserve its checked status (true)
      expect(unionNode.attrs.selectedColumns[0].checked).toBe(true);
    });

    it('should preserve checked status for columns that still exist', () => {
      const node1 = createMockNodeWithStructuredQuery('node1', [
        createColumnInfo('id', 'int'),
        createColumnInfo('name', 'string'),
        createColumnInfo('ts', 'int'),
      ]);
      const node2 = createMockNodeWithStructuredQuery('node2', [
        createColumnInfo('id', 'int'),
        createColumnInfo('name', 'string'),
        createColumnInfo('ts', 'int'),
      ]);

      const unionNode = new UnionNode(
        {
          inputNodes: [node1, node2],
          selectedColumns: [
            createColumnInfo('id', 'int'),
            {...createColumnInfo('name', 'string'), checked: false},
            createColumnInfo('ts', 'int'),
          ],
        },
        {},
      );

      unionNode.onPrevNodesUpdated();

      // All columns should still exist
      expect(unionNode.attrs.selectedColumns.length).toBe(3);

      // Check that checked status is preserved
      const idCol = unionNode.attrs.selectedColumns.find(
        (c) => c.name === 'id',
      );
      const nameCol = unionNode.attrs.selectedColumns.find(
        (c) => c.name === 'name',
      );
      const tsCol = unionNode.attrs.selectedColumns.find(
        (c) => c.name === 'ts',
      );

      expect(idCol?.checked).toBe(true);
      expect(nameCol?.checked).toBe(false);
      expect(tsCol?.checked).toBe(true);
    });
  });

  describe('validate', () => {
    it('should fail when there are fewer than 2 input nodes', () => {
      const node1 = createMockNodeWithStructuredQuery('node1', [
        createColumnInfo('id', 'int'),
      ]);

      const unionNode = new UnionNode(
        {
          inputNodes: [node1],
          selectedColumns: [],
        },
        {},
      );

      expectValidationError(unionNode, 'at least two sources');
    });

    it('should fail when there are no common columns', () => {
      const node1 = createMockNodeWithStructuredQuery('node1', [
        createColumnInfo('id', 'int'),
      ]);
      const node2 = createMockNodeWithStructuredQuery('node2', [
        createColumnInfo('value', 'int'),
      ]);

      const unionNode = new UnionNode(
        {
          inputNodes: [node1, node2],
          selectedColumns: [],
        },
        {},
      );

      expectValidationError(unionNode, 'common columns');
    });

    it('should fail when input nodes have disconnected inputs', () => {
      const node1 = createMockNodeWithStructuredQuery('node1', [
        createColumnInfo('id', 'int'),
      ]);

      const unionNode = new UnionNode(
        {
          inputNodes: [node1],
          selectedColumns: [],
        },
        {},
      );

      // Manually add undefined to connections
      unionNode.secondaryInputs.connections.set(
        1,
        undefined as unknown as QueryNode,
      );

      expectValidationError(unionNode, 'disconnected inputs');
    });

    it('should pass validation with valid inputs', () => {
      const node1 = createMockNodeWithStructuredQuery('node1', [
        createColumnInfo('id', 'int'),
        createColumnInfo('name', 'string'),
      ]);
      const node2 = createMockNodeWithStructuredQuery('node2', [
        createColumnInfo('id', 'int'),
        createColumnInfo('name', 'string'),
      ]);

      const unionNode = new UnionNode(
        {
          inputNodes: [node1, node2],
          selectedColumns: [
            createColumnInfo('id', 'int'),
            createColumnInfo('name', 'string'),
          ],
        },
        {},
      );

      expectValidationSuccess(unionNode);
    });

    it('should clear previous errors on successful validation', () => {
      const node1 = createMockNodeWithStructuredQuery('node1', [
        createColumnInfo('id', 'int'),
      ]);
      const node2 = createMockNodeWithStructuredQuery('node2', [
        createColumnInfo('id', 'int'),
      ]);

      const unionNode = new UnionNode(
        {
          inputNodes: [node1, node2],
          selectedColumns: [createColumnInfo('id', 'int')],
        },
        {},
      );

      // First validate with error
      unionNode.secondaryInputs.connections.clear();
      expect(unionNode.validate()).toBe(false);
      expect(unionNode.context.issues?.queryError).toBeDefined();

      // Restore connections and validate again
      unionNode.secondaryInputs.connections.set(0, node1);
      unionNode.secondaryInputs.connections.set(1, node2);
      expectValidationSuccess(unionNode);
      expect(unionNode.context.issues?.queryError).toBeUndefined();
    });
  });

  describe('getTitle', () => {
    it('should return "Union"', () => {
      const unionNode = new UnionNode(
        {
          inputNodes: [],
          selectedColumns: [],
        },
        {},
      );

      expect(unionNode.getTitle()).toBe('Union');
    });
  });

  describe('clone', () => {
    it('should create a deep copy of the node', () => {
      const node1 = createMockNodeWithStructuredQuery('node1', [
        createColumnInfo('id', 'int'),
      ]);
      const node2 = createMockNodeWithStructuredQuery('node2', [
        createColumnInfo('id', 'int'),
      ]);

      const unionNode = new UnionNode(
        {
          inputNodes: [node1, node2],
          selectedColumns: [createColumnInfo('id', 'int')],
        },
        {},
      );
      const cloned = unionNode.clone() as UnionNode;

      expect(cloned).not.toBe(unionNode);
      // Clone doesn't copy input connections (those are graph topology)
      expect((cloned as UnionNode).attrs.selectedColumns.length).toBe(1);
    });

    it('should not share state with original', () => {
      const node1 = createMockNodeWithStructuredQuery('node1', [
        createColumnInfo('id', 'int'),
      ]);
      const node2 = createMockNodeWithStructuredQuery('node2', [
        createColumnInfo('id', 'int'),
      ]);

      const unionNode = new UnionNode(
        {
          inputNodes: [node1, node2],
          selectedColumns: [createColumnInfo('id', 'int')],
        },
        {},
      );

      const cloned = unionNode.clone() as UnionNode;

      // Modify the cloned state
      (cloned as UnionNode).attrs.selectedColumns[0].checked = false;

      // Original should not be affected
      expect(unionNode.attrs.selectedColumns[0].checked).toBe(true);
    });
  });

  describe('getStructuredQuery', () => {
    it('should return undefined when there are fewer than 2 inputs', () => {
      const node1 = createMockNodeWithStructuredQuery('node1', [
        createColumnInfo('id', 'int'),
      ]);

      const unionNode = new UnionNode(
        {
          inputNodes: [node1],
          selectedColumns: [],
        },
        {},
      );

      expect(unionNode.getStructuredQuery()).toBeUndefined();
    });

    it('should return undefined when there are no checked columns', () => {
      const node1 = createMockNodeWithStructuredQuery('node1', [
        createColumnInfo('id', 'int'),
      ]);
      const node2 = createMockNodeWithStructuredQuery('node2', [
        createColumnInfo('id', 'int'),
      ]);

      const unionNode = new UnionNode(
        {
          inputNodes: [node1, node2],
          selectedColumns: [{...createColumnInfo('id', 'int'), checked: false}],
        },
        {},
      );

      expect(unionNode.getStructuredQuery()).toBeUndefined();
    });

    it('should return undefined when any input node is undefined', () => {
      const node1 = createMockNodeWithStructuredQuery('node1', [
        createColumnInfo('id', 'int'),
      ]);

      const unionNode = new UnionNode(
        {
          inputNodes: [node1],
          selectedColumns: [createColumnInfo('id', 'int')],
        },
        {},
      );

      unionNode.secondaryInputs.connections.set(
        1,
        undefined as unknown as QueryNode,
      );

      expect(unionNode.getStructuredQuery()).toBeUndefined();
    });

    it('should create union query with wrapped selects for common columns', () => {
      const node1 = createMockNodeWithStructuredQuery('node1', [
        createColumnInfo('id', 'int'),
        createColumnInfo('name', 'string'),
        createColumnInfo('ts', 'int'),
      ]);
      const node2 = createMockNodeWithStructuredQuery('node2', [
        createColumnInfo('id', 'int'),
        createColumnInfo('name', 'string'),
        createColumnInfo('value', 'int'),
      ]);

      const unionNode = new UnionNode(
        {
          inputNodes: [node1, node2],
          selectedColumns: [
            createColumnInfo('id', 'int'),
            createColumnInfo('name', 'string'),
          ],
        },
        {},
      );

      const sq = unionNode.getStructuredQuery();

      expect(sq).toBeDefined();
      expect(sq?.experimentalUnion).toBeDefined();
      expect(sq?.experimentalUnion?.useUnionAll).toBe(true);
      expect(sq?.experimentalUnion?.queries?.length).toBe(2);
    });

    it('should only select checked common columns', () => {
      const node1 = createMockNodeWithStructuredQuery('node1', [
        createColumnInfo('id', 'int'),
        createColumnInfo('name', 'string'),
        createColumnInfo('ts', 'int'),
      ]);
      const node2 = createMockNodeWithStructuredQuery('node2', [
        createColumnInfo('id', 'int'),
        createColumnInfo('name', 'string'),
        createColumnInfo('ts', 'int'),
      ]);

      const unionNode = new UnionNode(
        {
          inputNodes: [node1, node2],
          selectedColumns: [
            createColumnInfo('id', 'int'),
            {...createColumnInfo('name', 'string'), checked: false}, // unchecked
            createColumnInfo('ts', 'int'),
          ],
        },
        {},
      );

      const sq = unionNode.getStructuredQuery();

      expect(sq).toBeDefined();
      // The union queries reference queries with SELECT that only includes id and ts
      const queries = sq?.experimentalUnion?.queries;
      expect(queries?.length).toBe(2);

      // Each union query is a reference to the wrapped SELECT query.
      // Verify:
      // 1. Each query has an ID and innerQueryId
      // 2. The innerQueryIds are unique (different SELECT queries for each input)
      const queryIds = new Set<string>();
      const innerQueryIds = new Set<string>();
      queries?.forEach((query) => {
        expect(query.id).toBeDefined();
        expect(query.innerQueryId).toBeDefined();
        queryIds.add(query.id!);
        innerQueryIds.add(query.innerQueryId!);
      });
      // Should have 2 unique query IDs and 2 unique innerQueryIds
      expect(queryIds.size).toBe(2);
      expect(innerQueryIds.size).toBe(2);
    });
  });

  describe('serializeState', () => {
    it('should serialize all input node IDs and selected columns', () => {
      const node1 = createMockNodeWithStructuredQuery('node1', [
        createColumnInfo('id', 'int'),
      ]);
      const node2 = createMockNodeWithStructuredQuery('node2', [
        createColumnInfo('id', 'int'),
      ]);

      const unionNode = new UnionNode(
        {
          inputNodes: [node1, node2],
          selectedColumns: [createColumnInfo('id', 'int')],
        },
        {},
      );
      expect(unionNode.attrs.selectedColumns.length).toBe(1);
      expect(unionNode.attrs.selectedColumns[0].name).toBe('id');
    });
  });

  describe('deserialize', () => {
    it('should restore state via constructor', () => {
      const node = new UnionNode(
        {selectedColumns: [createColumnInfo('id', 'int')]},
        {},
      );

      expect(node.attrs.selectedColumns.length).toBe(1);
      expect(node.attrs.selectedColumns[0].name).toBe('id');
    });
  });
});
