sientia-dataops-scouter_temporal/e2e/helpers.py

"""
Shared helpers for Scouter end-to-end tests.
"""

from __future__ import annotations

import json
import re
import uuid
from datetime import UTC, datetime, timedelta
from pathlib import Path
from typing import Any
from urllib.parse import urlparse

from pymongo import MongoClient
from redis import Redis
from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ
from sqlalchemy import create_engine, text
from sqlalchemy.engine import Engine
from temporalio.client import Client

SCENARIO_INPUTS_DIR = Path(__file__).resolve().parent / 'scenario_inputs'
SCOUTER_TASK_QUEUE = 'scouter-test-queue'

_NOW_MARKER = re.compile(r'^@now(?:([+-])(\d+)([smhd]))?$')


def _resolve_timestamp_marker(value: str) -> datetime:
    """
    Resolve @now relative timestamp markers to timezone-aware datetimes.

    Args:
        - value: Marker string such as @now, @now-1h, or @now+30m

    Return:
        Resolved datetime in UTC
    """
    match = _NOW_MARKER.match(value.strip())
    if not match:
        raise ValueError(f'Invalid timestamp marker: {value}')

    now = datetime.now(UTC)
    if match.group(1) is None:
        return now

    sign, amount, unit = match.group(1), int(match.group(2)), match.group(3)
    delta_kwargs = {'seconds': 0, 'minutes': 0, 'hours': 0, 'days': 0}
    if unit == 's':
        delta_kwargs['seconds'] = amount
    elif unit == 'm':
        delta_kwargs['minutes'] = amount
    elif unit == 'h':
        delta_kwargs['hours'] = amount
    elif unit == 'd':
        delta_kwargs['days'] = amount

    delta = timedelta(**delta_kwargs)
    return now - delta if sign == '-' else now + delta


def _resolve_payload(node: Any) -> Any:
    """
    Recursively resolve @now markers inside JSON-loaded structures.

    Args:
        - node: JSON node (dict, list, or scalar)

    Return:
        Structure with markers replaced by datetimes or formatted strings
    """
    if isinstance(node, dict):
        return {key: _resolve_payload(value) for key, value in node.items()}
    if isinstance(node, list):
        return [_resolve_payload(item) for item in node]
    if isinstance(node, str) and node.startswith('@now'):
        resolved = _resolve_timestamp_marker(node)
        return resolved.strftime('%Y-%m-%d %H:%M:%S.%f%z')
    return node


def load_scenario_input(name: str, **overrides: Any) -> dict[str, Any]:
    """
    Load a scenario JSON file and apply optional overrides.

    Args:
        - name: Scenario slug with or without .json suffix
        - overrides: Top-level keys merged into the loaded document

    Return:
        Parsed scenario document with @now markers resolved
    """
    slug = name.removesuffix('.json')
    path = SCENARIO_INPUTS_DIR / f'{slug}.json'
    with path.open(encoding='utf-8') as handle:
        payload = json.load(handle)
    resolved = _resolve_payload(payload)
    if overrides:
        resolved.update(overrides)
    return resolved


def make_workflow_id(prefix: str) -> str:
    """
    Build a unique Temporal workflow id for E2E runs.

    Args:
        - prefix: Human-readable prefix for the workflow id

    Return:
        Unique workflow id string
    """
    return f'{prefix}-{uuid.uuid4().hex[:12]}'


async def start_and_await_workflow(
    client: Client,
    workflow_run: Any,
    input_data: dict[str, Any],
    workflow_id: str,
    *,
    task_queue: str = SCOUTER_TASK_QUEUE,
    timeout: float = 300.0,
) -> Any:
    """
    Start a workflow on the E2E task queue and await its result.

    Args:
        - client: Temporal client from WorkflowEnvironment
        - workflow_run: Workflow run method (e.g. Scouter.run)
        - input_data: Workflow input payload
        - workflow_id: Unique workflow id
        - task_queue: Task queue name
        - timeout: Maximum seconds to wait for completion

    Return:
        Workflow result (None for Scouter-family workflows)
    """
    return await client.execute_workflow(
        workflow_run,
        input_data,
        id=workflow_id,
        task_queue=task_queue,
        execution_timeout=timedelta(seconds=timeout),
    )


def _coerce_mongo_document(document: dict[str, Any]) -> dict[str, Any]:
    """
    Convert string timestamps in seed documents to BSON datetimes for Mongo filters.

    Args:
        - document: Raw document dict from scenario JSON

    Return:
        Document with inserted_at as datetime when present
    """
    doc = dict(document)
    inserted_at = doc.get('inserted_at')
    if isinstance(inserted_at, str):
        doc['inserted_at'] = datetime.strptime(inserted_at, DATETIME_FORMAT_MS_WITH_TZ).replace(
            tzinfo=UTC
        )
    return doc


def seed_raw_collection(
    mongo_uri: str,
    database: str,
    schedule_name: str,
    documents: list[dict[str, Any]],
) -> None:
    """
    Insert raw Mongo documents into raw_<schedule_name>.

    Args:
        - mongo_uri: MongoDB connection string
        - database: Database name
        - schedule_name: Schedule slug used in collection name
        - documents: Documents to insert
    """
    collection_name = f'raw_{schedule_name}'
    client = MongoClient(mongo_uri)
    try:
        collection = client[database][collection_name]
        if documents:
            collection.insert_many([_coerce_mongo_document(doc) for doc in documents])
    finally:
        client.close()


