from pathlib import Path
from collections.abc import Callable
from torch.utils.data import Dataset
from PIL import Image
import torch

from vigilcv import VisionSentinel
from vigilcv.exceptions import VigilCVError

sentinel = VisionSentinel(blur_threshold=80.0, min_entropy=2.5)


class QualityGuardedDataset(Dataset):
    """
    A PyTorch Dataset that validates images with VigilCV before loading.
    Corrupted or degraded images are replaced with a fallback tensor
    instead of crashing the DataLoader batch collation.
    """

    def __init__(
        self,
        image_dir: str | Path,
        transform: Callable | None = None,
    ) -> None:
        self.image_paths = sorted(Path(image_dir).glob("*.jpg"))
        self.transform = transform

    def __len__(self) -> int:
        return len(self.image_paths)

    def __getitem__(self, idx: int) -> tuple[torch.Tensor, int]:
        path = self.image_paths[idx]

        try:
            # Step 1: Pre-flight quality gate
            sentinel.guard(path)

            # Step 2: Safe to load and transform
            with Image.open(path) as img:
                img_rgb = img.convert("RGB")
                tensor = self.transform(img_rgb) if self.transform else torch.zeros(3, 224, 224)
                return tensor, 1  # label=1 for valid

        except VigilCVError:
            # Return fallback tensor — preserves batch collation
            return torch.zeros(3, 224, 224), -1  # label=-1 for invalid

from torch.utils.data import DataLoader
from torchvision import transforms

transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

dataset = QualityGuardedDataset("dataset/train/", transform=transform)
loader = DataLoader(dataset, batch_size=32, num_workers=4, pin_memory=True)

for images, labels in loader:
    valid_mask = labels != -1
    if valid_mask.sum() == 0:
        continue  # Skip fully-degraded batches
    
    images = images[valid_mask]
    labels = labels[valid_mask]
    
    outputs = model(images)
    loss = criterion(outputs, labels)
    loss.backward()

import cv2
from vigilcv import VisionSentinel

sentinel = VisionSentinel(
    blur_threshold=80.0,
    min_entropy=2.5,
    raise_on_fail=False,  # Don't raise; return metrics with passed=False
)

cap = cv2.VideoCapture("rtsp://camera.local/stream1")
# Or for webcam: cap = cv2.VideoCapture(0)

for frame, metrics in sentinel.audit_stream(cap):
    if metrics is None:
        break  # Stream ended

    # Overlay quality telemetry
    status = "PASS" if metrics.passed else "FAIL"
    color = (0, 200, 100) if metrics.passed else (0, 80, 240)  # BGR

    cv2.putText(
        frame,
        f"Blur: {metrics.blur_score:.1f} | Entropy: {metrics.shannon_entropy:.2f} | {status}",
        (10, 30),
        cv2.FONT_HERSHEY_SIMPLEX,
        0.7,
        color,
        2,
    )

    if not metrics.passed:
        # Draw red border on degraded frames
        cv2.rectangle(frame, (0, 0), (frame.shape[1], frame.shape[0]), (0, 0, 255), 4)

    cv2.imshow("VigilCV Monitor", frame)
    if cv2.waitKey(1) & 0xFF == ord("q"):
        break

cap.release()
cv2.destroyAllWindows()

for frame, metrics in sentinel.audit_stream(cap, sample_every=5):
    # Audit 1 in every 5 frames — 12 FPS effective audit rate
    ...

import cv2
from vigilcv import VisionSentinel

sentinel = VisionSentinel(blur_threshold=100.0)
cap = cv2.VideoCapture("rtsp://drone-camera/main")
writer = cv2.VideoWriter("clean_output.mp4", cv2.VideoWriter_fourcc(*"mp4v"), 30, (1920, 1080))

for frame, metrics in sentinel.audit_stream(cap):
    if metrics and metrics.passed:
        writer.write(frame)  # Only write pristine frames

writer.release()

import time
import cv2
from vigilcv import VisionSentinel

sentinel = VisionSentinel(blur_threshold=50.0, min_entropy=1.5)
cap = cv2.VideoCapture(0)

fail_streak = 0
ALERT_THRESHOLD = 30  # 30 consecutive failed frames

for frame, metrics in sentinel.audit_stream(cap):
    if not metrics or not metrics.passed:
        fail_streak += 1
        if fail_streak >= ALERT_THRESHOLD:
            print(f"ALERT: Camera {cap} potentially occluded — {fail_streak} consecutive failures")
            # Trigger PagerDuty / Slack alert here
    else:
        fail_streak = 0

from fastapi import FastAPI, UploadFile
from vigilcv import VisionSentinel
import numpy as np

