app/services/stream_processor.py

"""
Real-time Drone Stream Processor
=================================
Processes live video frames via WebSocket for real-time asset detection.
Supports webcam simulation and RTSP drone feeds.
"""
import base64
import time
from typing import Any, Dict, List

import cv2
import numpy as np
from PIL import Image


class StreamProcessor:
    """
    Processes individual video frames through YOLO models
    and returns annotated frames with detection overlays.
    """

    # Category colors for drawing (BGR for OpenCV)
    COLORS_BGR = {
        "Properties & Buildings": (60, 76, 231),
        "Trees & Green Cover":   (96, 174, 39),
        "Parks & Open Spaces":   (113, 204, 46),
        "Water Bodies":          (185, 128, 41),
        "Roads & Footpaths":     (166, 165, 149),
        "Drains & Sewage":       (173, 68, 142),
        "Vehicles & Parking":    (18, 156, 243),
        "Waste Dumps":           (0, 84, 211),
    }

    def __init__(self, detector):
        self.detector = detector
        self.total_frame_count = 0
        self.fps_frame_count = 0
        self.fps = 0
        self.last_fps_time = time.time()
        self.last_detections = []

    def process_frame(
        self,
        frame_bytes: bytes,
        confidence: float = 0.35,
        skip_frames: int = 3,
    ) -> Dict[str, Any]:
        """
        Process a single video frame.

        Args:
            frame_bytes: Raw JPEG/PNG frame bytes
            confidence: Detection confidence threshold
            skip_frames: Only run inference every N frames (for performance)

        Returns:
            Dict with annotated frame (base64), detections, and FPS.
        """
        skip_frames = max(1, int(skip_frames))
        self.total_frame_count += 1
        self.fps_frame_count += 1

        # Decode frame
        np_arr = np.frombuffer(frame_bytes, np.uint8)
        frame = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)

        if frame is None:
            return {"error": "Invalid frame data"}

        # Run inference only every N frames (performance optimization)
        run_inference = self.total_frame_count % skip_frames == 0

        if run_inference:
            # Convert BGR to RGB for PIL
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            pil_image = Image.fromarray(frame_rgb)

            # Run detection (without SAHI for speed)
            result = self.detector.detect(
                image=pil_image,
                confidence=confidence,
                use_sahi=False,
                geo_transform=None,
            )
            self.last_detections = result.get("detections", [])

        # Draw detections on frame
        annotated = self._draw_detections(frame, self.last_detections)

        # Calculate FPS
        now = time.time()
        elapsed = now - self.last_fps_time
        if elapsed >= 1.0:
            self.fps = round(self.fps_frame_count / elapsed, 1)
            self.fps_frame_count = 0
            self.last_fps_time = now

        # Draw HUD (heads-up display)
        annotated = self._draw_hud(annotated, len(self.last_detections))

        # Encode annotated frame to JPEG
        success, buffer = cv2.imencode(
            ".jpg",
            annotated,
            [cv2.IMWRITE_JPEG_QUALITY, 75],
        )
        if not success:
            return {"error": "Failed to encode annotated frame"}
        frame_b64 = base64.b64encode(buffer).decode("utf-8")

        # Build category summary
        cat_counts = {}
        for d in self.last_detections:
            cat = d["category"]
            cat_counts[cat] = cat_counts.get(cat, 0) + 1

        return {
            "frame": frame_b64,
            "fps": self.fps,
            "total_detections": len(self.last_detections),
            "categories": cat_counts,
            "frame_number": self.total_frame_count,
        }

    def _draw_detections(self, frame: np.ndarray, detections: List[Dict]) -> np.ndarray:
        """Draw bounding boxes and labels on the frame."""
        annotated = frame.copy()

        for det in detections:
            bbox = det.get("bbox_pixels", [0, 0, 0, 0])
            x1, y1, x2, y2 = [int(v) for v in bbox]
            category = det["category"]
            conf = det["confidence"]
            color = self.COLORS_BGR.get(category, (255, 255, 255))

            # Draw bounding box
            cv2.rectangle(annotated, (x1, y1), (x2, y2), color, 2)

            # Draw filled label background
            label = f"{category.split('&')[0].strip()} {conf:.0%}"
            font = cv2.FONT_HERSHEY_SIMPLEX
            font_scale = 0.45
            thickness = 1
            (tw, th), _ = cv2.getTextSize(label, font, font_scale, thickness)

            cv2.rectangle(annotated, (x1, y1 - th - 8), (x1 + tw + 6, y1), color, -1)
            cv2.putText(annotated, label, (x1 + 3, y1 - 4), font, font_scale,
                        (255, 255, 255), thickness, cv2.LINE_AA)

            # Draw mask polygon if available
            if det.get("mask_polygon") and len(det["mask_polygon"]) > 2:
                pts = np.array(det["mask_polygon"], dtype=np.int32)
                overlay = annotated.copy()
                cv2.fillPoly(overlay, [pts], color)
                cv2.addWeighted(overlay, 0.25, annotated, 0.75, 0, annotated)
                cv2.polylines(annotated, [pts], True, color, 1)

        return annotated

    def _draw_hud(self, frame: np.ndarray, det_count: int) -> np.ndarray:
        """Draw heads-up display with FPS, detection count, and status."""
        h, w = frame.shape[:2]
        font = cv2.FONT_HERSHEY_SIMPLEX

        # Semi-transparent top bar
        overlay = frame.copy()
        cv2.rectangle(overlay, (0, 0), (w, 40), (0, 0, 0), -1)
        cv2.addWeighted(overlay, 0.7, frame, 0.3, 0, frame)

        # FPS (left)
        fps_color = (0, 255, 0) if self.fps > 10 else (0, 165, 255) if self.fps > 5 else (0, 0, 255)
        cv2.putText(frame, f"FPS: {self.fps}", (10, 28), font, 0.6, fps_color, 2, cv2.LINE_AA)

        # Title (center)
        title = "SPATIAL ASSET INTELLIGENCE - LIVE"
        (tw, _), _ = cv2.getTextSize(title, font, 0.5, 1)
        cv2.putText(frame, title, ((w - tw) // 2, 28), font, 0.5, (255, 255, 255), 1, cv2.LINE_AA)

        # Detection count (right)
        count_text = f"Assets: {det_count}"
        (cw, _), _ = cv2.getTextSize(count_text, font, 0.6, 2)
        cv2.putText(frame, count_text, (w - cw - 10, 28), font, 0.6, (0, 200, 255), 2, cv2.LINE_AA)

        # Pulsing recording indicator
        if int(time.time() * 2) % 2 == 0:
            cv2.circle(frame, (w - cw - 30, 24), 5, (0, 0, 255), -1)

        return frame

    def process_video_file(
        self,
        video_path: str,
        confidence: float = 0.35,
        max_frames: int = 300,
    ) -> Dict[str, Any]:
        """
        Process a video file and return frame-by-frame detection summary.
        Used for uploaded drone video files.
        """
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            return {"error": "Failed to open video file"}

        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        video_fps = cap.get(cv2.CAP_PROP_FPS) or 30
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

        # Sample frames evenly across the video
        frames_to_sample = max(1, min(max_frames, total_frames or 1))
        sample_interval = max(1, total_frames // frames_to_sample) if total_frames else 1

        all_detections = {}
        frame_idx = 0
        processed = 0

        while cap.isOpened() and processed < max_frames:
            ret, frame = cap.read()
            if not ret:
                break

            if frame_idx % sample_interval == 0:
                frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                pil_image = Image.fromarray(frame_rgb)

                result = self.detector.detect(
                    image=pil_image,
                    confidence=confidence,
                    use_sahi=False,
                )

                for d in result.get("detections", []):
                    cat = d["category"]
                    all_detections[cat] = all_detections.get(cat, 0) + 1

                processed += 1

            frame_idx += 1

        cap.release()

        return {
            "video_info": {
                "total_frames": total_frames,
                "fps": video_fps,
                "resolution": f"{width}x{height}",
                "frames_analyzed": processed,
            },
            "category_totals": all_detections,
            "total_unique_assets": sum(all_detections.values()),
        }