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

import asyncio
from datetime import datetime
from decimal import Decimal
from typing import Any
from unittest.mock import MagicMock, patch

import pandas as pd
import pytest
from pytz import timezone
from sqlalchemy import text
from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Worker

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',  # Continue despite issues, not STOP
                '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_data(postgres_engine, model_id, values: list[Any]):
    with postgres_engine.begin() as conn:
        conn.execute(text(f"DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}"))
        
        # Insert data with some null values (quality issue)

        values_sql = []
        for i, value in enumerate(values):
            values_sql.append(f"""
                ({model_id}, 'sensor_{i+1}', {value}, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
            """)

        insert_sql = f"""
            INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
            VALUES
                {', '.join(values_sql)}
        """
        conn.execute(text(insert_sql))


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 data with some null values (quality issue)

        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')

async def start_and_await_workflow(client, input_data, workflow_id):
    handle = await client.start_workflow(
        PredictionsBatch.run,
        input_data,
        id=workflow_id,
        task_queue='test-queue',
    )
    print("[TEST] ✓ Workflow started")

    print("\n[TEST] 3. Waiting for workflow completion...")
    try:
        await asyncio.wait_for(handle.result(), timeout=60.0)
        print("[TEST] ✓ Workflow completed successfully")
    except asyncio.TimeoutError:
        pytest.fail("Workflow execution timed out after 60 seconds")

def assert_continue(
    postgres_engine, model_id, prediction_confidence: Decimal = Decimal(2),
    comments: str = 'Input data with bad quality',
):
    print("\n[TEST] 4. Verifying prediction was created despite warnings...")
    with postgres_engine.connect() as conn:
        result_query = conn.execute(
            text(f"SELECT model_id, prediction, prediction_confidence, prediction_status, comments FROM predictions_schema.predictions WHERE model_id = {model_id}")
        )
        prediction_rows = result_query.fetchall()
        assert len(prediction_rows) == 1, "Expected one prediction record despite warnings"

        # Assert prediction value is 0 and other fields
        row = prediction_rows[0]
        assert row[1] == 0, f"Expected prediction=0, got {row[1]}"
        assert row[2] == prediction_confidence, f"Expected prediction_confidence={prediction_confidence}, got {row[2]}"
        assert row[3] == 'Bad', f"Expected prediction_status='Bad', got {row[3]}"
        assert row[4] == comments, f"Expected comments='{comments}', got {row[4]}"


def assert_stop(postgres_engine, model_id):
    print("\n[TEST] 4. Verifying no predictions were created...")
    with postgres_engine.connect() as conn:
        result_query = conn.execute(
            text(f"SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = {model_id}")
        )
        count = result_query.scalar()
        assert count == 0, f"Expected no predictions, but found {count} records"

def assert_repeat(postgres_engine, model_id, last_prediction: list):
    print("\n[TEST] 4. Verifying prediction was repeated...")
    with postgres_engine.connect() as conn:
        result_query = conn.execute(
            text(f'SELECT model_id, prediction, prediction_confidence, prediction_status FROM predictions_schema.predictions WHERE model_id = {model_id}')
        )
        prediction_rows = result_query.fetchall()
        print(prediction_rows)
        assert len(prediction_rows) == 2, "Expected two prediction records"
        assert prediction_rows[0] == last_prediction, f"Expected first prediction to be the same as the last prediction, got {prediction_rows[0]}, expected {last_prediction}"
        assert prediction_rows[1] == last_prediction, f"Expected second prediction to be the same as the last prediction, got {prediction_rows[1]}, expected {last_prediction}"


@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,
):
    """
    Scenario 2.1.1: Input Gate Triggers CONTINUE
    
    Description:
    Input gate determines data should use previous prediction.
    
    Expected Behavior:
    - input_gate returns path_flag='CONTINUE'
    - path_flag_handler calls export workflow with input data directly
    - MLFlow transform and predict skipped
    - Data exported as-is
    
    Assertions:
    - input_gate called
    - MLFlow operations NOT called
    - Export workflow called with original data
    - Workflow completes
    """
    client = temporal_test_env.client

    model_id = 211

    print("\n[TEST] 1. Inserting test data...")

    insert_sample_data(postgres_engine, model_id, ['NULL', 78.2])

    print("[TEST] ✓ Data inserted successfully")

    input_data = get_base_input_data(model_id)


    print("\n[TEST] 2. Starting workflow with CONTINUE policy...")
    workflow_id = f'test-continue-policy-{datetime.now().timestamp()}'
    
    await start_and_await_workflow(client, input_data, workflow_id)
    assert_continue(postgres_engine, model_id)
    
    print("\n[TEST] ✓ All assertions passed!")


