#!/usr/bin/env python3
# Copyright (C) 2026 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.

from python.generators.diff_tests.testing import DataPath
from python.generators.diff_tests.testing import Csv
from python.generators.diff_tests.testing import DiffTestBlueprint
from python.generators.diff_tests.testing import TestSuite


class TreeFilter(TestSuite):
  """Tests for tree filtering with reparenting behavior."""

  def test_filter_leaf_nodes(self):
    """Filtering leaf nodes should just remove them."""
    return DiffTestBlueprint(
        trace=DataPath('counters.json'),
        query="""
          INCLUDE PERFETTO MODULE std.trees.table_conversion;
          INCLUDE PERFETTO MODULE std.trees.filter;

          CREATE PERFETTO TABLE input_tree AS
          SELECT 1 AS id, NULL AS parent_id, 'root' AS name
          UNION ALL SELECT 2, 1, 'keep'
          UNION ALL SELECT 3, 1, 'remove'
          UNION ALL SELECT 4, 2, 'child';

          -- Composition: table → tree → filter (constraint-based) → table
          SELECT _tree_id, _tree_parent_id, id, name
          FROM _tree_to_table!(
            _tree_filter(
              _tree_from_table!((SELECT * FROM input_tree), (name)),
              _tree_where(_tree_constraint('name', '!=', 'remove'))
            ),
            (name)
          )
          ORDER BY id;
        """,
        out=Csv("""
        "_tree_id","_tree_parent_id","id","name"
        0,"[NULL]",1,"root"
        1,0,2,"keep"
        2,1,4,"child"
        """))

  def test_filter_intermediate_with_reparenting(self):
    """Filtering intermediate nodes should reparent children to grandparent."""
    return DiffTestBlueprint(
        trace=DataPath('counters.json'),
        query="""
          INCLUDE PERFETTO MODULE std.trees.table_conversion;
          INCLUDE PERFETTO MODULE std.trees.filter;

          CREATE PERFETTO TABLE input_tree AS
          SELECT 1 AS id, NULL AS parent_id, 1 AS depth
          UNION ALL SELECT 2, 1, 2   -- Will be filtered
          UNION ALL SELECT 3, 2, 3   -- Should be reparented to node 1
          UNION ALL SELECT 4, 1, 2;  -- Sibling, unaffected

          SELECT _tree_id, _tree_parent_id, id, depth
          FROM _tree_to_table!(
            _tree_filter(
              _tree_from_table!((SELECT * FROM input_tree), (depth)),
              _tree_where(_tree_constraint('depth', '!=', 2))
            ),
            (depth)
          )
          ORDER BY id;
        """,
        out=Csv("""
        "_tree_id","_tree_parent_id","id","depth"
        0,"[NULL]",1,1
        1,0,3,3
        """))

  def test_filter_all_nodes(self):
    """Filtering all nodes should result in empty tree."""
    return DiffTestBlueprint(
        trace=DataPath('counters.json'),
        query="""
          INCLUDE PERFETTO MODULE std.trees.table_conversion;
          INCLUDE PERFETTO MODULE std.trees.filter;

          CREATE PERFETTO TABLE input_tree AS
          SELECT 1 AS id, NULL AS parent_id, 'a' AS tag
          UNION ALL SELECT 2, 1, 'b'
          UNION ALL SELECT 3, 2, 'c';

          SELECT _tree_id, _tree_parent_id, id, tag
          FROM _tree_to_table!(
            _tree_filter(
              _tree_from_table!((SELECT * FROM input_tree), (tag)),
              _tree_where(_tree_constraint('tag', '=', 'nonexistent'))
            ),
            (tag)
          )
          ORDER BY id;
        """,
        out=Csv("""
        "_tree_id","_tree_parent_id","id","tag"
        """))

  def test_filter_root_nodes(self):
    """Filtering root nodes promotes children to new roots."""
    return DiffTestBlueprint(
        trace=DataPath('counters.json'),
        query="""
          INCLUDE PERFETTO MODULE std.trees.table_conversion;
          INCLUDE PERFETTO MODULE std.trees.filter;

          CREATE PERFETTO TABLE input_tree AS
          SELECT 1 AS id, NULL AS parent_id, 0 AS level
          UNION ALL SELECT 2, 1, 1          -- Child of root 1, will become new root
          UNION ALL SELECT 3, NULL, 0       -- Another root (kept)
          UNION ALL SELECT 4, 3, 1;         -- Child of root 3 (kept)

          SELECT _tree_id, _tree_parent_id, id, level
          FROM _tree_to_table!(
            _tree_filter(
              _tree_from_table!((SELECT * FROM input_tree), (level)),
              _tree_where(_tree_constraint('id', '!=', 1))  -- Filter out root 1
            ),
            (level)
          )
          ORDER BY id;
        """,
        out=Csv("""
        "_tree_id","_tree_parent_id","id","level"
        0,"[NULL]",2,1
        1,"[NULL]",3,0
        2,1,4,1
        """))

  def test_filter_multiple_constraints_different_columns(self):
    """Multiple constraints on different columns (AND logic)."""
    return DiffTestBlueprint(
        trace=DataPath('counters.json'),
        query="""
          INCLUDE PERFETTO MODULE std.trees.table_conversion;
          INCLUDE PERFETTO MODULE std.trees.filter;

          CREATE PERFETTO TABLE input_tree AS
          SELECT 1 AS id, NULL AS parent_id, 'a' AS tag, 10 AS val
          UNION ALL SELECT 2, 1, 'b', 20
          UNION ALL SELECT 3, 1, 'c', 30
          UNION ALL SELECT 4, 2, 'b', 40
          UNION ALL SELECT 5, 3, 'c', 50;

          -- Keep only nodes where tag != 'a' AND val < 45
          SELECT _tree_id, _tree_parent_id, id, tag, val
          FROM _tree_to_table!(
            _tree_filter(
              _tree_from_table!((SELECT * FROM input_tree), (tag, val)),
              _tree_where(
                _tree_constraint('tag', '!=', 'a'),
                _tree_constraint('val', '<', 45)
              )
            ),
            (tag, val)
          )
          ORDER BY id;
        """,
        out=Csv("""
        "_tree_id","_tree_parent_id","id","tag","val"
        0,"[NULL]",2,"b",20
        1,"[NULL]",3,"c",30
        2,0,4,"b",40
        """))

  def test_filter_multiple_constraints_same_column_range(self):
    """Multiple constraints on the same column forming a range."""
    return DiffTestBlueprint(
        trace=DataPath('counters.json'),
        query="""
          INCLUDE PERFETTO MODULE std.trees.table_conversion;
          INCLUDE PERFETTO MODULE std.trees.filter;

          CREATE PERFETTO TABLE input_tree AS
          SELECT 1 AS id, NULL AS parent_id, 1 AS depth
          UNION ALL SELECT 2, 1, 2
          UNION ALL SELECT 3, 2, 3
          UNION ALL SELECT 4, 3, 4
          UNION ALL SELECT 5, 4, 5;

          -- Keep only nodes where 2 <= depth <= 4
          SELECT _tree_id, _tree_parent_id, id, depth
          FROM _tree_to_table!(
            _tree_filter(
              _tree_from_table!((SELECT * FROM input_tree), (depth)),
              _tree_where(
                _tree_constraint('depth', '>=', 2),
                _tree_constraint('depth', '<=', 4)
              )
            ),
            (depth)
          )
          ORDER BY id;
        """,
        out=Csv("""
        "_tree_id","_tree_parent_id","id","depth"
        0,"[NULL]",2,2
        1,0,3,3
        2,1,4,4
        """))

  def test_filter_multiple_constraints_filters_everything(self):
    """Multiple constraints whose intersection is empty."""
    return DiffTestBlueprint(
        trace=DataPath('counters.json'),
        query="""
          INCLUDE PERFETTO MODULE std.trees.table_conversion;
          INCLUDE PERFETTO MODULE std.trees.filter;

          CREATE PERFETTO TABLE input_tree AS
          SELECT 1 AS id, NULL AS parent_id, 'x' AS tag, 10 AS val
          UNION ALL SELECT 2, 1, 'y', 20
          UNION ALL SELECT 3, 2, 'z', 30;

          -- tag = 'x' AND val > 100: no node satisfies both
          SELECT _tree_id, _tree_parent_id, id, tag, val
          FROM _tree_to_table!(
            _tree_filter(
              _tree_from_table!((SELECT * FROM input_tree), (tag, val)),
              _tree_where(
                _tree_constraint('tag', '=', 'x'),
                _tree_constraint('val', '>', 100)
              )
            ),
            (tag, val)
          )
          ORDER BY id;
        """,
        out=Csv("""
        "_tree_id","_tree_parent_id","id","tag","val"
        """))

  def test_filter_multiple_constraints_with_reparenting(self):
    """Multiple constraints that remove intermediate nodes, causing reparenting."""
    return DiffTestBlueprint(
        trace=DataPath('counters.json'),
        query="""
          INCLUDE PERFETTO MODULE std.trees.table_conversion;
          INCLUDE PERFETTO MODULE std.trees.filter;

          CREATE PERFETTO TABLE input_tree AS
          SELECT 1 AS id, NULL AS parent_id, 'root' AS name, 0 AS val
          UNION ALL SELECT 2, 1, 'mid1', 10
          UNION ALL SELECT 3, 2, 'mid2', 20
          UNION ALL SELECT 4, 3, 'leaf', 30;

          -- Keep name != 'mid1' AND name != 'mid2': removes both
          -- intermediates, reparenting leaf -> root
          SELECT _tree_id, _tree_parent_id, id, name, val
          FROM _tree_to_table!(
            _tree_filter(
              _tree_from_table!((SELECT * FROM input_tree), (name, val)),
              _tree_where(
                _tree_constraint('name', '!=', 'mid1'),
                _tree_constraint('name', '!=', 'mid2')
              )
            ),
            (name, val)
          )
          ORDER BY id;
        """,
        out=Csv("""
        "_tree_id","_tree_parent_id","id","name","val"
        0,"[NULL]",1,"root",0
        1,0,4,"leaf",30
        """))

  def test_filter_chained_filters(self):
    """Two _tree_filter calls chained: first removes by tag, second by val."""
    return DiffTestBlueprint(
        trace=DataPath('counters.json'),
        query="""
          INCLUDE PERFETTO MODULE std.trees.table_conversion;
          INCLUDE PERFETTO MODULE std.trees.filter;

          CREATE PERFETTO TABLE input_tree AS
          SELECT 1 AS id, NULL AS parent_id, 'a' AS tag, 10 AS val
          UNION ALL SELECT 2, 1, 'b', 20
          UNION ALL SELECT 3, 1, 'a', 30
          UNION ALL SELECT 4, 2, 'b', 40
          UNION ALL SELECT 5, 3, 'c', 50
          UNION ALL SELECT 6, 4, 'c', 60;

          -- First filter: remove tag='a' (nodes 1,3).
          --   Node 2 becomes root, node 5 (child of 3) reparented to root.
          -- Second filter: remove val >= 50 (node 5, 6).
          --   Node 6 (child of 4) is removed, node 5 already gone.
          SELECT _tree_id, _tree_parent_id, id, tag, val
          FROM _tree_to_table!(
            _tree_filter(
              _tree_filter(
                _tree_from_table!((SELECT * FROM input_tree), (tag, val)),
                _tree_where(_tree_constraint('tag', '!=', 'a'))
              ),
              _tree_where(_tree_constraint('val', '<', 50))
            ),
            (tag, val)
          )
          ORDER BY id;
        """,
        out=Csv("""
        "_tree_id","_tree_parent_id","id","tag","val"
        0,"[NULL]",2,"b",20
        1,0,4,"b",40
        """))

  def test_filter_three_chained_filters_with_reparenting(self):
    """Three chained _tree_filter calls progressively pruning a deep tree.

    Tree: 1(root) -> {2(mid), 3(skip)},
          2 -> {4(keep), 5(skip)}, 3 -> {6(leaf)},
          4 -> {7(deep)}, 5 -> {8(deep)}, 7 -> {9(bottom)}

    Filter 1 (tag != 'skip'): removes 3,5; reparents 6->1, 8->2.
    Filter 2 (depth <= 4): removes 9.
    Filter 3 (val > 15): removes 1,2; reparents 4->root, 6->root, 8->root.
    Result: 4(root), 6(root), 7(child of 4), 8(root).
    """
    return DiffTestBlueprint(
        trace=DataPath('counters.json'),
        query="""
          INCLUDE PERFETTO MODULE std.trees.table_conversion;
          INCLUDE PERFETTO MODULE std.trees.filter;

          CREATE PERFETTO TABLE input_tree AS
          SELECT 1 AS id, NULL AS parent_id, 'root' AS tag, 0 AS val, 1 AS depth
          UNION ALL SELECT 2, 1, 'mid', 10, 2
          UNION ALL SELECT 3, 1, 'skip', 20, 2
          UNION ALL SELECT 4, 2, 'keep', 30, 3
          UNION ALL SELECT 5, 2, 'skip', 40, 3
          UNION ALL SELECT 6, 3, 'leaf', 50, 3
          UNION ALL SELECT 7, 4, 'deep', 60, 4
          UNION ALL SELECT 8, 5, 'deep', 70, 4
          UNION ALL SELECT 9, 7, 'bottom', 80, 5;

          SELECT _tree_id, _tree_parent_id, id, tag, val, depth
          FROM _tree_to_table!(
            _tree_filter(
              _tree_filter(
                _tree_filter(
                  _tree_from_table!(
                    (SELECT * FROM input_tree), (tag, val, depth)),
                  _tree_where(_tree_constraint('tag', '!=', 'skip'))
                ),
                _tree_where(_tree_constraint('depth', '<=', 4))
              ),
              _tree_where(_tree_constraint('val', '>', 15))
            ),
            (tag, val, depth)
          )
          ORDER BY id;
        """,
        out=Csv("""
        "_tree_id","_tree_parent_id","id","tag","val","depth"
        0,"[NULL]",4,"keep",30,3
        1,"[NULL]",6,"leaf",50,3
        2,0,7,"deep",60,4
        3,"[NULL]",8,"deep",70,4
        """))

  def test_filter_constraint_value_outlives_sqlite_register(self):
    """Constraint values built from expressions (e.g. printf) must be
    deep-copied: SQLite is free to reuse the register backing them once
    __intrinsic_tree_constraint returns, so the FilterConstraint cannot keep a
    borrowed pointer. Regression test for a heap-use-after-free where the
    constraint borrowed SQLite-owned string memory."""
    return DiffTestBlueprint(
        trace=DataPath('counters.json'),
        query="""
          INCLUDE PERFETTO MODULE std.trees.table_conversion;
          INCLUDE PERFETTO MODULE std.trees.filter;

          CREATE PERFETTO TABLE input_tree AS
          SELECT 1 AS id, NULL AS parent_id, 'keep_1' AS name
          UNION ALL SELECT 2, 1, 'remove_2'
          UNION ALL SELECT 3, 1, 'keep_1'
          UNION ALL SELECT 4, 2, 'keep_1';

          SELECT _tree_id, _tree_parent_id, id, name
          FROM _tree_to_table!(
            _tree_filter(
              _tree_from_table!((SELECT * FROM input_tree), (name)),
              _tree_where(
                _tree_constraint('name', '!=', printf('remove_%d', 2)),
                _tree_constraint('name', '=', printf('keep_%d', 1))
              )
            ),
            (name)
          )
          ORDER BY id;
        """,
        out=Csv("""
        "_tree_id","_tree_parent_id","id","name"
        0,"[NULL]",1,"keep_1"
        1,0,3,"keep_1"
        2,0,4,"keep_1"
        """))
