"""
End-to-end tests for PredictionsBatch workflow - Prediction Process scenarios.
"""

from decimal import Decimal
from unittest.mock import MagicMock, patch

import numpy as np
import pandas as pd
import pytest
from sqlalchemy import text
from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Worker

from e2e.helpers import (
    assert_continue,
    assert_repeat,
    assert_stop,
    insert_sample_data,
    make_workflow_id,
    start_and_await_workflow,
)
from laborious.activities.activities import Activities
from laborious.workflows.predictions_batch import PredictionsBatch

base_input_data = {
    'schedule_name': 'test-schedule',
    'model_name': 'test_model',
    'model_id': 201,
    'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 201',
    'schema': 'predictions_schema',
    'table_name': 'predictions',
    'transform_table_name': 'transformed_data',
    'input_filters': {
        'SPECIFIC_VARIABLES_NULL_VALUES': {
            'POLICY': 'CONTINUE',
            'CONFIG': {'variables': ['sensor_1']},
        },
    },
    'mlflow_transform_filters': {
        'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
    },
    'mlflow_predict_filters': {
        'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
    },
    'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
    'opc_output_config': {},
    'pi_web_api_output_config': {},
    'save_transform': True,
    'prediction_store_policy': 'lts:1',
    'model_config': {
        'retention_minutes': 0,
        'transform_flavor': 'sklearn',
        'predict_flavor': 'sklearn',
    },
    'datetime_columns': ['timestamp', 'created_at'],
}

base_query = 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = {model_id}'


def get_base_input_data(model_id):
    return {
        **base_input_data,
        'model_id': model_id,
        'query': base_query.format(model_id=model_id),
    }


def insert_sample_prediction(postgres_engine, model_id):
    with postgres_engine.begin() as conn:
        conn.execute(text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}'))
        insert_sql = f"""
            INSERT INTO predictions_schema.predictions (model_id, timestamp, prediction, prediction_confidence, prediction_status, comments, response_time)
            VALUES
               ({model_id}, '2024-01-01 12:00:00+00:00', 10, 0, 'Good', '', 0.1)
        """
        conn.execute(text(insert_sql))
    return (model_id, Decimal(10), Decimal(0), 'Good')


@pytest.fixture
def bad_data_model(patch_mlflow):
    model = MagicMock(predict=MagicMock(side_effect=Exception('Bad data model')))
    patch_mlflow.sklearn.load_model = MagicMock(return_value=model)
    return model


@pytest.fixture
def bad_predict_model(patch_mlflow, mock_mlflow_models):
    model = MagicMock(predict=MagicMock(side_effect=Exception('Bad predict model')))

    def mock_sklearn_load_model(model_uri):
        if 'data_model' in model_uri or 'transform' in model_uri.lower():
            return mock_mlflow_models['transform_model']
        return model

    patch_mlflow.sklearn = MagicMock()
    patch_mlflow.sklearn.load_model = MagicMock(side_effect=mock_sklearn_load_model)
    return model


@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_2_1_1_input_gate_triggers_continue(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker: Worker,
    test_activities: Activities,
    postgres_engine,
    mock_mlflow_models,
):
    """Input gate CONTINUE: export default prediction; MLflow transform/predict not used."""
    client = temporal_test_env.client
    model_id = 211
    insert_sample_data(postgres_engine, model_id, ['NULL', 78.2])
    input_data = get_base_input_data(model_id)
    await start_and_await_workflow(
        client, PredictionsBatch.run, input_data, make_workflow_id('test-continue-policy')
    )
    assert_continue(postgres_engine, model_id)
    mock_mlflow_models['transform_model'].predict.assert_not_called()
    mock_mlflow_models['predict_model'].predict.assert_not_called()


@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_2_1_2_input_gate_triggers_stop(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker: Worker,
    test_activities: Activities,
    postgres_engine,
    mock_mlflow_models,
):
    """Input gate STOP: no export, no MLflow."""
    client = temporal_test_env.client
    model_id = 212
    insert_sample_data(postgres_engine, model_id, ['NULL', 78.2])
    input_data = get_base_input_data(model_id)
    input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['POLICY'] = 'STOP'
    await start_and_await_workflow(
        client, PredictionsBatch.run, input_data, make_workflow_id('test-input-stop')
    )
    assert_stop(postgres_engine, model_id)
    mock_mlflow_models['transform_model'].predict.assert_not_called()


@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_2_1_3_input_gate_triggers_repeat(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker: Worker,
    test_activities: Activities,
    postgres_engine,
    mock_mlflow_models,
):
    """Input gate REPEAT with existing history."""
    client = temporal_test_env.client
    model_id = 213
    insert_sample_data(postgres_engine, model_id, ['NULL', 78.2])
    data = insert_sample_prediction(postgres_engine, model_id)
    input_data = get_base_input_data(model_id)
    input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['POLICY'] = 'REPEAT'
    await start_and_await_workflow(
        client, PredictionsBatch.run, input_data, make_workflow_id('test-input-repeat')
    )
    assert_repeat(postgres_engine, model_id, data)
    mock_mlflow_models['transform_model'].predict.assert_not_called()


@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_2_1_4_input_gate_repeat_without_prior_prediction(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker: Worker,
    test_activities: Activities,
    postgres_engine,
):
    """REPEAT when no prior row in predictions: repeat_last_prediction runs; still no new duplicate export path."""
    client = temporal_test_env.client
    model_id = 214
    insert_sample_data(postgres_engine, model_id, ['NULL', 78.2])
    with postgres_engine.begin() as conn:
        conn.execute(text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}'))
    input_data = get_base_input_data(model_id)
    input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['POLICY'] = 'REPEAT'
    await start_and_await_workflow(
        client, PredictionsBatch.run, input_data, make_workflow_id('test-input-repeat-no-history')
    )
    assert_stop(postgres_engine, model_id)


