# Copyright 2024 Google Inc.
#
# 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.

"""Tests for the ot_process_reminders module, verifying the formatting and logic of origin trial reminder emails."""

import logging
from unittest import mock

import testing_config  # Must be imported before the module under test.
from internals import ot_process_reminders
from internals.core_models import FeatureEntry, Stage


class OTProcessRemindersTest(testing_config.CustomTestCase):
    """Tests for OTProcessReminders."""

    def setUp(self):
        """Set up the test environment."""
        logging.disable(logging.CRITICAL)
        self.feature_1 = FeatureEntry(
            feature_type=1, name='feature one', summary='sum', category=1
        )
        self.feature_1.put()
        self.stage_1 = Stage(
            feature_id=1,
            stage_type=150,
            ot_owner_email='ot_owner1@google.com',
            ot_emails=['contact1@example.com', 'contact2@example.com'],
            origin_trial_id='4199606652522987521',
        )
        self.stage_1.put()
        self.stage_2 = Stage(
            feature_id=2,
            stage_type=150,
            ot_owner_email='ot_owner2@google.com',
            ot_emails=['sample_contact@example.com', 'another@example.com'],
            origin_trial_id='1',
        )
        self.stage_2.put()
        self.mock_get_trials_list_return_value = [
            {
                'id': '4199606652522987521',
                'display_name': 'feature one',
                'description': 'Another origin trial',
                'origin_trial_feature_name': 'ChromiumTrialName',
                'enabled': True,
                'status': 'ACTIVE',
                'chromestatus_url': 'https://chromestatus.com/feature/123',
                'start_milestone': '97',
                'end_milestone': '103',
                'original_end_milestone': '100',
                'end_time': '2020-11-10T23:59:59Z',
                'documentation_url': 'https://example.com/docs',
                'feedback_url': 'https://example.com/feedback',
                'intent_to_experiment_url': 'https://example.com/experiment',
                'trial_extensions': [
                    {
                        'endMilestone': '110',
                        'endTime': '2020-11-10T23:59:59Z',
                        'extensionIntentUrl': 'https://example.com/extension',
                    }
                ],
                'type': 'DEPRECATION',
                'allow_third_party_origins': True,
            },
            {
                'id': '1',
                'display_name': 'Sample trial 2',
                'description': 'Another origin trial 2',
                'origin_trial_feature_name': 'ChromiumTrialName2',
                'enabled': True,
                'status': 'ACTIVE',
                'chromestatus_url': 'https://chromestatus.com/feature/456',
                'start_milestone': '103',
                'end_milestone': '109',
                'original_end_milestone': '106',
                'end_time': '2020-11-10T23:59:59Z',
                'documentation_url': 'https://example.com/docs2',
                'feedback_url': 'https://example.com/feedback2',
                'intent_to_experiment_url': 'https://example.com/experiment2',
                'trial_extensions': [
                    {
                        'endMilestone': '109',
                        'endTime': '2020-11-10T23:59:59Z',
                        'extensionIntentUrl': 'https://example.com/extension2',
                    }
                ],
                'type': 'ORIGIN_TRIAL',
                'allow_third_party_origins': False,
            },
            {
                'id': '2',
                'display_name': 'Sample trial 3',
                'description': 'Another origin trial 3',
                'origin_trial_feature_name': 'ChromiumTrialName3',
                'enabled': True,
                'status': 'ACTIVE',
                'chromestatus_url': 'https://chromestatus.com/feature/789',
                'start_milestone': '110',
                'end_milestone': '116',
                'original_end_milestone': '116',
                'end_time': '2020-11-10T23:59:59Z',
                'documentation_url': 'https://example.com/docs3',
                'feedback_url': 'https://example.com/feedback3',
                'intent_to_experiment_url': 'https://example.com/experiment3',
                'trial_extensions': [],
                'type': 'ORIGIN_TRIAL',
                'allow_third_party_origins': False,
            },
        ]

    def tearDown(self):
        """Clean up the test environment."""
        logging.disable(logging.NOTSET)

    def test_build_trials__normal(self):
        """Test that trial data is formatted correctly."""
        expected_trials = [
            {
                'id': '4199606652522987521',
                'name': 'feature one',
                'start_milestone': 97,
                'end_milestone': 103,
                'contacts': [
                    'ot_owner1@google.com',
                    'contact1@example.com',
                    'contact2@example.com',
                ],
            },
            {
                'id': '1',
                'name': 'Sample trial 2',
                'start_milestone': 103,
                'end_milestone': 109,
                'contacts': [
                    'sample_contact@example.com',
                    'another@example.com',
                    'ot_owner2@google.com',
                ],
            },
        ]
        stage_map = {
            self.stage_1.origin_trial_id: self.stage_1,
            self.stage_2.origin_trial_id: self.stage_2,
        }
        formatted_1 = ot_process_reminders.build_trial_data(
            self.mock_get_trials_list_return_value[0], stage_map
        )
        formatted_2 = ot_process_reminders.build_trial_data(
            self.mock_get_trials_list_return_value[1], stage_map
        )
        formatted_3 = ot_process_reminders.build_trial_data(
            self.mock_get_trials_list_return_value[2], stage_map
        )
        # 3rd trial does not have associated stage, so should return as None.
        self.assertIsNone(formatted_3)

        self.assertEqual(expected_trials[0]['id'], formatted_1['id'])
        self.assertEqual(expected_trials[1]['id'], formatted_2['id'])
        self.assertEqual(expected_trials[0]['name'], formatted_1['name'])
        self.assertEqual(expected_trials[1]['name'], formatted_2['name'])
        self.assertEqual(
            expected_trials[0]['start_milestone'],
            formatted_1['start_milestone'],
        )
        self.assertEqual(
            expected_trials[1]['start_milestone'],
            formatted_2['start_milestone'],
        )
        self.assertEqual(
            set(expected_trials[0]['contacts']), set(formatted_1['contacts'])
        )
        self.assertEqual(
            set(expected_trials[1]['contacts']), set(formatted_2['contacts'])
        )

    @mock.patch('framework.origin_trials_client.get_trials_list')
    def test_get_trials(self, mock_get_trials_list):
        """Test get trials."""
        mock_get_trials_list.return_value = (
            self.mock_get_trials_list_return_value
        )
        starting_trials, ending_trials = ot_process_reminders.get_trials(103)
        expected_starting_trials = [
            {
                'id': '1',
                'name': 'Sample trial 2',
                'start_milestone': 103,
                'end_milestone': 109,
                'contacts': [
                    'sample_contact@example.com',
                    'another@example.com',
                    'ot_owner2@google.com',
                ],
            }
        ]
        expected_ending_trials = [
            {
                'id': '4199606652522987521',
                'name': 'feature one',
                'start_milestone': 97,
                'end_milestone': 103,
                'contacts': [
                    'contact1@example.com',
                    'contact2@example.com',
                    'ot_owner1@google.com',
                ],
            }
        ]

        self.assertEqual(len(starting_trials), 1)
        self.assertEqual(len(ending_trials), 1)
        starting_trial = starting_trials[0]
        ending_trial = ending_trials[0]
        self.assertEqual(
            starting_trial['id'], expected_starting_trials[0]['id']
        )
        self.assertEqual(ending_trial['id'], expected_ending_trials[0]['id'])
        self.assertEqual(
            starting_trial['name'], expected_starting_trials[0]['name']
        )
        self.assertEqual(
            ending_trial['name'], expected_ending_trials[0]['name']
        )
        self.assertEqual(
            set(starting_trial['contacts']),
            set(expected_starting_trials[0]['contacts']),
        )
        self.assertEqual(
            set(ending_trial['contacts']),
            set(expected_ending_trials[0]['contacts']),
        )

    @mock.patch('logging.exception')
    @mock.patch('framework.origin_trials_client.get_trials_list')
    def test_get_trials__null_ot_owner(
        self, mock_get_trials_list, mock_logging
    ):
        """If OT owner is not set, any other contacts should still be notified."""
        # OT owner email is not set.
        self.stage_2.ot_owner_email = None
        self.stage_2.put()
        mock_get_trials_list.return_value = (
            self.mock_get_trials_list_return_value
        )
        starting_trials, ending_trials = ot_process_reminders.get_trials(103)

        expected_starting_trial_contacts = [
            'sample_contact@example.com',
            'another@example.com',
        ]
        expected_ending_trial_contacts = [
            'contact1@example.com',
            'contact2@example.com',
            'ot_owner1@google.com',
        ]

        self.assertEqual(len(starting_trials), 1)
        self.assertEqual(len(ending_trials), 1)
        starting_trial = starting_trials[0]
        ending_trial = ending_trials[0]
        self.assertEqual(
            set(starting_trial['contacts']),
            set(expected_starting_trial_contacts),
        )
        self.assertEqual(
            set(ending_trial['contacts']), set(expected_ending_trial_contacts)
        )

    def test_get_trial_end_release_offset(self):
        """Test that the correct offset is returned based on release milestone."""
        self.assertEqual(
            ot_process_reminders.get_trial_end_release_offset(150), 2
        )
        self.assertEqual(
            ot_process_reminders.get_trial_end_release_offset(151), 2
        )
        self.assertEqual(
            ot_process_reminders.get_trial_end_release_offset(152), 3
        )
        self.assertEqual(
            ot_process_reminders.get_trial_end_release_offset(153), 4
        )
        self.assertEqual(
            ot_process_reminders.get_trial_end_release_offset(154), 4
        )

    def test_get_release_plus_n(self):
        """Test that get_release_plus_n correctly increments milestones."""
        self.assertEqual(ot_process_reminders.get_release_plus_n(150, 2), 152)
        self.assertEqual(ot_process_reminders.get_release_plus_n(150, 4), 154)
        # Test skipping milestone 82
        self.assertEqual(ot_process_reminders.get_release_plus_n(80, 2), 83)
