#!/usr/bin/env python3
"""
streaming_guardrail.py - Runtime Tripwire for Full Tensor Load Detection

WO-STREAM-FULL-BLOCK-01: Sharp Edge 5C

Runtime tripwire that fails if any code tries to load a full tensor
that exceeds the allowed threshold. This ensures true streaming mode
is maintained and catches accidental full tensor loads.

Usage:
    with StreamingGuardrail(max_tensor_mb=32):
        output = stream_attention_forward(X, ...)
    # If any single allocation > 32MB, raises GuardrailViolation

    # Or as a decorator
    @streaming_guardrail(max_tensor_mb=32)
    def my_streaming_function(X, ...):
        ...

Modes:
    - "fail": Raise GuardrailViolation on large allocation (default)
    - "warn": Log warning but continue
    - "log": Silent logging for analysis

Work Order: WO-STREAM-FULL-BLOCK-01 (Sharp Edge 5C)
"""

from __future__ import annotations

import functools
import sys
import threading
import tracemalloc
import warnings
from contextlib import contextmanager
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, Callable, Dict, List, Literal, Optional, TypeVar

F = TypeVar('F', bound=Callable[..., Any])


class GuardrailViolation(Exception):
    """Raised when a full tensor load is detected in streaming context."""

    def __init__(
        self,
        message: str,
        allocation_mb: float,
        max_allowed_mb: float,
        traceback_frames: List[str],
    ):
        super().__init__(message)
        self.allocation_mb = allocation_mb
        self.max_allowed_mb = max_allowed_mb
        self.traceback_frames = traceback_frames


@dataclass
class AllocationEvent:
    """Record of a single large allocation."""
    size_mb: float
    timestamp: str
    traceback_snippet: str


@dataclass
class GuardrailReport:
    """Summary report from guardrail monitoring."""
    max_single_allocation_mb: float = 0.0
    total_allocations: int = 0
    large_allocations: int = 0
    large_allocation_events: List[AllocationEvent] = field(default_factory=list)
    violations: int = 0
    mode: str = "fail"
    threshold_mb: float = 32.0
    status: str = "PASS"
    timestamp: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"))

    def to_dict(self) -> Dict[str, Any]:
        return {
            "max_single_allocation_mb": self.max_single_allocation_mb,
            "total_allocations": self.total_allocations,
            "large_allocations": self.large_allocations,
            "large_allocation_events": [
                {
                    "size_mb": e.size_mb,
                    "timestamp": e.timestamp,
                    "traceback_snippet": e.traceback_snippet,
                }
                for e in self.large_allocation_events
            ],
            "violations": self.violations,
            "mode": self.mode,
            "threshold_mb": self.threshold_mb,
            "status": self.status,
            "timestamp": self.timestamp,
        }


