"""Unit tests for Training activity."""

from io import BytesIO
from unittest.mock import MagicMock, patch

from pytest import mark

from model_manager.activities.training import Training
from model_manager.utils.models.train_model_result import TrainModelResult


@mark.asyncio
@patch('model_manager.activities.training.TrainingRepository')
async def test_train_model_success(mock_training_repository_class):
    """Test successful model training."""
    # Create mock repository instance
    mock_repository = MagicMock()
    mock_training_repository_class.return_value = mock_repository

    # Create mock train result
    mock_train_result = MagicMock(spec=TrainModelResult)
    mock_train_result.mse_val = 0.5
    mock_train_result.mae_val = 0.3
    mock_train_result.r2_val = 0.95

    mock_final_result = MagicMock(spec=TrainModelResult)
    mock_final_result.mse_val = 0.5
    mock_final_result.mae_val = 0.3
    mock_final_result.r2_val = 0.95

    # Setup repository mocks
    mock_repository.train.return_value = mock_train_result
    mock_repository.after_train_calculation.return_value = mock_final_result

    # Create Training instance
    logger = MagicMock()
    notification_handler = MagicMock()
    training = Training(logger=logger, notification_handler=notification_handler)

    # Mock inherited methods
    training.info = MagicMock()

    # Test data
    uploaded_file = BytesIO(b'test,data\n1,2\n3,4')
    train_params_dict = {
        'experiment_run_id': 123,
        'target_variable': 'price',
        'variable_columns': ['feature1', 'feature2'],
        'train_size': 80,
        'shuffle': True,
        'use_scaler': True,
        'include_ar': False,
        'bucket_name': 'test-bucket',
        'file_name': 'test.csv',
        'line_separator': '\n',
        'decimal_separator': '.',
        'lag_train': 1,
        'lag_val': 1,
        'rem_static_win': False,
        'low_lim': {'feature1': 0.0, 'feature2': 0.0},
        'upp_lim': {'feature1': 100.0, 'feature2': 100.0},
        'window': 10,
        'experiment_name': 'test_experiment',
        'experiment_description': 'Test experiment',
        'removed_intervals': [],
    }

    input_data = {
        'metadata': {'workflow_id': 'test-123'},
        'uploaded_file': uploaded_file,
        'train_params': train_params_dict,
    }

    # Execute
    result = await training.train_model(input_data)

    # Assertions
    assert result['success'] is True
    assert result['result'] == mock_final_result
    assert result['error_message'] is None

    # Verify repository calls
    mock_repository.train.assert_called_once()
    mock_repository.after_train_calculation.assert_called_once()


@mark.asyncio
@patch('model_manager.activities.training.TrainingRepository')
async def test_train_model_invalid_file_type(mock_training_repository_class):
    """Test training with invalid file type."""
    mock_repository = MagicMock()
    mock_training_repository_class.return_value = mock_repository

    logger = MagicMock()
    notification_handler = MagicMock()
    training = Training(logger=logger, notification_handler=notification_handler)

    # Invalid file type (string instead of BytesIO)
    input_data = {
        'metadata': {},
        'uploaded_file': 'not_a_bytesio',
        'train_params': {
            'experiment_run_id': 123,
            'target_variable': 'price',
            'variable_columns': ['feature1'],
            'train_size': 80,
            'shuffle': True,
            'use_scaler': False,
            'include_ar': False,
            'bucket_name': 'test',
            'file_name': 'test.csv',
            'line_separator': '\n',
            'decimal_separator': '.',
            'lag_train': 1,
            'lag_val': 1,
            'rem_static_win': False,
            'low_lim': {'feature1': 0.0},
            'upp_lim': {'feature1': 100.0},
            'window': 10,
            'experiment_name': 'test_experiment',
            'experiment_description': 'Test experiment',
            'removed_intervals': [],
        },
    }

    result = await training.train_model(input_data)

    assert result['success'] is False
    assert result['result'] is None
    assert 'uploaded_file must be BytesIO' in result['error_message']
    # Verify notification was sent (via BaseActivity)
    notification_handler.send_notification.assert_called_once()


