"""
End-to-end tests for the SimpleMetrics workflow.

Coverage focus:

- Happy path computes rmse/mse/mae/r2 from predictions joined against ``laborious_data``
  and persists rows to ``predictions_schema.simple_metrics_data`` with all required columns.
- Subset metric selection (only rmse) writes exactly the requested rows.
- Zero-variance target produces ``r2=0`` per division-by-zero guard.
- Empty join (no overlapping data) short-circuits without persisting anything.
"""

from datetime import datetime, timedelta, timezone
from decimal import Decimal

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

from e2e.helpers import load_scenario_input, make_workflow_id, start_and_await_workflow
from laborious.activities.activities import Activities
from laborious.workflows.simple_metrics import SimpleMetrics

EXPECTED_SIMPLE_METRICS_COLUMNS = [
    'id',
    'model_id',
    'metric',
    'value',
    'timestamp',
    'data_size',
    'interval_minutes',
    'created_at',
]


def _simple_metrics_input(model_id: int, **overrides) -> dict:
    """Load and override the simple-metrics base scenario."""
    payload = load_scenario_input('simple_metrics_base.json', model_id=model_id)
    payload.update(overrides)
    return payload


def _seed_predictions_and_targets(
    postgres_engine,
    model_id: int,
    pairs: list[tuple[float, float]],
    target_name: str = 'sensor_target',
    offset_minutes: int = 6,
) -> list[str]:
    """
    Insert matching prediction/target rows used by the SimpleMetrics SQL JOIN.

    For each ``(prediction, target)`` pair we write a row in ``predictions`` and
    a matching row in ``laborious_data`` with ``variable=target_name`` so the
    inner join in the workflow query yields one row per pair.

    Args:
     - postgres_engine: SQLAlchemy engine bound to the test container.
     - model_id: Model id stamped on every row.
     - pairs: ``(prediction, target)`` pairs, one per minute.
     - target_name: Variable name in ``laborious_data`` representing the target.
     - offset_minutes: Earliest row sits this many minutes ago so timestamps fall
       inside the workflow's recent-data window.

    Return:
     List of timestamp strings written for the inserted rows.
    """
    base = datetime.now(timezone.utc).replace(second=0, microsecond=0) - timedelta(
        minutes=offset_minutes
    )
    timestamps = [
        (base + timedelta(minutes=i)).strftime('%Y-%m-%d %H:%M:%S%z')
        for i in range(len(pairs))
    ]

    prediction_rows = []
    target_rows = []
    for index, (prediction, target_value) in enumerate(pairs):
        ts = timestamps[index]
        prediction_rows.append(
            f"({model_id}, {prediction}, 0, 0, 'Good', '{ts}', '{ts}')"
        )
        target_rows.append(
            f"({model_id}, '{target_name}', {target_value}, '{ts}', '{ts}')"
        )

    with postgres_engine.begin() as conn:
        conn.execute(text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}'))
        conn.execute(text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}'))
        # The SimpleMetrics SQL JOIN only filters ``predictions.model_id``; it does
        # NOT filter ``laborious_data.model_id`` (see ``e2e/CODE_ISSUES.md`` issue
        # SM-1). Without this cross-model cleanup, a previous test's target rows
        # under the same variable name would join into this test's predictions
        # whenever timestamps happened to overlap.
        conn.execute(
            text(
                "DELETE FROM predictions_schema.laborious_data "
                "WHERE variable IN (:sensor_default, :target_name) "
                "AND timestamp >= NOW() - INTERVAL '120 minutes'"
            ),
            {'sensor_default': 'sensor_target', 'target_name': target_name},
        )
        if prediction_rows:
            conn.execute(
                text(
                    'INSERT INTO predictions_schema.predictions '
                    '(model_id, prediction, prediction_confidence, response_time, '
                    'prediction_status, "timestamp", created_at) VALUES '
                    + ', '.join(prediction_rows)
                )
            )
            conn.execute(
                text(
                    'INSERT INTO predictions_schema.laborious_data '
                    '(model_id, variable, value, "timestamp", created_at) VALUES '
                    + ', '.join(target_rows)
                )
            )
    return timestamps