class StreamingGuardrail:
    """
    Runtime tripwire: fails if any code tries to load a full tensor.

    Uses tracemalloc to monitor memory allocations and detect when
    any single allocation exceeds the threshold.

    Usage:
        with StreamingGuardrail(max_tensor_mb=32):
            output = stream_attention_forward(X, ...)
        # If any single tensor > 32MB loaded, raises GuardrailViolation

    Args:
        max_tensor_mb: Maximum allowed single allocation in MB (default: 32)
        mode: "fail" | "warn" | "log"
            - fail: Raise GuardrailViolation (default)
            - warn: Log warning but continue
            - log: Silent logging for analysis
        check_interval: How often to check allocations (in frames, default: 100)
        record_top_n: How many large allocations to record (default: 10)

    Note:
        This uses tracemalloc which has some overhead (~10-30% slowdown).
        For production, consider using "log" mode for analysis first.
    """

    _active_guardrails: Dict[int, 'StreamingGuardrail'] = {}  # thread_id -> guardrail
    _lock = threading.Lock()

    def __init__(
        self,
        max_tensor_mb: float = 32.0,
        mode: Literal["fail", "warn", "log"] = "fail",
        check_interval: int = 100,
        record_top_n: int = 10,
    ):
        self.max_tensor_mb = max_tensor_mb
        self.mode = mode
        self.check_interval = check_interval
        self.record_top_n = record_top_n

        self._threshold_bytes = int(max_tensor_mb * 1024 * 1024)
        self._report = GuardrailReport(
            mode=mode,
            threshold_mb=max_tensor_mb,
        )
        self._prev_snapshot: Optional[tracemalloc.Snapshot] = None
        self._was_tracing = False

    def __enter__(self) -> 'StreamingGuardrail':
        """Start monitoring."""
        thread_id = threading.get_ident()

        with self._lock:
            if thread_id in self._active_guardrails:
                # Nested guardrail - use the stricter threshold
                existing = self._active_guardrails[thread_id]
                self.max_tensor_mb = min(self.max_tensor_mb, existing.max_tensor_mb)
                self._threshold_bytes = int(self.max_tensor_mb * 1024 * 1024)

            self._active_guardrails[thread_id] = self

        # Start or continue tracemalloc
        self._was_tracing = tracemalloc.is_tracing()
        if not self._was_tracing:
            tracemalloc.start()

        # Take initial snapshot
        self._prev_snapshot = tracemalloc.take_snapshot()

        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> bool:
        """Stop monitoring and finalize report."""
        thread_id = threading.get_ident()

        # Take final snapshot and check
        try:
            self._check_allocations()
        except GuardrailViolation:
            if exc_type is None:  # Only raise if no other exception
                raise

        # Cleanup
        with self._lock:
            if thread_id in self._active_guardrails:
                del self._active_guardrails[thread_id]

        if not self._was_tracing:
            tracemalloc.stop()

        # Finalize report
        if self._report.violations > 0:
            self._report.status = "FAIL"

        return False  # Don't suppress exceptions

    def _check_allocations(self) -> None:
        """Check for large allocations since last snapshot."""
        if self._prev_snapshot is None:
            return

        current_snapshot = tracemalloc.take_snapshot()

        # Compare snapshots
        diff = current_snapshot.compare_to(self._prev_snapshot, 'lineno')

        for stat in diff:
            size_mb = stat.size / (1024 * 1024)
            self._report.total_allocations += 1

            if stat.size > self._threshold_bytes:
                self._report.large_allocations += 1
                self._report.max_single_allocation_mb = max(
                    self._report.max_single_allocation_mb,
                    size_mb,
                )

                # Get traceback snippet
                frame = stat.traceback[0] if stat.traceback else None
                tb_snippet = f"{frame.filename}:{frame.lineno}" if frame else "unknown"

                event = AllocationEvent(
                    size_mb=size_mb,
                    timestamp=datetime.now(timezone.utc).isoformat(),
                    traceback_snippet=tb_snippet,
                )

                # Record top N
                if len(self._report.large_allocation_events) < self.record_top_n:
                    self._report.large_allocation_events.append(event)
                else:
                    # Replace smallest if this is larger
                    min_idx = min(
                        range(len(self._report.large_allocation_events)),
                        key=lambda i: self._report.large_allocation_events[i].size_mb
                    )
                    if size_mb > self._report.large_allocation_events[min_idx].size_mb:
                        self._report.large_allocation_events[min_idx] = event

                # Handle violation based on mode
                self._handle_violation(size_mb, tb_snippet, stat.traceback)

        self._prev_snapshot = current_snapshot

    def _handle_violation(
        self,
        size_mb: float,
        tb_snippet: str,
        traceback: tracemalloc.Traceback,
    ) -> None:
        """Handle a threshold violation based on mode."""
        self._report.violations += 1

        message = (
            f"StreamingGuardrail: Large allocation detected!\n"
            f"  Size: {size_mb:.2f} MB (threshold: {self.max_tensor_mb} MB)\n"
            f"  Location: {tb_snippet}\n"
            f"  This may indicate a full tensor load instead of streaming."
        )

        # Format traceback
        tb_frames = [
            f"{frame.filename}:{frame.lineno}"
            for frame in traceback
        ] if traceback else []

        if self.mode == "fail":
            raise GuardrailViolation(
                message=message,
                allocation_mb=size_mb,
                max_allowed_mb=self.max_tensor_mb,
                traceback_frames=tb_frames,
            )
        elif self.mode == "warn":
            warnings.warn(message, RuntimeWarning, stacklevel=4)
        # "log" mode: silent, just record in report

    def check_now(self) -> None:
        """Manually trigger an allocation check."""
        self._check_allocations()

    def get_report(self) -> GuardrailReport:
        """Get the current monitoring report."""
        return self._report


