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

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

import pytest
from pytest import mark

from model_manager.activities.training import Training
from model_manager.utils.models.train_model_params import TrainModelParams
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 = TrainModelParams(
        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',
        removed_intervals=[],
    )

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

    # Execute
    result = await training.train_model(input_data)

    # Assertions - now returns TrainModelResult directly
    assert result == mock_final_result
    assert result.mse_val == 0.5
    assert result.mae_val == 0.3
    assert result.r2_val == 0.95

    # 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)
    train_params = TrainModelParams(
        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',
        removed_intervals=[],
    )

    input_data = {
        'metadata': {},
        'uploaded_file': 'not_a_bytesio',
        'train_params': train_params,
    }

    # Should raise ValueError
    with pytest.raises(ValueError, match='uploaded_file must be BytesIO'):
        await training.train_model(input_data)

    # 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')
    train_params = TrainModelParams(
        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',
        removed_intervals=[],
    )

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

    # Should raise ValueError
    with pytest.raises(ValueError, match='Training data is empty'):
        await training.train_model(input_data)

    # 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')
    train_params = TrainModelParams(
        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',
        removed_intervals=[],
    )

    input_data = {
        'metadata': {'workflow_id': 'test-789', 'experiment_run_id': 789},
        'uploaded_file': uploaded_file,
        'train_params': train_params,
    }

    # Should raise Exception
    with pytest.raises(Exception, match='Database connection failed'):
        await training.train_model(input_data)

    # 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_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')
    train_params = TrainModelParams(
        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',
        removed_intervals=[],
    )

    input_data = {
        'metadata': {},
        'uploaded_file': uploaded_file,
        'train_params': train_params,
    }

    # Should raise Exception
    with pytest.raises(Exception, match='Metric calculation failed'):
        await training.train_model(input_data)

    # 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_invalid_train_params_type(mock_training_repository_class):
    """Test training with invalid train_params type (dict instead of TrainModelParams)."""
    mock_repository = MagicMock()
    mock_training_repository_class.return_value = mock_repository

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

    uploaded_file = BytesIO(b'test,data\n1,2\n3,4')

    # Invalid train_params type (dict instead of TrainModelParams object)
    input_data = {
        'metadata': {'workflow_id': 'test-invalid'},
        'uploaded_file': uploaded_file,
        'train_params': {
            'experiment_run_id': 123,
            'target_variable': 'price',
            'variable_columns': ['feature1'],
        },  # This is a dict, not TrainModelParams
    }

    # Should raise ValueError
    with pytest.raises(ValueError, match='train_params must be TrainModelParams.*dict'):
        await training.train_model(input_data)


# ============================================================================
# Tests for validate_train_params
# ============================================================================


@mark.asyncio
async def test_validate_train_params_success():
    """Test successful validation of training parameters."""
    logger = MagicMock()
    notification_handler = MagicMock()
    training = Training(logger=logger, notification_handler=notification_handler)

    training.info = MagicMock()

    input_data = {
        'metadata': {'workflow_id': 'test-123'},
        'experiment_run_id': 456,
        'target_variable': 'price',
        'variable_columns': ['feature1', 'feature2', 'price'],  # target_variable must be in list
        '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, 'price': 0.0},
        'upp_lim': {'feature1': 100.0, 'feature2': 100.0, 'price': 1000.0},
        'window': 10,
        'experiment_name': 'test_experiment',
        'removed_intervals': [],
    }

    result = await training.validate_train_params(input_data)

    assert isinstance(result, TrainModelParams)
    assert result.experiment_run_id == 456
    assert result.target_variable == 'price'
    assert result.variable_columns == ['feature1', 'feature2', 'price']
    assert result.train_size == 80
    assert result.experiment_name == 'test_experiment'
    assert training.info.call_count == 2


@mark.asyncio
async def test_validate_train_params_missing_required_field():
    """Test validation fails when required field is missing."""
    logger = MagicMock()
    notification_handler = MagicMock()
    training = Training(logger=logger, notification_handler=notification_handler)

    training.info = MagicMock()
    training.error = MagicMock()

    input_data = {
        'metadata': {'workflow_id': 'test-123'},
        'experiment_run_id': 456,
        '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',
        'removed_intervals': [],
    }

    with pytest.raises(ValueError, match='target_variable'):
        await training.validate_train_params(input_data)

    notification_handler.send_notification.assert_called_once()


@mark.asyncio
async def test_validate_train_params_invalid_type():
    """Test validation fails when field has invalid type."""
    logger = MagicMock()
    notification_handler = MagicMock()
    training = Training(logger=logger, notification_handler=notification_handler)

    training.info = MagicMock()
    training.error = MagicMock()

    input_data = {
        'metadata': {'workflow_id': 'test-invalid'},
        'experiment_run_id': 456,
        'target_variable': 'price',
        'variable_columns': ['feature1'],
        'train_size': 'invalid',
        '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',
        'removed_intervals': [],
    }

    with pytest.raises((ValueError, TypeError)):
        await training.validate_train_params(input_data)

    notification_handler.send_notification.assert_called_once()


@mark.asyncio
async def test_validate_train_params_empty_input():
    """Test validation fails with empty input."""
    logger = MagicMock()
    notification_handler = MagicMock()
    training = Training(logger=logger, notification_handler=notification_handler)

    training.info = MagicMock()
    training.error = MagicMock()

    input_data = {'metadata': {}}

    with pytest.raises(ValueError):
        await training.validate_train_params(input_data)

    notification_handler.send_notification.assert_called_once()


@mark.asyncio
async def test_validate_train_params_without_metadata():
    """Test validation works even without metadata key."""
    logger = MagicMock()
    notification_handler = MagicMock()
    training = Training(logger=logger, notification_handler=notification_handler)

    training.info = MagicMock()

    input_data = {
        'experiment_run_id': 789,
        'target_variable': 'temperature',
        'variable_columns': ['sensor1', 'temperature'],  # target_variable must be in list
        'train_size': 75,
        'shuffle': False,
        'use_scaler': True,
        'include_ar': True,
        'bucket_name': 'sensors',
        'file_name': 'data.csv',
        'line_separator': '\n',
        'decimal_separator': '.',
        'lag_train': 2,
        'lag_val': 2,
        'rem_static_win': True,
        'low_lim': {'sensor1': -50.0, 'temperature': -50.0},
        'upp_lim': {'sensor1': 150.0, 'temperature': 150.0},
        'window': 20,
        'experiment_name': 'sensor_experiment',
        'removed_intervals': [],
    }

    result = await training.validate_train_params(input_data)

    assert isinstance(result, TrainModelParams)
    assert result.experiment_run_id == 789
    assert result.target_variable == 'temperature'
    assert result.experiment_name == 'sensor_experiment'