"""
End-to-end tests for PredictionsBatch workflow - Main workflow scenarios.
"""

import asyncio

from sqlalchemy import text
from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Worker

import pytest

from e2e.helpers import make_workflow_id, start_and_await_workflow
from laborious.activities.activities import Activities
from laborious.workflows.predictions_batch import PredictionsBatch


@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_1_1_1_happy_path_complete_success(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker: Worker,
    test_activities: Activities,
    postgres_engine,
):
    """Scenario 1.1.1: Happy path with SQL load, MLflow mocks, Postgres predictions and transforms."""
    client = temporal_test_env.client

    with postgres_engine.begin() as conn:
        conn.execute(text('DELETE FROM predictions_schema.laborious_data WHERE model_id = 123'))
        insert_sql = """
            INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
            VALUES
                (123, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
                (123, 'sensor_2', 78.2, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
                (123, 'sensor_3', 120.8, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
        """
        conn.execute(text(insert_sql))

    input_data = {
        'metadata': {
            'metadata': {
                'model_id': 123,
                'model_name': 'test_model',
                'schedule_name': 'test-schedule',
                'workflow_name': 'predictions_batch',
            }
        },
        'schedule_name': 'test-schedule',
        'model_name': 'test_model',
        'model_id': 123,
        'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 123',
        'schema': 'predictions_schema',
        'table_name': 'predictions',
        'transform_table_name': 'transformed_data',
        'input_filters': {
            'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}},
        },
        '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': {
            'target': 'sensor_1',
            'retention_minutes': 0,
        },
        'datetime_columns': ['timestamp', 'created_at'],
    }

    await start_and_await_workflow(
        client,
        PredictionsBatch.run,
        input_data,
        make_workflow_id('test-predictions-batch'),
    )

    schema_name = 'predictions_schema'
    with postgres_engine.connect() as conn:
        result_query = conn.execute(
            text(
                f'SELECT model_id, prediction, prediction_confidence, response_time, prediction_status, comments '
                f'FROM {schema_name}.predictions WHERE model_id = 123'
            )
        )
        prediction_rows = result_query.fetchall()
        assert len(prediction_rows) == 1
        row = prediction_rows[0]
        assert row[0] == 123
        assert row[1] == 0.5
        assert row[2] == 0, f'Expected prediction_confidence=0, got {row[2]}'
        assert row[3] is not None
        assert row[4] == 'Good'
        assert row[5] == ''

        result_query = conn.execute(
            text(
                f'SELECT model_id, variable, value FROM {schema_name}.transformed_data WHERE model_id = 123'
            )
        )
        transformed_rows = result_query.fetchall()
        assert len(transformed_rows) == 2
        assert transformed_rows[0][0] == 123
        assert transformed_rows[0][1] == 'feature_1'
        assert float(transformed_rows[0][2]) == 0.234
        assert transformed_rows[1][0] == 123
        assert transformed_rows[1][1] == 'feature_2'
        assert float(transformed_rows[1][2]) == 0.783


@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_1_2_1_sql_query_execution_error(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker: Worker,
    test_activities: Activities,
    postgres_engine,
):
    """Invalid SQL: workflow may complete with early exit; no prediction rows."""
    client = temporal_test_env.client

    input_data = {
        'metadata': {
            'metadata': {
                'model_id': 128,
                'model_name': 'test_model',
                'schedule_name': 'test-schedule',
                'workflow_name': 'predictions_batch',
            }
        },
        'schedule_name': 'test-schedule',
        'model_name': 'test_model',
        'model_id': 128,
        'query': 'SELECT * FROM nonexistent_table WHERE invalid_syntax =',
        'schema': 'predictions_schema',
        'table_name': 'predictions',
        'transform_table_name': 'transformed_data',
        'input_filters': {
            'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}},
        },
        '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': {
            'target': 'sensor_1',
            'retention_minutes': 0,
        },
    }

    await start_and_await_workflow(
        client,
        PredictionsBatch.run,
        input_data,
        make_workflow_id('test-sql-error'),
    )

    with postgres_engine.connect() as conn:
        count = conn.execute(
            text('SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 128')
        ).scalar()
    assert count == 0


@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_1_2_2_missing_required_parameters(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker: Worker,
    test_activities: Activities,
    postgres_engine,
):
    """Missing query: workflow does not produce predictions and is terminated explicitly."""
    client = temporal_test_env.client

    input_data = {
        'metadata': {
            'metadata': {
                'model_id': 129,
                'model_name': 'test_model',
                'schedule_name': 'test-schedule',
                'workflow_name': 'predictions_batch',
            }
        },
        'schedule_name': 'test-schedule',
        'model_name': 'test_model',
        'model_id': 129,
        'schema': 'predictions_schema',
        'table_name': 'predictions',
        'transform_table_name': 'transformed_data',
    }

    handle = await client.start_workflow(
        PredictionsBatch.run,
        input_data,
        id=make_workflow_id('test-missing-param'),
        task_queue='test-queue',
    )

    # Let Temporal process a few workflow tasks; for this case, result() can hang.
    await asyncio.sleep(2.0)

    with postgres_engine.connect() as conn:
        count = conn.execute(
            text('SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 129')
        ).scalar()
    assert count == 0

    await handle.terminate('expected failure path in e2e test (missing required parameters)')


@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_1_2_3_invalid_datetime_column_specification(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker: Worker,
    test_activities: Activities,
    postgres_engine,
):
    """Invalid datetime column: no predictions persisted; workflow terminated after validation."""
    client = temporal_test_env.client

    with postgres_engine.begin() as conn:
        conn.execute(text('DELETE FROM predictions_schema.laborious_data WHERE model_id = 130'))
        conn.execute(
            text(
                """
            INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
            VALUES (130, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
        """
            )
        )

    input_data = {
        'metadata': {
            'metadata': {
                'model_id': 130,
                'model_name': 'test_model',
                'schedule_name': 'test-schedule',
                'workflow_name': 'predictions_batch',
            }
        },
        'schedule_name': 'test-schedule',
        'model_name': 'test_model',
        'model_id': 130,
        'query': 'SELECT timestamp, variable, value FROM predictions_schema.laborious_data WHERE model_id = 130',
        'schema': 'predictions_schema',
        'table_name': 'predictions',
        'transform_table_name': 'transformed_data',
        'input_filters': {
            'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}},
        },
        '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': {
            'target': 'sensor_1',
            'retention_minutes': 0,
        },
        'datetime_columns': ['nonexistent_column'],
    }

    handle = await client.start_workflow(
        PredictionsBatch.run,
        input_data,
        id=make_workflow_id('test-invalid-datetime-col'),
        task_queue='test-queue',
    )

    # Let Temporal process and surface the failure path internally.
    await asyncio.sleep(2.0)

    with postgres_engine.connect() as conn:
        count = conn.execute(
            text('SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 130')
        ).scalar()
    assert count == 0

    await handle.terminate('expected failure path in e2e test (invalid datetime column)')