#!/usr/bin/env python3
# Copyright 2025 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

from aggregation import AggregationKind
from aggregation import AggregationDetails
from aggregation import GetAggregationDetails
from aggregation import GenerateCCAggregation
from aggregation import GenerateHHAggregation
import unittest


class AggregationTest(unittest.TestCase):
  def testDefaults(self):
    self.assertEqual(
      GetAggregationDetails({}),
      AggregationDetails(AggregationKind.NONE, None, True, {}, None),
    )

  def testGenerateArray(self):
    self.assertEqual(
      GetAggregationDetails(
        {"elements": {}, "aggregation": {"type": "array", "name": "TestArray"}}
      ),
      AggregationDetails(AggregationKind.ARRAY, "TestArray", True, {}, None),
    )

  def testGenerateArrayAndHideElements(self):
    self.assertEqual(
      GetAggregationDetails(
        {
          "elements": {},
          "aggregation": {
            "type": "array",
            "name": "TestArray",
            "export_items": False,
          },
        }
      ),
      AggregationDetails(AggregationKind.ARRAY, "TestArray", False, {}, None),
    )

  def testGenerateMapWithDefaultKeyType(self):
    self.assertEqual(
      GetAggregationDetails(
        {"elements": {}, "aggregation": {"type": "map", "name": "TestMap"}}
      ),
      AggregationDetails(
        AggregationKind.MAP, "TestMap", True, {}, "std::string_view"
      ),
    )

  def testGenerateMapWithSuppliedKeyType(self):
    self.assertEqual(
      GetAggregationDetails(
        {
          "elements": {},
          "aggregation": {
            "type": "map",
            "name": "TestMap",
            "map_key_type": "chrome::CustomKey",
          },
        }
      ),
      AggregationDetails(
        AggregationKind.MAP, "TestMap", True, {}, "chrome::CustomKey"
      ),
    )

  def testGenerateMapAndHideElements(self):
    self.assertEqual(
      GetAggregationDetails(
        {
          "elements": {},
          "aggregation": {
            "type": "map",
            "name": "TestMap",
            "export_items": False,
          },
        }
      ),
      AggregationDetails(
        AggregationKind.MAP, "TestMap", False, {}, "std::string_view"
      ),
    )

  def testGetSortedArrayElements(self):
    aggregation = GetAggregationDetails(
      {
        "elements": {"c": [], "d": [], "b": [], "a": []},
        "aggregation": {"type": "array", "name": "TestMap"},
      }
    )

    self.assertEqual(["a", "b", "c", "d"], aggregation.GetSortedArrayElements())

  def testGetSortedMapElementsNoAliases(self):
    aggregation = GetAggregationDetails(
      {
        "elements": {"c": [], "d": [], "b": [], "a": []},
        "aggregation": {"type": "map", "name": "TestMap"},
      }
    )

    # Expect the result to be sorted as well.
    self.assertEqual(
      [("a", "a"), ("b", "b"), ("c", "c"), ("d", "d")],
      aggregation.GetSortedMapElements(),
    )

  def testGetSortedMapElementsWithAliases(self):
    aggregation = GetAggregationDetails(
      {
        "elements": {
          "c": [],
          "d": [],
        },
        "aggregation": {
          "type": "map",
          "name": "TestMap",
          "map_aliases": {
            "a": "d",
            "b": "c",
            "e": "c",
          },
        },
      }
    )

    # Expect the result to be sorted as well.
    self.assertEqual(
      [("a", "d"), ("b", "c"), ("c", "c"), ("d", "d"), ("e", "c")],
      aggregation.GetSortedMapElements(),
    )

  def testGenerateMapWithConflictingAliases(self):
    self.assertRaisesRegex(
      Exception,
      "Alias `b` already defined as element.",
      lambda: GetAggregationDetails(
        {
          "elements": {
            "a": [],
            "b": [],
          },
          "aggregation": {
            "type": "map",
            "name": "TestMap",
            "map_aliases": {"b": "a"},
          },
        }
      ),
    )

  def testGenerateMapWithIncorrectAliases(self):
    self.assertRaisesRegex(
      Exception,
      "Aliased element `b` does not exist.",
      lambda: GetAggregationDetails(
        {
          "elements": {
            "a": [],
          },
          "aggregation": {
            "type": "map",
            "name": "TestMap",
            "map_aliases": {"b": "a", "c": "b"},
          },
        }
      ),
    )

  def testUnrecognizedAggregationType(self):
    self.assertRaisesRegex(
      Exception,
      "'collection' is not a valid AggregationKind",
      lambda: GetAggregationDetails(
        {
          "elements": {},
          "aggregation": {
            "type": "collection",
          },
        }
      ),
    )

  def testArrayMustHaveAName(self):
    self.assertRaisesRegex(
      Exception,
      "Aggregation container needs a `name`.",
      lambda: GetAggregationDetails(
        {
          "elements": {},
          "aggregation": {
            "type": "array",
          },
        }
      ),
    )

  def testMapMustHaveAName(self):
    self.assertRaisesRegex(
      Exception,
      "Aggregation container needs a `name`.",
      lambda: GetAggregationDetails(
        {
          "elements": {},
          "aggregation": {
            "type": "map",
          },
        }
      ),
    )

  def BuildTestElements(
    self, source: str, target: str, count: int
  ) -> dict[str, str]:
    """Generate a map of `count` test elements.

    Elements are named `<source>_<id>` and point to `<target>_<id>`.
    """
    res = {}
    for index in range(count):
      res[f'{source}_{index + 1}'] = f'{target}_{index + 1}'
    return res

  def testHHNoAggregation(self):
    agg = AggregationDetails(
      kind=AggregationKind.NONE,
      name='kTestAggregation',
      elements=self.BuildTestElements('item', 'item', 5),
      export_items=True,
      map_key_type=None,
    )

    self.assertIsNone(GenerateHHAggregation('TypeName', agg))

  def testHHArrayAggregation(self):
    agg = AggregationDetails(
      kind=AggregationKind.ARRAY,
      name='kTestArray',
      elements=self.BuildTestElements('item', 'item', 5),
      export_items=True,
      map_key_type=None,
    )

    self.assertEqual(
      GenerateHHAggregation('TypeName', agg).strip(),
      'extern const std::array<raw_ptr<const TypeName>, 5> kTestArray;',
    )

  def testHHMapAggregation(self):
    agg = AggregationDetails(
      kind=AggregationKind.MAP,
      name='kTestMap',
      elements=self.BuildTestElements('src', 'tgt', 5),
      export_items=True,
      map_key_type='std::string_view',
    )

    self.assertEqual(
      GenerateHHAggregation('ValueTypeName', agg).strip(),
      'extern const base::fixed_flat_map<std::string_view, raw_ptr<const ValueTypeName>, 5> kTestMap;',
    )

  def testCCNoAggregation(self):
    agg = AggregationDetails(
      kind=AggregationKind.NONE,
      name='kTestAggregation',
      elements=self.BuildTestElements('item', 'item', 5),
      export_items=True,
      map_key_type=None,
    )

    self.assertIsNone(GenerateCCAggregation('TypeName', agg))

  def testCCArrayAggregation(self):
    agg = AggregationDetails(
      kind=AggregationKind.ARRAY,
      name='kTestArray',
      elements=self.BuildTestElements('item', 'item', 3),
      export_items=True,
      map_key_type=None,
    )

    self.assertEqual(
      GenerateCCAggregation('TypeName', agg),
      '''
const std::array<raw_ptr<const TypeName>, 3> kTestArray ({{
  &item_1,
  &item_2,
  &item_3,
}});
''',
    )

  def testCCMapAggregation(self):
    agg = AggregationDetails(
      kind=AggregationKind.MAP,
      name='kTestMap',
      elements=self.BuildTestElements('src', 'tgt', 3),
      export_items=True,
      map_key_type='std::string',
    )

    self.assertEqual(
      GenerateCCAggregation('ValueTypeName', agg),
      '''
const base::fixed_flat_map<std::string, raw_ptr<const ValueTypeName>, 3> kTestMap =
    base::MakeFixedFlatMap<std::string, raw_ptr<const ValueTypeName>>({
  {std::string("src_1"), &tgt_1},
  {std::string("src_2"), &tgt_2},
  {std::string("src_3"), &tgt_3},
});
''',
    )


if __name__ == '__main__':
  unittest.main()