"""
E2E-style tests for MinIO offload using a real MinIO testcontainer.
"""

from unittest.mock import patch

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

from e2e.helpers import (
    insert_sample_data,
    load_scenario_input,
    make_workflow_id,
    start_and_await_workflow,
)
from laborious.activities.activities import Activities
from laborious.utils.models import minio_dataframe_payload as mdp
from laborious.workflows.predictions_batch import PredictionsBatch


@pytest.mark.asyncio
@pytest.mark.integration
async def test_load_query_with_minio_offload_writes_object_to_bucket(
    postgres_engine,
    minio_container,
    test_activities_real_minio: Activities,
):
    """
    With a tiny offload threshold, query results are uploaded as Parquet to MinIO.

    Uses real MinioRepository against testcontainers MinIO (no MinIO mock).
    """
    model_id = 501
    with postgres_engine.begin() as conn:
        conn.execute(text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}'))
    insert_sample_data(postgres_engine, model_id, [1.0, 2.0])

    scenario_input = load_scenario_input('minio_offload_load_query.json', model_id=model_id)
    metadata = {'metadata': scenario_input['metadata']}
    with patch.object(mdp, 'OFFLOAD_THRESHOLD_BYTES', 1):
        payload = test_activities_real_minio.load_query_with_minio_offload(scenario_input)
    assert payload.object_key, 'offloaded payload must reference a MinIO object'
    assert payload.data is None or payload.data == {}, 'large payloads should not inline tabular dict'

    df = payload.retrieve(test_activities_real_minio.minio_repository, metadata['metadata'])
    assert len(df) >= 1

    client = minio_container.get_client()
    listed = list(client.list_objects('test-bucket', recursive=True))
    names = [getattr(o, 'object_name', None) or getattr(o, '_object_name', '') for o in listed]
    assert any(n and 'prediction_datasets' in n for n in names), f'unexpected object listing: {names!r}'


@pytest.mark.asyncio
@pytest.mark.integration
async def test_predictions_batch_with_minio_offload_path(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker_real_minio: Worker,
    postgres_engine,
    test_activities_real_minio: Activities,
):
    """
    Full PredictionsBatch run with offload: load step stores Parquet in MinIO; pipeline completes.
    """
    model_id = 502
    with postgres_engine.begin() as conn:
        conn.execute(text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}'))
        conn.execute(text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}'))
        conn.execute(text(f'DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}'))
    insert_sample_data(postgres_engine, model_id, [10.0, 20.0, 30.0])

    input_data = load_scenario_input('minio_offload_workflow.json', model_id=model_id)

    with patch.object(mdp, 'OFFLOAD_THRESHOLD_BYTES', 1):
        await start_and_await_workflow(
            temporal_test_env.client,
            PredictionsBatch.run,
            input_data,
            make_workflow_id('test-batch-minio-offload'),
        )

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


@pytest.mark.asyncio
@pytest.mark.integration
async def test_load_query_with_inline_payload_when_below_threshold(
    postgres_engine,
    test_activities_real_minio: Activities,
):
    """Scenario 4.2.1: payload stays inline when threshold is high enough."""
    model_id = 503
    with postgres_engine.begin() as conn:
        conn.execute(text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}'))
    insert_sample_data(postgres_engine, model_id, [1.0, 2.0])

    scenario_input = load_scenario_input('minio_offload_load_query.json', model_id=model_id)
    with patch.object(mdp, 'OFFLOAD_THRESHOLD_BYTES', 10**9):
        payload = test_activities_real_minio.load_query_with_minio_offload(scenario_input)

    assert payload.object_key is None
    assert payload.data is not None