import cv2
import threading
import time
import asyncio
import base64
from typing import Dict, Optional
import logging

logger = logging.getLogger(__name__)

class LowCostStream(threading.Thread):
    def __init__(self, camera_id: int, rtsp_url: str):
        threading.Thread.__init__(self)
        self.camera_id = camera_id
        self.rtsp_url = rtsp_url
        self.daemon = True
        self.running = False
        self.cap = None
        self.lock = threading.Lock()
        self.last_frame = None
        self.status = 0 # 0: offline, 1: online
        
    def run(self):
        self.running = True
        logger.info(f"Starting stream for camera {self.camera_id}: {self.rtsp_url}")
        while self.running:
            if self.cap is None or not self.cap.isOpened():
                self.cap = cv2.VideoCapture(self.rtsp_url)
                if not self.cap.isOpened():
                    self.status = 0
                    time.sleep(5)
                    continue
                self.status = 1
                # Optimize buffer
                self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
            
            # Grab only - fast, low CPU
            ret = self.cap.grab()
            if not ret:
                self.status = 0
                self.cap.release()
                time.sleep(5)
                continue
            
            self.status = 1
            time.sleep(0.01) # Prevent CPU spin
            
    def get_snapshot(self) -> Optional[bytes]:
        if not self.status or self.cap is None:
            return None
        
        with self.lock:
            # Decode the actual frame
            ret, frame = self.cap.retrieve()
            if not ret:
                return None
            
            # Compress to JPEG
            _, buffer = cv2.imencode('.jpg', frame, [int(cv2.IMWRITE_JPEG_QUALITY), 60])
            return buffer.tobytes()

    def stop(self):
        self.running = False
        if self.cap:
            self.cap.release()

class VideoManager:
    _instance = None
    
    def __new__(cls):
        if cls._instance is None:
            cls._instance = super(VideoManager, cls).__new__(cls)
            cls._instance.streams: Dict[int, LowCostStream] = {}
        return cls._instance
    
    def add_stream(self, camera_id: int, rtsp_url: str):
        if camera_id in self.streams:
            self.stop_stream(camera_id)
        
        stream = LowCostStream(camera_id, rtsp_url)
        stream.start()
        self.streams[camera_id] = stream
        
    def stop_stream(self, camera_id: int):
        if camera_id in self.streams:
            self.streams[camera_id].stop()
            del self.streams[camera_id]
            
    def get_snapshot(self, camera_id: int) -> Optional[bytes]:
        if camera_id in self.streams:
            return self.streams[camera_id].get_snapshot()
        return None

    def get_all_snapshots_base64(self) -> Dict[str, str]:
        result = {}
        for cam_id, stream in self.streams.items():
            img_bytes = stream.get_snapshot()
            if img_bytes:
                b64 = base64.b64encode(img_bytes).decode('utf-8')
                result[f"cam_{cam_id}"] = f"data:image/jpeg;base64,{b64}"
            else:
                result[f"cam_{cam_id}"] = None
        return result
    
    def test_camera_connection(self, rtsp_url: str) -> Optional[str]:
        """Try to connect to RTSP and return a base64 snapshot"""
        cap = cv2.VideoCapture(rtsp_url)
        try:
            if not cap.isOpened():
                return None
            
            # Read a frame
            ret, frame = cap.read()
            if not ret:
                return None
            
            # Resize for preview if too large (optional, but good for bandwidth)
            # height, width = frame.shape[:2]
            # if width > 640:
            #     frame = cv2.resize(frame, (640, int(640 * height / width)))
            
            _, buffer = cv2.imencode('.jpg', frame, [int(cv2.IMWRITE_JPEG_QUALITY), 60])
            b64 = base64.b64encode(buffer).decode('utf-8')
            return f"data:image/jpeg;base64,{b64}"
        finally:
            cap.release()

video_manager = VideoManager()
