"""
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 load_scenario_input, 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 = load_scenario_input('main_happy_path.json', model_id=123)

    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 = load_scenario_input('main_sql_error.json', model_id=128)

    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 = load_scenario_input('main_missing_required.json', model_id=129)

    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 = load_scenario_input('main_invalid_datetime.json', model_id=130)

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