import pytest
import argparse
from unittest.mock import Mock, patch, MagicMock
from app.main import main
from app.dtos import User
from app.migrate_to_v2_permissions import OldRole, NewRole, FeatureFlag


class TestMainMigrationIntegration:
    """Testing CLI integration for migration functionality"""
    
    @patch('app.main.perform_migration')
    @patch('app.main.fill_users_data_from_neo4j')
    @patch('app.main.get_users_by_emails')
    @patch('app.main.check_vpn_is_working')
    @patch('app.main.build_users_data_table')
    @patch('app.main.get_pdp_permissions_for_identities')
    @patch('app.main.check_treatments')
    @patch('sys.argv', ['main.py', '-e', 'test@example.com', '--migrate'])
    def test_main_migrate_flag(self, mock_check_treatments, mock_get_permissions, mock_build_table,
                              mock_check_vpn, mock_get_users, mock_fill_users, mock_perform_migration):
        """Test main function with --migrate flag"""
        # Setup mocks
        mock_check_vpn.return_value = True
        mock_get_users.return_value = [{'email': 'test@example.com', 'user_id': 'uuid1'}]
        mock_get_permissions.return_value = {'uuid1': {'tenant1': ['admin']}}
        mock_check_treatments.return_value = {'uuid1': ['orchard_suite_show_audience_app']}
        mock_perform_migration.return_value = (True, 'logs/migration_log.csv')
        
        with patch('app.main.Console') as mock_console_cls:
            mock_console = Mock()
            mock_console_cls.return_value = mock_console
            
            # Should not raise any exceptions
            main()
            
            # Verify migration was called
            mock_perform_migration.assert_called_once()
            args, kwargs = mock_perform_migration.call_args
            users, dry_run = args[0], kwargs.get('dry_run', False)
            
            assert dry_run is False
            assert len(users) == 1
            assert users[0].email == 'test@example.com'
            
            # Verify success message
            mock_console.print.assert_any_call(
                "[green]Migration completed successfully! Log saved to: logs/migration_log.csv[/green]"
            )

    @patch('app.main.perform_migration')
    @patch('app.main.fill_users_data_from_neo4j')
    @patch('app.main.get_users_by_emails')
    @patch('app.main.check_vpn_is_working')
    @patch('app.main.build_users_data_table')
    @patch('app.main.get_pdp_permissions_for_identities')
    @patch('app.main.check_treatments')
    @patch('sys.argv', ['main.py', '-e', 'test@example.com', '--migrate-dry-run'])
    def test_main_migrate_dry_run_flag(self, mock_check_treatments, mock_get_permissions, mock_build_table,
                                      mock_check_vpn, mock_get_users, mock_fill_users, mock_perform_migration):
        """Test main function with --migrate-dry-run flag"""
        # Setup mocks
        mock_check_vpn.return_value = True
        mock_get_users.return_value = [{'email': 'test@example.com', 'user_id': 'uuid1'}]
        mock_get_permissions.return_value = {'uuid1': {'tenant1': ['admin']}}
        mock_check_treatments.return_value = {'uuid1': ['orchard_suite_show_audience_app']}
        mock_perform_migration.return_value = (True, None)
        
        with patch('app.main.Console') as mock_console_cls:
            mock_console = Mock()
            mock_console_cls.return_value = mock_console
            
            main()
            
            # Verify dry run migration was called
            mock_perform_migration.assert_called_once()
            args, kwargs = mock_perform_migration.call_args
            users, dry_run = args[0], kwargs.get('dry_run', False)
            
            assert dry_run is True
            
            # Verify dry run message
            mock_console.print.assert_any_call(
                "[green]Migration analysis completed successfully![/green]"
            )

    @patch('app.main.perform_migration')
    @patch('app.main.fill_users_data_from_neo4j')
    @patch('app.main.get_users_by_identities')
    @patch('app.main.check_vpn_is_working')
    @patch('app.main.build_users_data_table')
    @patch('app.main.get_pdp_permissions_for_identities')
    @patch('app.main.check_treatments')
    @patch('sys.argv', ['main.py', '-i', 'uuid1,uuid2', '--migrate'])
    def test_main_migrate_with_identities(self, mock_check_treatments, mock_get_permissions, mock_build_table,
                                         mock_check_vpn, mock_get_users, mock_fill_users, mock_perform_migration):
        """Test main function with --migrate flag using identities"""
        # Setup mocks
        mock_check_vpn.return_value = True
        mock_get_users.return_value = [
            {'email': 'test1@example.com', 'user_id': 'uuid1'},
            {'email': 'test2@example.com', 'user_id': 'uuid2'}
        ]
        mock_get_permissions.return_value = {
            'uuid1': {'tenant1': ['admin']},
            'uuid2': {'tenant1': ['analyst']}
        }
        mock_check_treatments.return_value = {
            'uuid1': ['orchard_suite_show_audience_app'],
            'uuid2': ['orchard_suite_show_audience_app']
        }
        mock_perform_migration.return_value = (True, 'logs/migration_log.csv')
        
        with patch('app.main.Console') as mock_console_cls:
            mock_console = Mock()
            mock_console_cls.return_value = mock_console
            
            main()
            
            # Verify migration was called with correct users
            mock_perform_migration.assert_called_once()
            args, kwargs = mock_perform_migration.call_args
            users = args[0]
            
            assert len(users) == 2
            assert users[0].identity == 'uuid1'
            assert users[1].identity == 'uuid2'

    @patch('app.main.perform_migration')
    @patch('app.main.fill_users_data_from_neo4j')
    @patch('app.main.get_users_from_csv')
    @patch('app.main.check_vpn_is_working')
    @patch('app.main.build_users_data_table')
    @patch('app.main.get_pdp_permissions_for_identities')
    @patch('app.main.check_treatments')
    @patch('sys.argv', ['main.py', '-f', 'users.csv', '--migrate'])
    def test_main_migrate_with_csv_file(self, mock_check_treatments, mock_get_permissions, mock_build_table,
                                       mock_check_vpn, mock_get_csv_users, mock_fill_users, mock_perform_migration):
        """Test main function with --migrate flag using CSV file"""
        # Setup mocks
        mock_check_vpn.return_value = True
        mock_csv_users = [
            User(email='user1@example.com', identity='uuid1', role='admin'),
            User(email='user2@example.com', identity='uuid2', role='analyst')
        ]
        mock_get_csv_users.return_value = mock_csv_users
        
        # Mock Neo4j data
        neo4j_users = [
            {'email': 'user1@example.com', 'user_id': 'uuid1', 'name': 'User One'},
            {'email': 'user2@example.com', 'user_id': 'uuid2', 'name': 'User Two'}
        ]
        
        with patch('app.main.get_users_by_emails') as mock_get_users:
            mock_get_users.return_value = neo4j_users
            mock_get_permissions.return_value = {
                'uuid1': {'tenant1': ['admin']},
                'uuid2': {'tenant1': ['analyst']}
            }
            mock_check_treatments.return_value = {
                'uuid1': ['orchard_suite_show_audience_app'],
                'uuid2': ['orchard_suite_show_audience_app']
            }
            mock_perform_migration.return_value = (True, 'logs/migration_log.csv')
            
            with patch('app.main.Console') as mock_console_cls:
                mock_console = Mock()
                mock_console_cls.return_value = mock_console
                
                main()
                
                # Verify migration was called
                mock_perform_migration.assert_called_once()
                args, kwargs = mock_perform_migration.call_args
                users = args[0]
                
                assert len(users) == 2

    @patch('app.main.perform_migration')
    @patch('app.main.fill_users_data_from_neo4j')
    @patch('app.main.get_users_by_emails')
    @patch('app.main.check_vpn_is_working')
    @patch('app.main.build_users_data_table')
    @patch('app.main.get_pdp_permissions_for_identities')
    @patch('app.main.check_treatments')
    @patch('app.main.alert_message_and_exit')
    @patch('sys.argv', ['main.py', '-e', 'test@example.com', '--migrate'])
    def test_main_migrate_with_non_existent_users(self, mock_alert_exit, mock_check_treatments, 
                                                  mock_get_permissions, mock_build_table, mock_check_vpn,
                                                  mock_get_users, mock_fill_users, mock_perform_migration):
        """Test main function with --migrate flag when users don't exist"""
        # Setup mocks - no users found in Neo4j
        mock_check_vpn.return_value = True
        mock_get_users.return_value = []  # No users found
        mock_get_permissions.return_value = {}
        mock_check_treatments.return_value = {}
        
        with patch('app.main.Console') as mock_console_cls:
            mock_console = Mock()
            mock_console_cls.return_value = mock_console
            
            main()
            
            # Should call alert_message_and_exit for non-existent users
            mock_alert_exit.assert_called_with("We could migrate only users that exists")
            mock_perform_migration.assert_not_called()

    @patch('app.main.perform_migration')
    @patch('app.main.fill_users_data_from_neo4j')
    @patch('app.main.get_users_by_emails')
    @patch('app.main.check_vpn_is_working')
    @patch('app.main.build_users_data_table')
    @patch('app.main.get_pdp_permissions_for_identities')
    @patch('app.main.check_treatments')
    @patch('app.main.alert_message_and_exit')
    @patch('sys.argv', ['main.py', '-e', 'test@example.com', '--migrate'])
    def test_main_migrate_failure(self, mock_alert_exit, mock_check_treatments, mock_get_permissions,
                                 mock_build_table, mock_check_vpn, mock_get_users, mock_fill_users,
                                 mock_perform_migration):
        """Test main function when migration fails"""
        # Setup mocks
        mock_check_vpn.return_value = True
        mock_get_users.return_value = [{'email': 'test@example.com', 'user_id': 'uuid1'}]
        mock_get_permissions.return_value = {'uuid1': {'tenant1': ['admin']}}
        mock_check_treatments.return_value = {'uuid1': ['orchard_suite_show_audience_app']}
        mock_perform_migration.return_value = (False, None)  # Migration failed
        
        with patch('app.main.Console') as mock_console_cls:
            mock_console = Mock()
            mock_console_cls.return_value = mock_console
            
            main()
            
            # Should call alert_message_and_exit on failure
            mock_alert_exit.assert_called_with("Migration failed. Check error messages above.")

    def test_argparse_migrate_flags(self):
        """Test that argparse correctly handles migration flags"""
        parser = argparse.ArgumentParser()
        
        # Add the same arguments as in main.py
        users_input_group = parser.add_mutually_exclusive_group(required=True)
        users_input_group.add_argument("-e", "--emails", metavar="EMAIL1,EMAIL2,...", type=str)
        users_input_group.add_argument("-i", "--identities", metavar="IDENTITY1,IDENTITY2,...", type=str)
        users_input_group.add_argument("-f", "--csv-file", metavar="FILE", type=str)
        
        parser.add_argument("-m", "--migrate", action='store_true')
        parser.add_argument("--md", "--migrate-dry-run", action='store_true', dest='migrate_dry_run')
        
        # Test --migrate flag
        args = parser.parse_args(['-e', 'test@example.com', '--migrate'])
        assert args.migrate is True
        assert args.migrate_dry_run is False
        
        # Test --migrate-dry-run flag
        args = parser.parse_args(['-e', 'test@example.com', '--migrate-dry-run'])
        assert args.migrate is False
        assert args.migrate_dry_run is True
        
        # Test --md short form
        args = parser.parse_args(['-e', 'test@example.com', '--md'])
        assert args.migrate is False
        assert args.migrate_dry_run is True

    @patch('app.main.perform_migration')
    @patch('app.main.fill_users_data_from_neo4j')
    @patch('app.main.get_users_by_emails')
    @patch('app.main.check_vpn_is_working')
    @patch('app.main.build_users_data_table')
    @patch('app.main.get_pdp_permissions_for_identities')
    @patch('app.main.check_treatments')
    @patch('sys.argv', ['main.py', '-e', 'test@example.com'])
    def test_main_without_migration_flags(self, mock_check_treatments, mock_get_permissions, mock_build_table,
                                         mock_check_vpn, mock_get_users, mock_fill_users, mock_perform_migration):
        """Test main function without migration flags - should not call perform_migration"""
        # Setup mocks
        mock_check_vpn.return_value = True
        mock_get_users.return_value = [{'email': 'test@example.com', 'user_id': 'uuid1'}]
        mock_get_permissions.return_value = {'uuid1': {'tenant1': ['admin']}}
        mock_check_treatments.return_value = {'uuid1': ['orchard_suite_show_audience_app']}
        
        with patch('app.main.Console') as mock_console_cls:
            mock_console = Mock()
            mock_console_cls.return_value = mock_console
            
            main()
            
            # Migration should not be called
            mock_perform_migration.assert_not_called()
            
            # But user data table should be built (normal flow)
            mock_build_table.assert_called_once()


