from unittest.mock import ANY, MagicMock, patch

import pytest
from pytest import fixture, mark
from sientia_do.notifications.models import NotificationLevel

from model_manager.activities.mlflow import MLFlow


@patch('model_manager.activities.mlflow.MLFlowRepository')
def test___init__(mock_mlflow_repository):
    mlflow = MLFlow(
        mlflow_host='http://localhost',
        mlflow_port=5000,
        mlflow_username='admin',
        mlflow_password='admin',
        logger=MagicMock(),
        notification_handler=MagicMock(),
    )

    assert mlflow.mlflow_host == 'http://localhost'
    assert mlflow.mlflow_port == 5000
    assert mlflow.mlflow_username == 'admin'
    assert mlflow.mlflow_password == 'admin'

    mock_mlflow_repository.assert_called_once_with('http://localhost:5000', 'admin', 'admin', ANY)


@fixture
@patch('model_manager.activities.mlflow.MLFlowRepository')
def mlflow(mock_mlflow_repository):
    mlflow = MLFlow(
        mlflow_host='http://localhost:5000',
        mlflow_port=5000,
        mlflow_username='admin',
        mlflow_password='admin',
        logger=MagicMock(),
        notification_handler=MagicMock(),
    )

    mlflow.send_notification = MagicMock()

    return mlflow


metadata = {
    'metadata': {
        'model_id': 'test_model',
        'model_name': 'test_model',
        'workflow_name': 'test_workflow',
        'schema_name': 'test_schedule',
    },
}


@mark.asyncio
async def test_save_model_success(mlflow):
    """Test save_model successfully saves model and artifacts to MLflow."""
    from model_manager.utils.models.train_model_params import TrainModelParams
    from model_manager.utils.models.train_model_result import TrainModelResult

    # Mock train result
    params = MagicMock(spec=TrainModelParams)
    params.experiment_name = 'test_experiment'

    train_result = MagicMock(spec=TrainModelResult)
    train_result.params = params
    train_result.run_name = None  # Will be set by get_next_run_name

    # Mock repository methods
    mlflow.model_monitoring_repository.get_next_run_name.return_value = 'test_experiment-1'
    mlflow.model_monitoring_repository.generate_artifacts.return_value = train_result
    mlflow.model_monitoring_repository.save_run.return_value = None

    input_data = {
        **metadata,
        'train_result': train_result,
    }

    # Call the method
    response = await mlflow.save_model(input_data)

    # Verify repository methods were called
    mlflow.model_monitoring_repository.get_next_run_name.assert_called_once_with('test_experiment')
    mlflow.model_monitoring_repository.generate_artifacts.assert_called_once_with(train_result)
    mlflow.model_monitoring_repository.save_run.assert_called_once_with(train_result)

    # Verify response - now returns TrainModelResult directly
    assert response == train_result
    assert response.run_name == 'test_experiment-1'
    assert train_result.run_name == 'test_experiment-1'


@mark.asyncio
async def test_save_model_get_next_run_name_error(mlflow):
    """Test save_model handles error during get_next_run_name."""
    from model_manager.utils.models.train_model_params import TrainModelParams
    from model_manager.utils.models.train_model_result import TrainModelResult

    params = MagicMock(spec=TrainModelParams)
    params.experiment_name = 'test_experiment'

    train_result = MagicMock(spec=TrainModelResult)
    train_result.params = params

    # Mock error in get_next_run_name
    mlflow.model_monitoring_repository.get_next_run_name.side_effect = Exception(
        'MLflow connection error'
    )

    input_data = {
        **metadata,
        'train_result': train_result,
    }

    # Call the method - should raise exception
    with pytest.raises(Exception, match='MLflow connection error'):
        await mlflow.save_model(input_data)

    # Verify notification was sent
    mlflow.send_notification.assert_called_once_with(
        metadata=metadata['metadata'],
        notification_id='SAVE_MODEL_ERROR',
        message=ANY,
        block='save_model',
        level=NotificationLevel.ERROR,
        attachment_content=ANY,
    )


@mark.asyncio
async def test_save_model_generate_artifacts_error(mlflow):
    """Test save_model handles error during generate_artifacts."""
    from model_manager.utils.models.train_model_params import TrainModelParams
    from model_manager.utils.models.train_model_result import TrainModelResult

    params = MagicMock(spec=TrainModelParams)
    params.experiment_name = 'test_experiment'

    train_result = MagicMock(spec=TrainModelResult)
    train_result.params = params

    # Mock successful get_next_run_name but error in generate_artifacts
    mlflow.model_monitoring_repository.get_next_run_name.return_value = 'test_experiment-1'
    mlflow.model_monitoring_repository.generate_artifacts.side_effect = FileNotFoundError(
        'Reports directory does not exist'
    )

    input_data = {
        **metadata,
        'train_result': train_result,
    }

    # Call the method - should raise exception
    with pytest.raises(FileNotFoundError, match='Reports directory does not exist'):
        await mlflow.save_model(input_data)

    # Verify notification was sent
    mlflow.send_notification.assert_called_once_with(
        metadata=metadata['metadata'],
        notification_id='SAVE_MODEL_ERROR',
        message=ANY,
        block='save_model',
        level=NotificationLevel.ERROR,
        attachment_content=ANY,
    )