@pytest.mark.asyncio
@pytest.mark.integration
async def test_simple_metrics_happy_path_persists_all_metrics_and_columns(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker_simple_metrics: Worker,
    test_activities: Activities,
    postgres_engine,
):
    """
    Scenario S.1.1: rmse/mse/mae/r2 are calculated from a deterministic
    prediction/target pair set and written one row per metric. Every column
    expected by ``predictions_schema.simple_metrics_data`` must be populated and
    the numerical values must match closed-form expectations.
    """
    client = temporal_test_env.client
    model_id = 511

    pairs = [
        (1.0, 2.0),
        (2.0, 4.0),
        (3.0, 5.0),
        (4.0, 9.0),
        (5.0, 12.0),
    ]
    diffs = [target - prediction for prediction, target in pairs]
    n = len(diffs)
    expected_rmse = math.sqrt(sum(d * d for d in diffs) / n)
    expected_mse = sum(d * d for d in diffs) / n
    expected_mae = sum(abs(d) for d in diffs) / n
    target_mean = sum(t for _, t in pairs) / n
    ss_res = sum((target - prediction) ** 2 for prediction, target in pairs)
    ss_tot = sum((t - target_mean) ** 2 for _, t in pairs)
    expected_r2 = 1.0 - (ss_res / ss_tot)

    _seed_predictions_and_targets(postgres_engine, model_id=model_id, pairs=pairs)

    input_data = _simple_metrics_input(model_id)
    await start_and_await_workflow(
        client,
        SimpleMetrics.run,
        input_data,
        make_workflow_id('test-simple-metrics-happy'),
    )

    with postgres_engine.connect() as conn:
        rows = (
            conn.execute(
                text(
                    'SELECT * FROM predictions_schema.simple_metrics_data '
                    'WHERE model_id = :m ORDER BY metric'
                ),
                {'m': model_id},
            )
            .mappings()
            .all()
        )

    assert len(rows) == 4, f'Expected 4 metric rows, got {len(rows)}'
    for column in EXPECTED_SIMPLE_METRICS_COLUMNS:
        assert column in rows[0], f'Missing simple_metrics column: {column}'
    for row in rows:
        for column in EXPECTED_SIMPLE_METRICS_COLUMNS:
            assert row[column] is not None, f"Column '{column}' is NULL in {dict(row)}"

    by_metric = {row['metric']: row for row in rows}
    assert set(by_metric) == {'rmse', 'mse', 'mae', 'r2'}

    def _decimal_close(actual, expected, places: int = 6) -> bool:
        return abs(float(actual) - expected) < 10 ** (-places)

    assert _decimal_close(by_metric['rmse']['value'], expected_rmse)
    assert _decimal_close(by_metric['mse']['value'], expected_mse)
    assert _decimal_close(by_metric['mae']['value'], expected_mae)
    assert _decimal_close(by_metric['r2']['value'], expected_r2)

    assert all(row['data_size'] == n for row in rows), 'data_size must equal target row count'
    assert all(row['interval_minutes'] == 60 for row in rows)
    assert all(row['model_id'] == model_id for row in rows)


@pytest.mark.asyncio
@pytest.mark.integration
async def test_simple_metrics_subset_metrics_writes_only_requested_rows(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker_simple_metrics: Worker,
    test_activities: Activities,
    postgres_engine,
):
    """
    Scenario S.1.2: Requesting ``metrics=['rmse']`` must persist exactly one row
    with metric ``rmse`` and skip mse/mae/r2.
    """
    client = temporal_test_env.client
    model_id = 512

    pairs = [(1.0, 2.0), (2.0, 4.0), (3.0, 6.0)]
    _seed_predictions_and_targets(postgres_engine, model_id=model_id, pairs=pairs)

    input_data = _simple_metrics_input(model_id, metrics=['rmse'])
    await start_and_await_workflow(
        client,
        SimpleMetrics.run,
        input_data,
        make_workflow_id('test-simple-metrics-subset'),
    )

    with postgres_engine.connect() as conn:
        metrics = [
            r[0]
            for r in conn.execute(
                text(
                    'SELECT metric FROM predictions_schema.simple_metrics_data '
                    'WHERE model_id = :m'
                ),
                {'m': model_id},
            ).all()
        ]
    assert metrics == ['rmse']


@pytest.mark.asyncio
@pytest.mark.integration
async def test_simple_metrics_zero_variance_target_returns_zero_r2(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker_simple_metrics: Worker,
    test_activities: Activities,
    postgres_engine,
):
    """
    Scenario S.2.1: When the target column has zero variance the activity must
    return ``r2 = 0`` (division-by-zero guard) and still persist all four metrics.
    """
    client = temporal_test_env.client
    model_id = 521

    pairs = [(0.0, 5.0), (1.0, 5.0), (2.0, 5.0), (3.0, 5.0)]
    _seed_predictions_and_targets(postgres_engine, model_id=model_id, pairs=pairs)

    input_data = _simple_metrics_input(model_id)
    await start_and_await_workflow(
        client,
        SimpleMetrics.run,
        input_data,
        make_workflow_id('test-simple-metrics-zero-variance'),
    )

    with postgres_engine.connect() as conn:
        r2_value = conn.execute(
            text(
                "SELECT value FROM predictions_schema.simple_metrics_data "
                "WHERE model_id = :m AND metric = 'r2'"
            ),
            {'m': model_id},
        ).scalar()
    assert r2_value is not None
    assert Decimal(str(r2_value)) == Decimal('0'), f'expected r2=0, got {r2_value!r}'


@pytest.mark.asyncio
@pytest.mark.integration
async def test_simple_metrics_no_overlapping_data_short_circuits(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker_simple_metrics: Worker,
    test_activities: Activities,
    postgres_engine,
):
    """
    Scenario S.3.1: When the join produces no rows (no matching laborious_data
    row for the configured ``target``), the workflow returns early without
    invoking ``calculate_simple_metrics`` and writes nothing.
    """
    client = temporal_test_env.client
    model_id = 531

    # Insert predictions but no matching target rows for the configured variable.
    _seed_predictions_and_targets(
        postgres_engine,
        model_id=model_id,
        pairs=[(1.0, 1.0)],
        target_name='wrong_variable_name',
    )

    input_data = _simple_metrics_input(model_id)
    await start_and_await_workflow(
        client,
        SimpleMetrics.run,
        input_data,
        make_workflow_id('test-simple-metrics-empty-join'),
    )

    with postgres_engine.connect() as conn:
        count = conn.execute(
            text(
                'SELECT COUNT(*) FROM predictions_schema.simple_metrics_data '
                'WHERE model_id = :m'
            ),
            {'m': model_id},
        ).scalar()
    assert count == 0, 'Empty target data must short-circuit and skip persistence'