def streaming_guardrail(
    max_tensor_mb: float = 32.0,
    mode: Literal["fail", "warn", "log"] = "fail",
) -> Callable[[F], F]:
    """
    Decorator version of StreamingGuardrail.

    Usage:
        @streaming_guardrail(max_tensor_mb=32)
        def my_streaming_function(X, ...):
            ...
    """
    def decorator(func: F) -> F:
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            with StreamingGuardrail(max_tensor_mb=max_tensor_mb, mode=mode):
                return func(*args, **kwargs)
        return wrapper  # type: ignore
    return decorator


@contextmanager
def assert_streaming_mode(
    max_tensor_mb: float = 32.0,
    context_name: str = "operation",
):
    """
    Context manager that asserts streaming mode is maintained.

    More user-friendly than StreamingGuardrail for tests.

    Usage:
        with assert_streaming_mode(max_tensor_mb=32, context_name="attention"):
            output = stream_attention_forward(X, ...)
        # Raises AssertionError with friendly message if violated
    """
    guardrail = StreamingGuardrail(max_tensor_mb=max_tensor_mb, mode="log")

    try:
        with guardrail:
            yield guardrail
    finally:
        report = guardrail.get_report()
        if report.violations > 0:
            largest = report.max_single_allocation_mb
            events = report.large_allocation_events
            locations = [e.traceback_snippet for e in events[:3]]

            raise AssertionError(
                f"Streaming mode violated in {context_name}!\n"
                f"  Largest allocation: {largest:.2f} MB (threshold: {max_tensor_mb} MB)\n"
                f"  Violations: {report.violations}\n"
                f"  Locations: {', '.join(locations)}\n"
                f"  This indicates a full tensor was loaded instead of being streamed."
            )


def get_current_guardrail() -> Optional[StreamingGuardrail]:
    """Get the currently active guardrail for this thread, if any."""
    thread_id = threading.get_ident()
    with StreamingGuardrail._lock:
        return StreamingGuardrail._active_guardrails.get(thread_id)


if __name__ == "__main__":
    import numpy as np

    print("=== Streaming Guardrail Demo ===")
    print()

    # Demo 1: Safe allocation (should pass)
    print("Demo 1: Safe allocation (10 MB < 32 MB threshold)")
    try:
        with StreamingGuardrail(max_tensor_mb=32) as guard:
            small_array = np.zeros((1000, 1000), dtype=np.float32)  # ~4 MB
            print(f"  Allocated {small_array.nbytes / (1024 * 1024):.1f} MB")
        print("  PASS: No violation")
    except GuardrailViolation as e:
        print(f"  FAIL: {e}")
    print()

    # Demo 2: Large allocation (should fail in "fail" mode)
    print("Demo 2: Large allocation (100 MB > 32 MB threshold) in 'fail' mode")
    try:
        with StreamingGuardrail(max_tensor_mb=32, mode="fail") as guard:
            large_array = np.zeros((5000, 5000), dtype=np.float32)  # ~100 MB
            print(f"  Allocated {large_array.nbytes / (1024 * 1024):.1f} MB")
        print("  UNEXPECTED: Should have raised GuardrailViolation")
    except GuardrailViolation as e:
        print(f"  PASS: Caught violation - {e.allocation_mb:.1f} MB > {e.max_allowed_mb} MB")
    print()

    # Demo 3: Large allocation in "warn" mode
    print("Demo 3: Large allocation in 'warn' mode")
    try:
        with StreamingGuardrail(max_tensor_mb=32, mode="warn") as guard:
            large_array = np.zeros((5000, 5000), dtype=np.float32)  # ~100 MB
            print(f"  Allocated {large_array.nbytes / (1024 * 1024):.1f} MB")
        report = guard.get_report()
        print(f"  Report: {report.violations} violations, max={report.max_single_allocation_mb:.1f} MB")
    except GuardrailViolation:
        print("  UNEXPECTED: Should not raise in 'warn' mode")
    print()

    print("Demo complete.")
