import pytest
from unittest.mock import Mock, patch
import requests
from sagemaker_gen_ai_jupyterlab_extension.extract_utils import (
    get_sm_environment_metadata,
    build_origin_validation_patterns,
    EnvironmentMetadata
)


class TestGetSMEnvironmentMetadata:
    """Tests for get_sm_environment_metadata() function"""
    
    @patch('sagemaker_gen_ai_jupyterlab_extension.extract_utils.requests.get')
    def test_success_with_complete_metadata(self, mock_get):
        """Test successful metadata fetch with all fields present"""
        mock_response = Mock()
        mock_response.json.return_value = {
            'domain_id': 'd-abc123',
            'project_id': 'p-xyz789',
            'aws_region': 'us-east-1',
            'environment_id': 'env-123',
            'repository_name': 'test-repo',
            'user_id': 'user-456',
            'dz_stage': 'prod',
            'sm_domain_id': 'd-abc123',
            'sm_space_name': 'test-space',
            'sm_user_profile_name': 'test-user',
            'sm_project_path': '/projects/test'
        }
        mock_response.raise_for_status = Mock()
        mock_get.return_value = mock_response
        
        result = get_sm_environment_metadata('test.sagemaker.us-east-1.on.aws')
        
        assert result['domain_id'] == 'd-abc123'
        assert result['aws_region'] == 'us-east-1'
        assert result['dz_stage'] == 'prod'
        mock_get.assert_called_once_with(
            'http://test.sagemaker.us-east-1.on.aws/jupyterlab/default/api/env',
            timeout=5
        )
    
    @patch('sagemaker_gen_ai_jupyterlab_extension.extract_utils.requests.get')
    def test_http_error(self, mock_get):
        """Test handling of HTTP errors during metadata fetch"""
        mock_get.side_effect = requests.HTTPError("404 Not Found")
        
        with pytest.raises(requests.RequestException):
            get_sm_environment_metadata('test.sagemaker.us-east-1.on.aws')
    
    @patch('sagemaker_gen_ai_jupyterlab_extension.extract_utils.requests.get')
    def test_timeout_error(self, mock_get):
        """Test handling of timeout during metadata fetch"""
        mock_get.side_effect = requests.Timeout("Connection timeout")
        
        with pytest.raises(requests.Timeout):
            get_sm_environment_metadata('test.sagemaker.us-east-1.on.aws')
    
    @patch('sagemaker_gen_ai_jupyterlab_extension.extract_utils.requests.get')
    def test_invalid_json_response(self, mock_get):
        """Test handling of invalid JSON response"""
        mock_response = Mock()
        mock_response.json.side_effect = ValueError("Invalid JSON")
        mock_response.raise_for_status = Mock()
        mock_get.return_value = mock_response
        
        with pytest.raises(ValueError):
            get_sm_environment_metadata('test.sagemaker.us-east-1.on.aws')
    
    @patch('sagemaker_gen_ai_jupyterlab_extension.extract_utils.requests.get')
    def test_missing_required_field(self, mock_get):
        """Test handling of response missing required fields"""
        mock_response = Mock()
        mock_response.json.return_value = {
            'domain_id': 'd-abc123',
            'aws_region': 'us-east-1'
            # Missing 'dz_stage' - required field
        }
        mock_response.raise_for_status = Mock()
        mock_get.return_value = mock_response
        
        with pytest.raises(ValueError, match="Missing required field"):
            get_sm_environment_metadata('test.sagemaker.us-east-1.on.aws')
    
    @patch('sagemaker_gen_ai_jupyterlab_extension.extract_utils.requests.get')
    def test_with_localhost_host(self, mock_get):
        """Test with localhost-style host"""
        mock_response = Mock()
        mock_response.json.return_value = {
            'domain_id': 'd-local',
            'aws_region': 'us-west-2',
            'dz_stage': 'gamma'
        }
        mock_response.raise_for_status = Mock()
        mock_get.return_value = mock_response
        
        result = get_sm_environment_metadata('localhost:8888')
        
        assert result['domain_id'] == 'd-local'
        mock_get.assert_called_once_with(
            'http://localhost:8888/jupyterlab/default/api/env',
            timeout=5
        )