@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,
):
    """
    Scenario 2.1.2: Input Gate Triggers STOP
    
    Description:
    Input data quality gate fails with STOP policy.
    
    Expected Behavior:
    - input_gate returns path_flag='STOP'
    - path_flag_handler detects STOP
    - Workflow returns early without calling MLFlow
    - No prediction exported
    
    Assertions:
    - input_gate called
    - path_flag_handler returns True (early exit)
    - MLFlow transform NOT called
    - Export workflow NOT called
    - Workflow completes without error
    """
    client = temporal_test_env.client

    model_id = 212

    print("\n[TEST] 1. Inserting test data...")
    insert_sample_data(postgres_engine, model_id, ['NULL', 78.2])
    print("[TEST] ✓ Data inserted successfully")

    input_data = get_base_input_data(model_id)
    input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['POLICY'] = 'STOP'

    print("\n[TEST] 2. Starting workflow that should stop at input gate...")
    workflow_id = f'test-input-stop-{datetime.now().timestamp()}'
    
    await start_and_await_workflow(client, input_data, workflow_id)

    assert_stop(postgres_engine, model_id)
    
    print("\n[TEST] ✓ All assertions passed!")


@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,
):
    """
    Scenario 2.1.3: Input Gate Triggers REPEAT
    
    Description:
    Input gate determines data should repeat last prediction.
    
    Expected Behavior:
    - input_gate returns path_flag='REPEAT'
    - path_flag_handler calls repeat_last_prediction activity
    - MLFlow transform and predict skipped
    - Last prediction repeated and exported
    
    Assertions:
    - input_gate called
    - MLFlow operations NOT called
    - repeat_last_prediction activity called
    - Workflow completes
    """
    client = temporal_test_env.client

    model_id = 213

    print("\n[TEST] 1. Inserting test data and previous prediction...")
    insert_sample_data(postgres_engine, model_id, ['NULL', 78.2])
    data = insert_sample_prediction(postgres_engine, model_id)

    print("[TEST] ✓ Data and previous prediction inserted")

    input_data = get_base_input_data(model_id)
    input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['POLICY'] = 'REPEAT'

    print("\n[TEST] 2. Starting workflow that should trigger REPEAT...")
    workflow_id = f'test-input-repeat-{datetime.now().timestamp()}'
    
    await start_and_await_workflow(client, input_data, workflow_id)
    assert_repeat(postgres_engine, model_id, data)

    print("\n[TEST] ✓ All assertions passed!")

@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.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,
):
    """
    Scenario 2.2.1: Transform Gate Triggers CONTINUE
    
    Description:
    Transform response gate determines data should continue despite issues.
    
    Expected Behavior:
    - request_transform succeeds
    - mlflow_response_gate for transform returns path_flag='CONTINUE'
    - path_flag_handler calls export workflow with transform data
    - MLFlow predict skipped
    - Transform data exported as-is
    
    Assertions:
    - Transform completed
    - mlflow_response_gate called for transform
    - MLFlow predict NOT called
    - Export workflow called with transform data
    - Workflow completes
    """
    client = temporal_test_env.client

    model_id = 221

    print("\n[TEST] 1. Inserting test data...")
    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'

    print("\n[TEST] 2. Starting workflow that should trigger CONTINUE at transform gate...")
    workflow_id = f'test-transform-continue-{datetime.now().timestamp()}'
    
    await start_and_await_workflow(client, input_data, workflow_id)
    assert_continue(
        postgres_engine=postgres_engine,
        model_id=model_id,
        prediction_confidence=Decimal(10),
        comments='Unknown MLFlow API error',
    )
    
    print("\n[TEST] ✓ All assertions passed!")


@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,
):
    """
    Scenario 2.2.2: Transform Gate Triggers STOP
    
    Description:
    Transform response validation fails with STOP policy.
    
    Expected Behavior:
    - request_transform succeeds but response invalid
    - mlflow_response_gate for transform returns path_flag='STOP'
    - Workflow exits without calling predict or export
    
    Assertions:
    - Transform completed but validation failed
    - mlflow_response_gate called for transform
    - MLFlow predict NOT called
    - Export workflow NOT called
    - Workflow completes without error
    """
    client = temporal_test_env.client

    model_id = 222

    print("\n[TEST] 1. Inserting test data...")
    insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
    print("[TEST] ✓ Data inserted successfully")

    input_data = get_base_input_data(model_id)
    input_data['mlflow_transform_filters']['API_ERROR']['POLICY'] = 'STOP'

    print("\n[TEST] 2. Starting workflow that should trigger STOP at transform gate...")
    workflow_id = f'test-transform-stop-{datetime.now().timestamp()}'
    
    await start_and_await_workflow(client, input_data, workflow_id)
    assert_stop(postgres_engine, model_id)

    print("\n[TEST] ✓ All assertions passed!")