class TestMigrationFlagValidation:
    """Test validation and edge cases for migration flags"""
    
    def test_migration_flags_mutually_exclusive_with_other_actions(self):
        """Test that migration flags work correctly with other action flags"""
        parser = argparse.ArgumentParser()
        
        users_input_group = parser.add_mutually_exclusive_group(required=True)
        users_input_group.add_argument("-e", "--emails", type=str)
        
        # Add migration flags and other action flags
        parser.add_argument("-m", "--migrate", action='store_true')
        parser.add_argument("--md", "--migrate-dry-run", action='store_true', dest='migrate_dry_run')
        parser.add_argument("-p", "--permissions", action='store_true')
        parser.add_argument("-d", "--database-pr", type=str)
        
        # Should be able to combine migration with email input
        args = parser.parse_args(['-e', 'test@example.com', '--migrate'])
        assert args.migrate is True
        assert args.permissions is False
        
        # Should be able to combine other flags
        args = parser.parse_args(['-e', 'test@example.com', '--permissions'])
        assert args.migrate is False
        assert args.permissions is True
        
        # Can use both migration flags (though logically one would override)
        args = parser.parse_args(['-e', 'test@example.com', '--migrate', '--migrate-dry-run'])
        assert args.migrate is True
        assert args.migrate_dry_run is True