from unittest.mock import ANY, AsyncMock, MagicMock, call, patch

import pytest_asyncio
from pandas import DataFrame
from pytest import mark
from sientia_do.notifications.models import NotificationLevel

from laborious.activities.opc import OPC

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


def test__init__():
    servers = {'server1': {'id': 'server1'}}
    opc = OPC(
        opc_servers=servers,
        logger=MagicMock(),
        notification_handler=MagicMock(),
        metrics_controller=AsyncMock(),
    )

    assert opc.opc_servers == servers
    assert opc.opc_repository == {}


@mark.asyncio
@patch('laborious.activities.opc.OpcRepository')
@patch('laborious.activities.opc.OPC.send_notification_async')
async def test_init_opc(mock_send_notification, mock_opc_repository):
    mock_logger = MagicMock()
    mock_metrics_controller = AsyncMock()
    server1 = MagicMock(
        connect=AsyncMock(return_value=(True, {})), write_data=AsyncMock(return_value=(True, {}))
    )
    server2 = MagicMock(
        connect=AsyncMock(return_value=(True, {})), write_data=AsyncMock(return_value=(True, {}))
    )
    server3 = MagicMock(
        connect=AsyncMock(
            return_value=(
                False,
                {
                    'notification_id': 'OPC_CONNECTION_ERROR_server3',
                    'message': 'Failed to connect to OPC server: Test error',
                    'block': 'opc_repository',
                    'level': NotificationLevel.ERROR,
                    'attachment_content': 'Test error',
                },
            )
        ),
        write_data=AsyncMock(return_value=(True, {})),
    )
    mock_opc_repository.side_effect = [server1, server2, server3]
    mock_notification_handler = MagicMock()
    servers = {
        'server1': {
            'server_name': 'server1',
            'id': 'server1',
            'url': 'http://localhost:8080',
            'server_uri': 'opc.tcp://localhost:4840',
            'cert_path': '',
            'private_key_path': '',
            'server_cert_path': '',
            'reconnection_interval': 60,
        },
        'server2': {
            'server_name': 'server2',
            'id': 'server2',
            'url': 'http://localhost:8080',
            'server_uri': 'opc.tcp://localhost:4840',
            'cert_path': '',
            'private_key_path': '',
            'server_cert_path': '',
            'reconnection_interval': 60,
        },
        'server3': {
            'server_name': 'server3',
            'id': 'server3',
            'url': 'http://localhost:8080',
            'server_uri': 'opc.tcp://localhost:4840',
            'cert_path': '',
            'private_key_path': '',
            'server_cert_path': '',
            'reconnection_interval': 60,
        },
    }
    opc = OPC(
        opc_servers=servers,
        logger=mock_logger,
        notification_handler=mock_notification_handler,
        metrics_controller=mock_metrics_controller,
    )
    await opc.init_opc()

    assert opc.opc_servers == servers
    assert opc.logger == mock_logger
    assert opc.notification_handler == mock_notification_handler
    assert opc.opc_repository['server1'] == server1
    assert opc.opc_repository['server2'] == server2

    mock_opc_repository.assert_has_calls(
        [
            call(
                opc_id='server1',
                server_name='server1',
                url='http://localhost:8080',
                logger=mock_logger,
                server_uri='opc.tcp://localhost:4840',
                cert_path='',
                private_key_path='',
                server_cert_path='',
                notification_handler=mock_notification_handler,
                reconnection_interval=60,
                metrics_controller=mock_metrics_controller,
            ),
        ]
    )
    mock_opc_repository.assert_has_calls(
        [
            call(
                opc_id='server2',
                server_name='server2',
                url='http://localhost:8080',
                logger=mock_logger,
                server_uri='opc.tcp://localhost:4840',
                cert_path='',
                private_key_path='',
                server_cert_path='',
                notification_handler=mock_notification_handler,
                reconnection_interval=60,
                metrics_controller=mock_metrics_controller,
            )
        ]
    )

    server1.connect.assert_called_once()
    server2.connect.assert_called_once()

    mock_send_notification.assert_has_calls(
        [
            call(
                metadata={
                    'model_id': '-',
                    'model_name': '-',
                    'workflow_name': '-',
                    'schedule_name': 'INITIALIZATION',
                },
                notification_id='OPC_CONNECTION_ERROR_server3',
                message='Failed to connect to OPC server: Test error',
                block='opc_repository',
                level=NotificationLevel.ERROR,
                attachment_content=ANY,
            )
        ]
    )


