# preprocessing/monitoring/telemetry.py
import logging
import time
from contextlib import contextmanager
from typing import Dict, Any

try:
    import psutil
except ImportError:  # pragma: no cover
    psutil = None

log = logging.getLogger(__name__)


class StageTelemetry:
    """Track performance metrics per stage"""

    def __init__(self):
        self.metrics: Dict[str, Any] = {}

    @contextmanager
    def measure_stage(self, stage_name: str, row_count: int):
        """Context manager to measure stage performance"""
        start_time = time.time()
        start_memory = (
            psutil.Process().memory_info().rss / 1024 / 1024 if psutil else 0.0
        )

        yield

        elapsed = time.time() - start_time
        end_memory = (
            psutil.Process().memory_info().rss / 1024 / 1024 if psutil else 0.0
        )

        rows_per_sec = row_count / elapsed if elapsed > 0 else 0

        self.metrics[stage_name] = {
            "elapsed_sec": elapsed,
            "rows_per_sec": rows_per_sec,
            "memory_delta_mb": end_memory - start_memory,
            "row_count": row_count,
        }

        log.info(
            f"Stage '{stage_name}': {row_count:,} rows in {elapsed:.1f}s "
            f"({rows_per_sec:,.0f} rows/sec, {end_memory - start_memory:+.1f}MB)"
        )