def seed_last_data_timestamp(
    redis_client: Redis,
    workflow_name: str,
    schedule_name: str,
    value: str,
) -> None:
    """
    Pre-seed last_data_timestamp Redis key the same way RedisRepository.set stores strings.

    Args:
        - redis_client: Connected Redis client
        - workflow_name: Workflow name segment in the key
        - schedule_name: Schedule name segment in the key
        - value: Timestamp string to store
    """
    key = f'last_data_timestamp:{workflow_name}:{schedule_name}'
    redis_client.set(key, json.dumps(value))


def count_laborious_rows(engine: Engine, model_id: str | int | None = None) -> int:
    """
    Count rows in sientia_data.laborious_data, optionally filtered by model_id.

    Args:
        - engine: SQLAlchemy engine bound to the Postgres testcontainer
        - model_id: Optional model id filter

    Return:
        Row count
    """
    query = 'SELECT COUNT(*) FROM sientia_data.laborious_data'
    params: dict[str, Any] = {}
    if model_id is not None:
        query += ' WHERE model_id = :model_id'
        params['model_id'] = int(model_id)

    with engine.connect() as conn:
        return conn.execute(text(query), params).scalar() or 0


def fetch_laborious_rows(engine: Engine, model_id: str | int | None = None) -> list[dict[str, Any]]:
    """
    Fetch laborious_data rows as plain dicts.

    Args:
        - engine: SQLAlchemy engine bound to the Postgres testcontainer
        - model_id: Optional model id filter

    Return:
        List of row dicts with variable and value keys
    """
    query = 'SELECT model_id, variable, value, timestamp FROM sientia_data.laborious_data'
    params: dict[str, Any] = {}
    if model_id is not None:
        query += ' WHERE model_id = :model_id'
        params['model_id'] = int(model_id)
    query += ' ORDER BY variable'

    with engine.connect() as conn:
        rows = conn.execute(text(query), params).mappings().all()
    return [dict(row) for row in rows]


def count_held_data_keys(redis_client: Redis) -> int:
    """
    Count Redis keys matching held_data_*.

    Args:
        - redis_client: Connected Redis client

    Return:
        Number of matching keys
    """
    return len(redis_client.keys('held_data_*'))


def count_notifications(
    mongo_uri: str,
    database: str,
    *,
    notification_id: str | None = None,
    level: str | None = None,
) -> int:
    """
    Count notification documents in the E2E notification collection.

    Args:
        - mongo_uri: MongoDB connection string
        - database: Database name
        - notification_id: Optional notification_id filter
        - level: Optional level filter (WARNING, ERROR, ...)

    Return:
        Matching document count
    """
    client = MongoClient(mongo_uri)
    try:
        collection = client[database]['notification_queue']
        query: dict[str, Any] = {}
        if notification_id:
            query['notification_id'] = notification_id
        if level:
            query['level'] = level
        return collection.count_documents(query)
    finally:
        client.close()


def default_model_tags(
    *,
    names: list[str],
    aggr: str = 'avg',
    data_range: list[float] | None = None,
    frequency: int = 60000,
) -> dict[str, dict[str, Any]]:
    """
    Build a minimal model_tags map for E2E scenarios.

    Args:
        - names: Tag names to include
        - aggr: Aggregation function (aggr_func field)
        - data_range: Optional [min, max] validation range
        - frequency: Collection frequency in milliseconds

    Return:
        model_tags dict keyed by tag name
    """
    if data_range is None:
        data_range = [0, 100]
    return {
        name: {
            'webid': f'webid_{name}',
            'aggr_func': aggr,
            'data_range': data_range,
            'frequency': frequency,
        }
        for name in names
    }


def default_scouter_input(
    *,
    model_id: str = '1',
    model_name: str = 'Test Model',
    schedule_name: str = 'test-schedule',
    workflow_name: str = 'scouter',
    **overrides: Any,
) -> dict[str, Any]:
    """
    Return a base workflow input dict for Scouter / CoreScouter E2E runs.

    Args:
        - model_id: Model identifier
        - model_name: Human-readable model name
        - schedule_name: Schedule slug
        - workflow_name: Parent workflow name
        - overrides: Additional keys merged into the payload

    Return:
        Workflow input dictionary
    """
    payload: dict[str, Any] = {
        'topic': 'e2e-topic',
        'model_id': model_id,
        'model_name': model_name,
        'schedule_name': schedule_name,
        'workflow_name': workflow_name,
        'trigger_laborious': False,
        'filters': {},
        'schema': 'sientia_data',
        'table_name': 'laborious_data',
        'retention_time': 3600,
        'fill_missing_tags': False,
        'model_tags': default_model_tags(names=['tag1']),
    }
    payload.update(overrides)
    return payload


def apply_pi_web_api_server_config(server: Any, config: dict[str, Any] | None) -> None:
    """
    Configure the in-process PI Web API stub from a scenario pi_web_api_server block.

    Args:
        - server: PIWebAPITestServer instance
        - config: Optional mode/rows/timeout_sleep_seconds dict from scenario JSON
    """
    if not config:
        return
    server.set_mode(
        config.get('mode', 'success'),
        rows=config.get('rows'),
        timeout_sleep_seconds=config.get('timeout_sleep_seconds', 60),
    )


def postgres_connection_parts(connection_url: str) -> dict[str, Any]:
    """
    Parse a SQLAlchemy Postgres URL into Activities postgres_config fields.

    Args:
        - connection_url: SQLAlchemy connection URL from testcontainers

    Return:
        Dict with host, port, user, password, dbname keys
    """
    parsed = urlparse(connection_url)
    return {
        'host': parsed.hostname or 'localhost',
        'port': parsed.port or 5432,
        'user': parsed.username or 'test',
        'password': parsed.password or 'test',
        'dbname': (parsed.path or '/test').lstrip('/'),
    }