@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,
):
    """
    Scenario 2.2.3: Transform Gate Triggers REPEAT
    
    Description:
    Transform response gate determines data should repeat last prediction.
    
    Expected Behavior:
    - request_transform succeeds but response has issues
    - mlflow_response_gate for transform returns path_flag='REPEAT'
    - path_flag_handler calls repeat_last_prediction activity
    - MLFlow predict skipped
    - Last prediction repeated and exported
    
    Assertions:
    - Transform completed but validation triggered REPEAT
    - mlflow_response_gate called for transform
    - MLFlow predict NOT called
    - repeat_last_prediction activity called
    - Workflow completes
    """
    client = temporal_test_env.client

    model_id = 223

    print("\n[TEST] 1. Inserting test data...")
    insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
    data = insert_sample_prediction(postgres_engine, model_id)

    print("[TEST] ✓ Data inserted successfully")

    input_data = get_base_input_data(model_id)
    input_data['mlflow_transform_filters']['API_ERROR']['POLICY'] = 'REPEAT'

    print("\n[TEST] 2. Starting workflow that should trigger REPEAT at transform gate...")
    workflow_id = f'test-transform-repeat-{datetime.now().timestamp()}'

    await start_and_await_workflow(client, input_data, workflow_id)
    assert_repeat(postgres_engine, model_id, data)

    print("\n[TEST] ✓ All assertions passed!")


@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_3_1_predict_gate_triggers_continue(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker: Worker,
    test_activities: Activities,
    postgres_engine,
    bad_predict_model,
):
    """
    Scenario 2.3.1: Predict Gate Triggers CONTINUE
    
    Description:
    Predict response gate determines data should continue despite issues.
    
    Expected Behavior:
    - request_predict succeeds
    - mlflow_response_gate for predict returns path_flag='CONTINUE'
    - path_flag_handler calls export workflow with predict data
    - Prediction exported despite quality issues
    
    Assertions:
    - Transform and predict completed
    - mlflow_response_gate called for predict
    - Export workflow called with predict data
    - Workflow completes
    """
    client = temporal_test_env.client

    model_id = 231

    print("\n[TEST] 1. Inserting test data...")
    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'

    print("\n[TEST] 2. Starting workflow that should trigger CONTINUE at predict gate...")
    workflow_id = f'test-predict-continue-{datetime.now().timestamp()}'
    
    await start_and_await_workflow(client, input_data, workflow_id)
    assert_continue(
        postgres_engine=postgres_engine,
        model_id=model_id,
        prediction_confidence=Decimal(10),
        comments='Unknown MLFlow API error',
    )
    
    print("\n[TEST] ✓ All assertions passed!")


@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,
):
    """
    Scenario 2.3.2: Predict Gate Triggers STOP
    
    Description:
    Prediction validation fails with STOP policy.
    
    Expected Behavior:
    - request_predict succeeds but response invalid
    - mlflow_response_gate for predict returns path_flag='STOP'
    - Workflow exits without export
    
    Assertions:
    - Transform completed
    - Predict completed but validation failed
    - Export workflow NOT called
    - Workflow completes without error
    """
    client = temporal_test_env.client

    model_id = 232

    print("\n[TEST] 1. Inserting test data...")
    insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
    print("[TEST] ✓ Data inserted successfully")

    input_data = get_base_input_data(model_id)
    input_data['mlflow_predict_filters']['API_ERROR']['POLICY'] = 'STOP'

    print("\n[TEST] 2. Starting workflow that should stop at predict gate...")
    workflow_id = f'test-predict-stop-{datetime.now().timestamp()}'
    
    await start_and_await_workflow(client, input_data, workflow_id)
    assert_stop(postgres_engine, model_id)

    print("\n[TEST] ✓ All assertions passed!")


@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,
):
    """
    Scenario 2.3.3: Predict Gate Triggers REPEAT
    
    Description:
    Predict response gate determines data should repeat last prediction.
    
    Expected Behavior:
    - request_predict succeeds but response has issues
    - mlflow_response_gate for predict returns path_flag='REPEAT'
    - path_flag_handler calls repeat_last_prediction activity
    - Last prediction repeated and exported
    
    Assertions:
    - Transform and predict completed but validation triggered REPEAT
    - mlflow_response_gate called for predict
    - repeat_last_prediction activity called
    - Export workflow NOT called with current prediction
    - Workflow completes
    """
    client = temporal_test_env.client

    model_id = 233

    print("\n[TEST] 1. Inserting test data and previous prediction...")
    insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
    data = insert_sample_prediction(postgres_engine, model_id)

    print("[TEST] ✓ Data and previous prediction inserted")

    input_data = get_base_input_data(model_id)
    input_data['mlflow_predict_filters']['API_ERROR']['POLICY'] = 'REPEAT'
    input_data['path_priority'] = ['REPEAT', 'STOP', 'CONTINUE']  # REPEAT first

    print("\n[TEST] 2. Starting workflow that should trigger REPEAT at predict gate...")
    workflow_id = f'test-predict-repeat-{datetime.now().timestamp()}'

    await start_and_await_workflow(client, input_data, workflow_id)
    assert_repeat(postgres_engine, model_id, data)

    print("\n[TEST] ✓ All assertions passed!")