app = FastAPI()
sentinel = VisionSentinel()

@app.post("/predict")
async def predict(file: UploadFile):
    payload = await file.read()
    raw_array = np.frombuffer(payload, dtype=np.uint8)
    
    # Validation occurs in microseconds
    metrics = sentinel.audit_stream(raw_array)
    if metrics.is_degraded:
        return {"status": "rejected", "reason": metrics.failure_reason}
        
    return {"status": "passed"}

from torch.utils.data import Dataset
import torch
from vigilcv.integrations.torch import FastTensorSentinel

class ResilientVisionDataset(Dataset):
    def __init__(self, records):
        self.records = records
        self.sentinel = FastTensorSentinel(strict_mode=False)

    def __getitem__(self, index: int):
        tensor = self._load_tensor(self.records[index])
        
        is_valid, diag = self.sentinel.validate_tensor(tensor)
        if not is_valid:
            # Mask corrupted sample with running batch median to protect gradient stability
            return torch.zeros_like(tensor), -1
            
        return tensor, 1

import cv2
from vigilcv import VisionSentinel

sentinel = VisionSentinel()
cap = cv2.VideoCapture("rtsp://camera_ip:554/stream")

while True:
    ret, frame = cap.read()
    if not ret:
        break
        
    metrics = sentinel.audit(frame)
    if metrics.is_degraded:
        # Trigger alarm or skip frame
        cv2.putText(frame, "CORRUPTED", (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
        
    cv2.imshow("Stream", frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

VisionSentinel(
    blur_threshold: float = 100.0,
    min_entropy: float = 3.0,
    max_underexposure_ratio: float = 0.20,
    max_overexposure_ratio: float = 0.20,
    raise_on_fail: bool = True,
)

metrics = sentinel.guard("path/to/image.jpg")
metrics = sentinel.guard(Path("image.png"))
metrics = sentinel.guard(io.BytesIO(image_bytes))  # FastAPI UploadFile

for frame, metrics in sentinel.audit_stream(cap):
    if metrics and metrics.passed:
        process(frame)

BatchAuditor(
    directory: str | Path,
    workers: int = 4,
    sentinel: VisionSentinel | None = None,
    glob_pattern: str = "**/*.{jpg,jpeg,png,webp}",
)

auditor = BatchAuditor("dataset/train/", workers=8)
summary = auditor.run()

@dataclass(frozen=True)
class QualityMetrics:
    path: Path | None           # Source file path (None for BytesIO)
    blur_score: float           # Laplacian variance
    shannon_entropy: float      # Shannon entropy (bits)
    underexposure_ratio: float  # Fraction of dark-clipped pixels
    overexposure_ratio: float   # Fraction of saturated pixels
    width: int                  # Image width in pixels
    height: int                 # Image height in pixels
    passed: bool                # True if all gates passed
    warnings: list[str]         # Human-readable failure reasons
    latency_ms: float           # Wall-clock audit time in milliseconds

@dataclass(frozen=True)
class AuditSummary:
    total_images: int
    valid_images: int
    corrupted_count: int
    blurred_count: int
    underexposed_count: int
    overexposed_count: int
    throughput_fps: float
    total_duration_s: float
    drift_report: DriftReport | None  # Only if baseline provided
    per_image_metrics: list[QualityMetrics]

@dataclass(frozen=True)
class DriftReport:
    is_drifted: bool
    wasserstein_distance: float
    mmd_score: float
    top_drift_features: list[str]
    reference_n: int
    query_n: int
    threshold: float

from vigilcv.exceptions import (
    VigilCVError,              # Base exception
    CorruptImageError,         # Unreadable / truncated file
    QualityThresholdExceeded,  # Quality gate failure
    DriftDetectedError,        # Distribution drift exceeds threshold
)

try:
    sentinel.guard("image.jpg")
except QualityThresholdExceeded as e:
    print(e.metrics.blur_score)  # Access the actual scores
    print(e.metrics.warnings)

from vigilcv.core.drift import DriftEngine

engine = DriftEngine()
engine.fit(reference_directory="dataset/train/")
engine.save("baseline.pkl")

report = engine.detect(query_directory="dataset/new_batch/")

if report.is_drifted:
    print(f"DRIFT DETECTED: {report.wasserstein_distance:.4f} W1")

# Inspect a single image
vigilcv inspect photo.jpg --blur-threshold 150

# Audit a massive directory in parallel
vigilcv audit dataset/train/ --workers 16 --report quality_report.html

# Detect distribution drift against baseline
vigilcv drift dataset/new/ --baseline baselines/prod_v1.pkl --threshold 0.20