@mark.asyncio
@patch('model_manager.activities.training.TrainingRepository')
async def test_train_model_training_error(mock_training_repository_class):
    """Test training failure during model training."""
    mock_repository = MagicMock()
    mock_training_repository_class.return_value = mock_repository

    # Setup repository to raise error
    mock_repository.train.side_effect = ValueError('Training data is empty')

    logger = MagicMock()
    notification_handler = MagicMock()
    training = Training(logger=logger, notification_handler=notification_handler)

    uploaded_file = BytesIO(b'test,data\n')
    input_data = {
        'metadata': {'workflow_id': 'test-456'},
        'uploaded_file': uploaded_file,
        'train_params': {
            'experiment_run_id': 456,
            'target_variable': 'price',
            'variable_columns': ['feature1'],
            'train_size': 80,
            'shuffle': True,
            'use_scaler': False,
            'include_ar': False,
            'bucket_name': 'test',
            'file_name': 'test.csv',
            'line_separator': '\n',
            'decimal_separator': '.',
            'lag_train': 1,
            'lag_val': 1,
            'rem_static_win': False,
            'low_lim': {'feature1': 0.0},
            'upp_lim': {'feature1': 100.0},
            'window': 10,
            'experiment_name': 'test_experiment',
            'experiment_description': 'Test experiment',
            'removed_intervals': [],
        },
    }

    result = await training.train_model(input_data)

    assert result['success'] is False
    assert result['result'] is None
    assert 'Training data is empty' in result['error_message']
    # Verify notification was sent (via BaseActivity)
    notification_handler.send_notification.assert_called_once()


@mark.asyncio
@patch('model_manager.activities.training.TrainingRepository')
async def test_train_model_sends_notification_on_error(mock_training_repository_class):
    """Test that notification is sent when training fails."""
    mock_repository = MagicMock()
    mock_training_repository_class.return_value = mock_repository

    mock_repository.train.side_effect = Exception('Database connection failed')

    logger = MagicMock()
    notification_handler = MagicMock()
    training = Training(logger=logger, notification_handler=notification_handler)

    uploaded_file = BytesIO(b'test,data\n1,2')
    input_data = {
        'metadata': {'workflow_id': 'test-789', 'experiment_run_id': 789},
        'uploaded_file': uploaded_file,
        'train_params': {
            'experiment_run_id': 789,
            'target_variable': 'price',
            'variable_columns': ['feature1'],
            'train_size': 80,
            'shuffle': True,
            'use_scaler': False,
            'include_ar': False,
            'bucket_name': 'test',
            'file_name': 'test.csv',
            'line_separator': '\n',
            'decimal_separator': '.',
            'lag_train': 1,
            'lag_val': 1,
            'rem_static_win': False,
            'low_lim': {'feature1': 0.0},
            'upp_lim': {'feature1': 100.0},
            'window': 10,
            'experiment_name': 'test_experiment',
            'experiment_description': 'Test experiment',
            'removed_intervals': [],
        },
    }

    result = await training.train_model(input_data)

    # Verify notification was sent (via BaseActivity)
    notification_handler.send_notification.assert_called_once()

    # Verify result - the important part is that error was caught and returned
    assert result['success'] is False
    assert result['result'] is None
    assert 'Database connection failed' in result['error_message']


@mark.asyncio
@patch('model_manager.activities.training.TrainingRepository')
async def test_train_model_after_calculation_error(mock_training_repository_class):
    """Test training failure during post-training calculations."""
    mock_repository = MagicMock()
    mock_training_repository_class.return_value = mock_repository

    # Train succeeds but after_calculation fails
    mock_train_result = MagicMock(spec=TrainModelResult)
    mock_repository.train.return_value = mock_train_result
    mock_repository.after_train_calculation.side_effect = Exception('Metric calculation failed')

    logger = MagicMock()
    notification_handler = MagicMock()
    training = Training(logger=logger, notification_handler=notification_handler)

    uploaded_file = BytesIO(b'test,data\n1,2\n3,4')
    input_data = {
        'metadata': {},
        'uploaded_file': uploaded_file,
        'train_params': {
            'experiment_run_id': 999,
            'target_variable': 'price',
            'variable_columns': ['feature1'],
            'train_size': 80,
            'shuffle': True,
            'use_scaler': False,
            'include_ar': False,
            'bucket_name': 'test',
            'file_name': 'test.csv',
            'line_separator': '\n',
            'decimal_separator': '.',
            'lag_train': 1,
            'lag_val': 1,
            'rem_static_win': False,
            'low_lim': {'feature1': 0.0},
            'upp_lim': {'feature1': 100.0},
            'window': 10,
            'experiment_name': 'test_experiment',
            'experiment_description': 'Test experiment',
            'removed_intervals': [],
        },
    }

    result = await training.train_model(input_data)

    assert result['success'] is False
    assert result['result'] is None
    assert 'Metric calculation failed' in result['error_message']
    # Verify notification was sent (via BaseActivity)
    notification_handler.send_notification.assert_called_once()