@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_2_2_1_transform_gate_triggers_continue(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker: Worker,
    test_activities: Activities,
    postgres_engine,
    bad_data_model,
):
    client = temporal_test_env.client
    model_id = 221
    insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
    input_data = get_base_input_data(model_id)
    input_data['mlflow_transform_filters']['API_ERROR']['POLICY'] = 'CONTINUE'
    await start_and_await_workflow(
        client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-continue')
    )
    assert_continue(
        postgres_engine=postgres_engine,
        model_id=model_id,
        prediction_confidence=Decimal(10),
        comments='Unknown MLFlow API error',
    )


@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_2_2_2_transform_gate_triggers_stop(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker: Worker,
    test_activities: Activities,
    postgres_engine,
    bad_data_model,
    mock_mlflow_models,
):
    client = temporal_test_env.client
    model_id = 222
    insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
    input_data = get_base_input_data(model_id)
    input_data['mlflow_transform_filters']['API_ERROR']['POLICY'] = 'STOP'
    await start_and_await_workflow(
        client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-stop')
    )
    assert_stop(postgres_engine, model_id)
    mock_mlflow_models['predict_model'].predict.assert_not_called()


@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_2_2_3_transform_gate_triggers_repeat(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker: Worker,
    test_activities: Activities,
    postgres_engine,
    bad_data_model,
):
    client = temporal_test_env.client
    model_id = 223
    insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
    data = insert_sample_prediction(postgres_engine, model_id)
    input_data = get_base_input_data(model_id)
    input_data['mlflow_transform_filters']['API_ERROR']['POLICY'] = 'REPEAT'
    await start_and_await_workflow(
        client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-repeat')
    )
    assert_repeat(postgres_engine, model_id, data)


@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_2_2_4_transform_content_gate_nan_values_stop(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker: Worker,
    test_activities: Activities,
    postgres_engine,
    mock_mlflow_models,
):
    """mlflow_content_gate triggers STOP when transform output is all NaN (NAN_VALUES filter)."""
    client = temporal_test_env.client
    model_id = 224

    def all_nan_transform(data):
        num_rows = max(len(data), 1) if hasattr(data, '__len__') else 1
        result = pd.DataFrame({'feature_1': [np.nan] * num_rows, 'feature_2': [np.nan] * num_rows})
        result.index = data.index
        return result

    mock_mlflow_models['transform_model'].predict = MagicMock(side_effect=all_nan_transform)

    insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
    input_data = get_base_input_data(model_id)
    input_data['mlflow_transform_filters'] = {
        'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
        'NAN_VALUES': {'POLICY': 'STOP', 'CONFIG': {}},
    }
    await start_and_await_workflow(
        client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-content-stop')
    )
    assert_stop(postgres_engine, model_id)
    mock_mlflow_models['predict_model'].predict.assert_not_called()


@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_2_3_1_predict_gate_triggers_continue(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker: Worker,
    test_activities: Activities,
    postgres_engine,
    bad_predict_model,
):
    client = temporal_test_env.client
    model_id = 231
    insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
    input_data = get_base_input_data(model_id)
    input_data['mlflow_predict_filters']['API_ERROR']['POLICY'] = 'CONTINUE'
    await start_and_await_workflow(
        client, PredictionsBatch.run, input_data, make_workflow_id('test-predict-continue')
    )
    assert_continue(
        postgres_engine=postgres_engine,
        model_id=model_id,
        prediction_confidence=Decimal(10),
        comments='Unknown MLFlow API error',
    )


@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_2_3_2_predict_gate_triggers_stop(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker: Worker,
    test_activities: Activities,
    postgres_engine,
    bad_predict_model,
):
    client = temporal_test_env.client
    model_id = 232
    insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
    input_data = get_base_input_data(model_id)
    input_data['mlflow_predict_filters']['API_ERROR']['POLICY'] = 'STOP'
    await start_and_await_workflow(
        client, PredictionsBatch.run, input_data, make_workflow_id('test-predict-stop')
    )
    assert_stop(postgres_engine, model_id)


@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_2_3_3_predict_gate_triggers_repeat(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker: Worker,
    test_activities: Activities,
    postgres_engine,
    bad_predict_model,
):
    client = temporal_test_env.client
    model_id = 233
    insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
    data = insert_sample_prediction(postgres_engine, model_id)
    input_data = get_base_input_data(model_id)
    input_data['mlflow_predict_filters']['API_ERROR']['POLICY'] = 'REPEAT'
    input_data['path_priority'] = ['REPEAT', 'STOP', 'CONTINUE']
    await start_and_await_workflow(
        client, PredictionsBatch.run, input_data, make_workflow_id('test-predict-repeat')
    )
    assert_repeat(postgres_engine, model_id, data)


@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_2_4_1_input_empty_data_stop(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker: Worker,
    test_activities: Activities,
    postgres_engine,
):
    """EMPTY_DATA filter with STOP when query returns no rows (offload payload empty)."""
    client = temporal_test_env.client
    model_id = 241
    with postgres_engine.begin() as conn:
        conn.execute(text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}'))
    input_data = get_base_input_data(model_id)
    input_data['input_filters'] = {'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}}
    await start_and_await_workflow(
        client, PredictionsBatch.run, input_data, make_workflow_id('test-empty-data-stop')
    )
    assert_stop(postgres_engine, model_id)