@mark.asyncio
async def test_save_model_save_run_error(mlflow):
    """Test save_model handles error during save_run."""
    from model_manager.utils.models.train_model_params import TrainModelParams
    from model_manager.utils.models.train_model_result import TrainModelResult

    params = MagicMock(spec=TrainModelParams)
    params.experiment_name = 'test_experiment'

    train_result = MagicMock(spec=TrainModelResult)
    train_result.params = params

    # Mock successful get_next_run_name and generate_artifacts but error in save_run
    mlflow.model_monitoring_repository.get_next_run_name.return_value = 'test_experiment-1'
    mlflow.model_monitoring_repository.generate_artifacts.return_value = train_result
    mlflow.model_monitoring_repository.save_run.side_effect = ValueError(
        'One or more metrics (MSE, R2, MAE) are None'
    )

    input_data = {
        **metadata,
        'train_result': train_result,
    }

    # Call the method - should raise exception
    with pytest.raises(ValueError, match=r'One or more metrics \(MSE, R2, MAE\) are None'):
        await mlflow.save_model(input_data)

    # Verify notification was sent
    mlflow.send_notification.assert_called_once_with(
        metadata=metadata['metadata'],
        notification_id='SAVE_MODEL_ERROR',
        message=ANY,
        block='save_model',
        level=NotificationLevel.ERROR,
        attachment_content=ANY,
    )


@mark.asyncio
async def test_save_model_missing_metadata(mlflow):
    """Test save_model handles missing metadata gracefully."""
    from model_manager.utils.models.train_model_params import TrainModelParams
    from model_manager.utils.models.train_model_result import TrainModelResult

    params = MagicMock(spec=TrainModelParams)
    params.experiment_name = 'test_experiment'

    train_result = MagicMock(spec=TrainModelResult)
    train_result.params = params

    # Mock repository methods
    mlflow.model_monitoring_repository.get_next_run_name.return_value = 'test_experiment-1'
    mlflow.model_monitoring_repository.generate_artifacts.return_value = train_result
    mlflow.model_monitoring_repository.save_run.return_value = None

    # Input data without metadata
    input_data = {
        'train_result': train_result,
    }

    # Call the method
    response = await mlflow.save_model(input_data)

    # Verify it still works (metadata defaults to {}) - returns TrainModelResult directly
    assert response == train_result
    assert response.run_name == 'test_experiment-1'


@mark.asyncio
async def test_save_model_complete_flow(mlflow):
    """Test save_model complete flow with all steps."""
    from model_manager.utils.models.train_model_params import TrainModelParams
    from model_manager.utils.models.train_model_result import TrainModelResult

    params = MagicMock(spec=TrainModelParams)
    params.experiment_name = 'production_model'

    train_result = MagicMock(spec=TrainModelResult)
    train_result.params = params
    train_result.run_name = None
    train_result.run_dir = None
    train_result.report_path = None

    # Mock complete flow
    mlflow.model_monitoring_repository.get_next_run_name.return_value = 'production_model-5'

    # After generate_artifacts, paths should be set
    updated_result = MagicMock(spec=TrainModelResult)
    updated_result.params = params
    updated_result.run_name = 'production_model-5'
    updated_result.run_dir = '/reports/production_model-5_20231010'
    updated_result.report_path = '/reports/production_model-5_20231010/report.html'
    updated_result.train_data_path = '/reports/production_model-5_20231010/train_data.csv'
    updated_result.test_data_path = '/reports/production_model-5_20231010/test_data.csv'

    mlflow.model_monitoring_repository.generate_artifacts.return_value = updated_result
    mlflow.model_monitoring_repository.save_run.return_value = None

    input_data = {
        **metadata,
        'train_result': train_result,
    }

    # Call the method
    response = await mlflow.save_model(input_data)

    # Verify complete flow
    mlflow.model_monitoring_repository.get_next_run_name.assert_called_once_with('production_model')
    mlflow.model_monitoring_repository.generate_artifacts.assert_called_once()
    mlflow.model_monitoring_repository.save_run.assert_called_once_with(updated_result)

    # Verify response - returns TrainModelResult directly
    assert response == updated_result
    assert response.run_name == 'production_model-5'
    assert response.run_dir is not None
    assert response.report_path is not None