@pytest_asyncio.fixture
@patch('laborious.activities.opc.OpcRepository')
async def opc(mock_opc_repository):
    servers = {
        'server1': {
            'id': 'server1',
            'server_name': 'server1',
            'url': 'http://localhost:8080',
            'server_uri': 'opc.tcp://localhost:4840',
            'cert_path': '',
            'private_key_path': '',
            'server_cert_path': '',
            'reconnection_interval': 60,
        }
    }

    mock_opc_repository.return_value.write_data = AsyncMock(return_value=(True, {}))
    mock_opc_repository.return_value.connect = AsyncMock(return_value=(True, {}))
    opc = OPC(
        opc_servers=servers,
        logger=MagicMock(),
        notification_handler=MagicMock(),
        metrics_controller=AsyncMock(),
    )
    await opc.init_opc()
    opc.send_notification = MagicMock()
    opc.send_notification_async = AsyncMock()
    opc.emit_metric = AsyncMock()
    return opc


WRITE_DATA_CASES = [
    ('tag1', 'int', 50),
    ('tag2', 'float', 50.5),
    ('tag3', 'bool', True),
    ('tag4', 'string', 'test'),
]


@mark.parametrize('tag,data_type,data', WRITE_DATA_CASES)
@mark.asyncio
async def test_write_data_success(opc, tag, data_type, data):
    opc.opc_repository['server1'].write_data.return_value = (True, {'response_time': 0.1})

    result = await opc.write_data(
        server_id='server1',
        tag=tag,
        data=data,
        data_type=data_type,
        tag_type='prediction',
        metadata=metadata,
    )
    assert result == 0.1
    opc.opc_repository['server1'].write_data.assert_called_once_with(
        tag, data, data_type, opc.logger, metadata
    )


@mark.asyncio
async def test_write_data_failed(opc):
    opc.opc_repository['server1'].write_data.return_value = (
        False,
        {
            'notification_id': 'OPC_WRITE_DATA_ERROR_server1',
            'message': 'Failed to write data to OPC server: Test error',
            'block': 'opc_repository',
            'level': NotificationLevel.ERROR,
            'attachment_content': 'Test error',
        },
    )

    result = await opc.write_data(
        server_id='server1',
        tag='tag1',
        data=50,
        data_type='int',
        tag_type='prediction',
        metadata=metadata,
    )
    assert result is None

    opc.send_notification_async.assert_called_once_with(
        metadata=metadata,
        notification_id='OPC_WRITE_DATA_ERROR_server1',
        message='Failed to write data to OPC server: Test error',
        block='opc_repository',
        level=NotificationLevel.ERROR,
        attachment_content=ANY,
    )


@mark.asyncio
async def test_write_data_exception(opc):
    opc.opc_repository['server1'].write_data.side_effect = Exception('Test error')

    try:
        await opc.write_data(
            server_id='server1',
            tag='tag1',
            data=50,
            data_type='int',
            tag_type='prediction',
            metadata=metadata,
        )

    except Exception:
        opc.send_notification_async.assert_called_once_with(
            metadata=metadata,
            notification_id='WRITE_OPC_PREDICTION_ERROR',
            message='Error writing data to OPC server: Test error',
            block='write_opc_data',
            level=NotificationLevel.ERROR,
            attachment_content=ANY,
        )

    else:
        raise AssertionError('Expected an exception to be raised')


@mark.asyncio
async def test_manage_output_tags_success(opc):
    opc.write_data = AsyncMock(return_value=0.1)

    data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
    config = {
        'prediction_tags': {'tag1': {'data_type': 'float'}},
        'confidence_tags': {'tag2': {'data_type': 'float'}},
    }

    output_data, opc_metrics = await opc.manage_output_tags(
        server_id='server1',
        config=config,
        data=data,
        metadata=metadata['metadata'],
    )

    assert output_data is True
    assert opc_metrics == {'tag1': 0.1, 'tag2': 0.1}
    opc.write_data.assert_has_calls(
        [
            call(
                server_id='server1',
                tag='tag1',
                data=0.75,
                data_type='float',
                tag_type='prediction',
                metadata=metadata['metadata'],
            ),
            call(
                server_id='server1',
                tag='tag2',
                data=0.95,
                data_type='float',
                tag_type='confidence',
                metadata=metadata['metadata'],
            ),
        ]
    )


