"""
Direct E2E execution of child workflows (smaller surface than PredictionsBatch).
"""

from decimal import Decimal

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

from e2e.helpers import make_workflow_id, start_and_await_workflow
from laborious.workflows.sub_workflows.format_and_export_prediction import FormatAndExportPrediction


@pytest.mark.asyncio
@pytest.mark.integration
async def test_format_and_export_prediction_default_path_e2e(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker: Worker,
    postgres_engine,
):
    """
    Run FormatAndExportPrediction with path_flag set (format_default_prediction path).
    """
    client = temporal_test_env.client
    model_id = 401
    with postgres_engine.begin() as conn:
        conn.execute(text(f'DELETE FROM sientia_data.predictions WHERE model_id = {model_id}'))

    metadata = {
        'metadata': {
            'model_id': model_id,
            'model_name': 'test_model',
            'schedule_name': 'test-schedule',
            'workflow_name': 'subworkflow.format_and_export_prediction',
        }
    }
    input_data = {
        'metadata': metadata,
        'path_flag': 'CONTINUE',
        'data': {'last_timestamp': '2024-01-01 12:00:00+00:00'},
        'prediction_confidence': 2,
        'timestamp': '2024-01-01 12:00:00+00:00',
        'model_id': model_id,
        'model_name': 'test_model',
        'schema': 'sientia_data',
        'table_name': 'predictions',
        'transform_table_name': 'transformed_data',
        'comment': 'e2e child workflow default path',
        'opc_output_config': {},
        'pi_web_api_output_config': {},
        'prediction_store_policy': 'lts:1',
    }

    await start_and_await_workflow(
        client,
        FormatAndExportPrediction.run,
        input_data,
        make_workflow_id('e2e-format-export-child'),
    )

    with postgres_engine.connect() as conn:
        row = conn.execute(
            text(
                f'SELECT prediction, prediction_confidence, prediction_status, comments '
                f'FROM sientia_data.predictions WHERE model_id = {model_id}'
            )
        ).fetchone()
    assert row is not None
    assert row[0] == 0
    assert row[1] == Decimal(2)
    assert row[2] == 'Bad'
    assert row[3] == 'e2e child workflow default path'