# ============================================================================
# Tests for cleanup_run_directory
# ============================================================================


@mark.asyncio
async def test_cleanup_run_directory_success(mlflow):
    """Test successful cleanup of run directory."""
    input_data = {
        **metadata,
        'run_dir': 'test_run_dir',
    }

    # Patch os and shutil inside the activity method
    with (
        patch('os.path.exists', return_value=True) as mock_exists,
        patch('shutil.rmtree') as mock_rmtree,
    ):
        # Call the method
        await mlflow.cleanup_run_directory(input_data)

        # Verify directory existence was checked
        mock_exists.assert_called_once_with('test_run_dir')

        # Verify shutil.rmtree was called
        mock_rmtree.assert_called_once_with('test_run_dir')


@mark.asyncio
async def test_cleanup_run_directory_already_deleted(mlflow):
    """Test cleanup when directory is already deleted (idempotent)."""
    input_data = {
        **metadata,
        'run_dir': 'already_deleted_dir',
    }

    # Patch os and shutil inside the activity method
    with (
        patch('os.path.exists', return_value=False) as mock_exists,
        patch('shutil.rmtree') as mock_rmtree,
    ):
        # Call the method - should not raise error
        await mlflow.cleanup_run_directory(input_data)

        # Verify directory existence was checked
        mock_exists.assert_called_once_with('already_deleted_dir')

        # Verify shutil.rmtree was NOT called
        mock_rmtree.assert_not_called()


@mark.asyncio
async def test_cleanup_run_directory_no_run_dir(mlflow):
    """Test cleanup when no run_dir is provided."""
    input_data = {
        **metadata,
        # No run_dir key
    }

    # Call the method - should not raise error
    await mlflow.cleanup_run_directory(input_data)

    # Should complete without errors


@mark.asyncio
async def test_cleanup_run_directory_empty_run_dir(mlflow):
    """Test cleanup when run_dir is empty string."""
    input_data = {
        **metadata,
        'run_dir': '',
    }

    # Call the method - should not raise error
    await mlflow.cleanup_run_directory(input_data)

    # Should complete without errors


@mark.asyncio
async def test_cleanup_run_directory_none_run_dir(mlflow):
    """Test cleanup when run_dir is None."""
    input_data = {
        **metadata,
        'run_dir': None,
    }

    # Call the method - should not raise error
    await mlflow.cleanup_run_directory(input_data)

    # Should complete without errors


@mark.asyncio
async def test_cleanup_run_directory_error(mlflow):
    """Test cleanup handles errors correctly."""
    input_data = {
        **metadata,
        'run_dir': 'error_dir',
    }

    # Patch os and shutil inside the activity method
    with (
        patch('os.path.exists', return_value=True),
        patch('shutil.rmtree', side_effect=PermissionError('Permission denied')),
    ):
        # Call the method - should raise exception
        with pytest.raises(PermissionError, match='Permission denied'):
            await mlflow.cleanup_run_directory(input_data)

        # Verify notification was sent
        mlflow.send_notification.assert_called_once_with(
            metadata=metadata['metadata'],
            notification_id='CLEANUP_RUN_DIRECTORY_ERROR',
            message=ANY,
            block='cleanup_run_directory',
            level=NotificationLevel.ERROR,
            attachment_content=ANY,
        )


@mark.asyncio
async def test_cleanup_run_directory_missing_metadata(mlflow):
    """Test cleanup handles missing metadata gracefully."""
    input_data = {
        'run_dir': 'no_metadata_dir',
    }

    # Patch os and shutil inside the activity method
    with patch('os.path.exists', return_value=True), patch('shutil.rmtree') as mock_rmtree:
        # Call the method
        await mlflow.cleanup_run_directory(input_data)

        # Verify directory was deleted
        mock_rmtree.assert_called_once_with('no_metadata_dir')


@mark.asyncio
async def test_cleanup_run_directory_oserror(mlflow):
    """Test cleanup handles OSError correctly."""
    input_data = {
        **metadata,
        'run_dir': 'os_error_dir',
    }

    # Patch os and shutil inside the activity method
    with (
        patch('os.path.exists', return_value=True),
        patch('shutil.rmtree', side_effect=OSError('Directory not empty')),
    ):
        # Call the method - should raise exception
        with pytest.raises(OSError, match='Directory not empty'):
            await mlflow.cleanup_run_directory(input_data)

        # Verify notification was sent
        mlflow.send_notification.assert_called_once()
        call_args = mlflow.send_notification.call_args[1]
        assert call_args['notification_id'] == 'CLEANUP_RUN_DIRECTORY_ERROR'
        assert call_args['level'] == NotificationLevel.ERROR
        assert 'Directory not empty' in call_args['message']