@mark.asyncio
@mark.parametrize('side_effect', [[0.1, None], [None, 0.2]])
async def test_manage_output_tags_failed(opc, side_effect):
    opc.write_data = AsyncMock(side_effect=side_effect)
    data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
    config = {
        'prediction_tags': {'tag1': {'data_type': 'float'}},
        'confidence_tags': {'tag2': {'data_type': 'float'}},
    }
    output_data, opc_metrics = await opc.manage_output_tags(
        server_id='server1',
        config=config,
        data=data,
        metadata=metadata['metadata'],
    )
    assert output_data is False
    assert opc_metrics == {'tag1': side_effect[0], 'tag2': side_effect[1]}
    opc.write_data.assert_has_calls(
        [
            call(
                server_id='server1',
                tag='tag1',
                data=0.75,
                data_type='float',
                tag_type='prediction',
                metadata=metadata['metadata'],
            ),
            call(
                server_id='server1',
                tag='tag2',
                data=0.95,
                data_type='float',
                tag_type='confidence',
                metadata=metadata['metadata'],
            ),
        ]
    )


@mark.asyncio
async def test_manage_output_tags_do_nothing(opc):
    opc.write_data = AsyncMock(return_value=0.1)
    data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
    config = {
        '_invalid_key': {'tag1': {'data_type': 'float'}},
    }
    output_data, opc_metrics = await opc.manage_output_tags(
        server_id='server1',
        config=config,
        data=data,
        metadata=metadata['metadata'],
    )
    assert output_data is True
    assert opc_metrics == {}
    opc.write_data.assert_not_called()


@mark.asyncio
@patch('laborious.activities.opc.DataFrame')
async def test_write_opc_data_success(mock_dataframe, opc):
    # Arrange
    input_data = {
        **metadata,
        'data': {'prediction': [0.75], 'prediction_confidence': [0.95]},
        'opc_output_config': {
            'server1': {
                'prediction_tags': {'tag1': {'data_type': 'float'}},
                'confidence_tags': {'tag2': {'data_type': 'float'}},
            }
        },
    }

    # Act
    opc.manage_output_tags = AsyncMock(return_value=(True, {'tag1': 0.1, 'tag2': 0.2}))

    opc.process_confidence = MagicMock(return_value={'data': 'data'})
    output_data, opc_metrics = await opc.write_opc_data(input_data)

    # Assert
    assert output_data == {'data': 'data'}
    assert opc_metrics == {'server1': {'tag1': 0.1, 'tag2': 0.2}}
    opc.manage_output_tags.assert_called_once_with(
        'server1',
        input_data['opc_output_config']['server1'],
        mock_dataframe.return_value,
        metadata['metadata'],
    )
    opc.process_confidence.assert_called_once_with(
        mock_dataframe.return_value,
        True,
        metadata['metadata'],
    )


@mark.asyncio
async def test_write_opc_data_empty_config(opc):
    # Arrange
    input_data = {
        **metadata,
        'data': {'prediction': [0.75], 'prediction_confidence': [0.95]},
        'opc_servers': ['server1'],
        'opc_output_config': {'server1': {'prediction_tags': {}, 'confidence_tags': {}}},
    }

    # Act
    await opc.write_opc_data(input_data)

    # Assert
    opc.opc_repository['server1'].write_data.assert_not_called()


@mark.asyncio
async def test_write_opc_data_no_validate_server(opc):
    opc.validate_server = AsyncMock(return_value=False)
    input_data = {
        **metadata,
        'data': {'prediction': [0.75], 'prediction_confidence': [0.95]},
        'opc_output_config': {
            'server1': {
                'prediction_tags': {'tag1': {'data_type': 'float'}},
                'confidence_tags': {'tag2': {'data_type': 'float'}},
            }
        },
    }

    # Act
    await opc.write_opc_data(input_data)

    # Assert
    opc.opc_repository['server1'].write_data.assert_not_called()


@mark.parametrize(
    'data,success,expected',
    [
        (DataFrame({'prediction_confidence': [0]}), True, 0),
        (DataFrame({'prediction_confidence': [0]}), False, 12),
    ],
)
def test_process_confidence(opc, data, success, expected):
    # Act
    result = opc.process_confidence(data, success, metadata)

    # Assert
    assert result['prediction_confidence'][0] == expected


@mark.asyncio
async def test_validate_server(opc):
    assert await opc.validate_server('server1', metadata) is True
    assert await opc.validate_server('server2', metadata) is False


@mark.asyncio
async def test_close(opc):
    opc.opc_repository['server1'].disconnect = AsyncMock(return_value=True)
    await opc.close()
    opc.opc_repository['server1'].disconnect.assert_called_once()