class TestBuildOriginValidationPatterns:
    """Tests for build_origin_validation_patterns() function"""
    
    def test_prod_stage_single_domain(self):
        """Test pattern building for production stage with single domain"""
        env = {
            'domain_id': 'd-abc123',
            'aws_region': 'us-east-1',
            'dz_stage': 'prod',
            'sm_domain_id': 'd-abc123'  # Same as domain_id
        }
        
        patterns = build_origin_validation_patterns(env)
        
        assert len(patterns) == 1
        assert patterns[0] == 'https://d-abc123.sagemaker.us-east-1.on.aws'
    
    def test_gamma_stage_single_domain(self):
        """Test pattern building for gamma stage"""
        env = {
            'domain_id': 'd-xyz789',
            'aws_region': 'eu-west-1',
            'dz_stage': 'gamma',
            'sm_domain_id': 'd-xyz789'
        }
        
        patterns = build_origin_validation_patterns(env)
        
        assert len(patterns) == 1
        assert patterns[0] == 'https://d-xyz789.sagemaker-gamma.eu-west-1.on.aws'
    
    def test_beta_stage_single_domain(self):
        """Test pattern building for beta stage"""
        env = {
            'domain_id': 'd-beta',
            'aws_region': 'ap-south-1',
            'dz_stage': 'beta'
        }
        
        patterns = build_origin_validation_patterns(env)
        
        assert len(patterns) == 1
        assert patterns[0] == 'https://d-beta.sagemaker-beta.ap-south-1.on.aws'
    
    def test_multiple_domains(self):
        """Test pattern building with different domain_id and sm_domain_id"""
        env = {
            'domain_id': 'd-primary',
            'aws_region': 'us-west-2',
            'dz_stage': 'prod',
            'sm_domain_id': 'd-secondary'  # Different from domain_id
        }
        
        patterns = build_origin_validation_patterns(env)
        
        assert len(patterns) == 2
        assert 'https://d-primary.sagemaker.us-west-2.on.aws' in patterns
        assert 'https://d-secondary.sagemaker.us-west-2.on.aws' in patterns
    
    def test_missing_domain_id(self):
        """Test pattern building when domain_id is missing"""
        env = {
            'aws_region': 'us-east-1',
            'dz_stage': 'prod',
            'sm_domain_id': 'd-only-sm'
        }
        
        patterns = build_origin_validation_patterns(env)
        
        # Should still build pattern for sm_domain_id
        assert len(patterns) == 1
        assert patterns[0] == 'https://d-only-sm.sagemaker.us-east-1.on.aws'
    
    def test_missing_region(self):
        """Test pattern building when aws_region is missing"""
        env = {
            'domain_id': 'd-abc123',
            'dz_stage': 'prod'
        }
        
        patterns = build_origin_validation_patterns(env)
        
        # Should return empty list if region missing
        assert len(patterns) == 0
    
    def test_empty_domain_id(self):
        """Test pattern building with empty domain_id"""
        env = {
            'domain_id': '',
            'aws_region': 'us-east-1',
            'dz_stage': 'prod'
        }
        
        patterns = build_origin_validation_patterns(env)
        
        assert len(patterns) == 0
    
    def test_none_values(self):
        """Test pattern building with None values"""
        env = {
            'domain_id': None,
            'aws_region': 'us-east-1',
            'dz_stage': 'prod',
            'sm_domain_id': None
        }
        
        patterns = build_origin_validation_patterns(env)
        
        assert len(patterns) == 0
    
    def test_all_regions(self):
        """Test pattern building works for various AWS regions"""
        regions = [
            'us-east-1', 'us-west-2', 'eu-west-1', 'eu-central-1',
            'ap-south-1', 'ap-southeast-1', 'ap-northeast-1'
        ]
        
        for region in regions:
            env = {
                'domain_id': 'd-test',
                'aws_region': region,
                'dz_stage': 'prod'
            }
            
            patterns = build_origin_validation_patterns(env)
            
            assert len(patterns) == 1
            assert patterns[0] == f'https://d-test.sagemaker.{region}.on.aws'
    
    def test_stage_variations(self):
        """Test all stage variations produce correct patterns"""
        stages = {
            'prod': '',
            'gamma': '-gamma',
            'beta': '-beta',
            'alpha': '-alpha'
        }
        
        for stage_name, stage_suffix in stages.items():
            env = {
                'domain_id': 'd-test',
                'aws_region': 'us-east-1',
                'dz_stage': stage_name
            }
            
            patterns = build_origin_validation_patterns(env)
            
            assert len(patterns) == 1
            expected = f'https://d-test.sagemaker{stage_suffix}.us-east-1.on.aws'
            assert patterns[0] == expected, f"Stage {stage_name} should produce {expected}"


class TestIntegrationScenarios:
    """Integration tests combining metadata fetch and pattern building"""
    
    @patch('sagemaker_gen_ai_jupyterlab_extension.extract_utils.requests.get')
    def test_full_workflow_prod(self, mock_get):
        """Test complete workflow from fetch to pattern building (prod)"""
        mock_response = Mock()
        mock_response.json.return_value = {
            'domain_id': 'd-prod123',
            'aws_region': 'us-east-1',
            'dz_stage': 'prod',
            'sm_domain_id': 'd-prod123',
            'project_id': 'p-123',
            'environment_id': 'env-123',
            'repository_name': 'test',
            'user_id': 'user-123',
            'sm_space_name': 'space',
            'sm_user_profile_name': 'profile',
            'sm_project_path': '/path'
        }
        mock_response.raise_for_status = Mock()
        mock_get.return_value = mock_response
        
        # Fetch metadata
        env = get_sm_environment_metadata('d-prod123.notebook.us-east-1.sagemaker.aws')
        
        # Build patterns
        patterns = build_origin_validation_patterns(env)
        
        assert len(patterns) == 1
        assert patterns[0] == 'https://d-prod123.sagemaker.us-east-1.on.aws'
    
    @patch('sagemaker_gen_ai_jupyterlab_extension.extract_utils.requests.get')
    def test_full_workflow_gamma_multiple_domains(self, mock_get):
        """Test complete workflow with gamma stage and multiple domains"""
        mock_response = Mock()
        mock_response.json.return_value = {
            'domain_id': 'd-primary',
            'aws_region': 'eu-west-1',
            'dz_stage': 'gamma',
            'sm_domain_id': 'd-secondary',  # Different
            'project_id': 'p-123',
            'environment_id': 'env-123',
            'repository_name': 'test',
            'user_id': 'user-123',
            'sm_space_name': 'space',
            'sm_user_profile_name': 'profile',
            'sm_project_path': '/path'
        }
        mock_response.raise_for_status = Mock()
        mock_get.return_value = mock_response
        
        # Fetch metadata
        env = get_sm_environment_metadata('d-primary.notebook.eu-west-1.sagemaker.aws')
        
        # Build patterns
        patterns = build_origin_validation_patterns(env)
        
        assert len(patterns) == 2
        assert 'https://d-primary.sagemaker-gamma.eu-west-1.on.aws' in patterns
        assert 'https://d-secondary.sagemaker-gamma.eu-west-1.on.aws' in patterns