# uv pip install imutils pypylon ultralytics

from pypylon import pylon
import cv2
from ultralytics import YOLO
from http.server import BaseHTTPRequestHandler, HTTPServer
from socketserver import ThreadingMixIn
import imutils
import threading
import time
import os

# --- Global frame buffer ---
latest_frame = None
frame_lock = threading.Lock()
frame_sequence = 0  # Optional: for debugging frame freshness

# --- 1. Load YOLOE Prompt-Free Model ---
model = YOLO('yoloe-11l-seg-pf.pt')  # ← Prompt-Free! Detects anything

# --- 2. Setup Basler Camera ---
camera = pylon.InstantCamera(pylon.TlFactory.GetInstance().CreateFirstDevice())
camera.Open()

# >>> LOAD .PFS CONFIGURATION FILE DIRECTLY <<<<
pfs_file_path = "./a2A1920-51gcPRO_v1_calo.pfs"  # ← Make sure this exists!
if pfs_file_path and os.path.exists(pfs_file_path):
    try:
        pylon.FeaturePersistence.Load(pfs_file_path, camera.GetNodeMap(), True)
        print(f"Loaded camera config: {pfs_file_path}")
    except Exception as e:
        print(f"Error loading .pfs: {e}")

# Start grabbing — always get latest frame
camera.StartGrabbing(pylon.GrabStrategy_LatestImageOnly)

print("Camera started. Streaming to http://192.168.1.99:7860/video")

# --- 3. Frame Capture & Processing Thread ---
def capture_and_process():
    global latest_frame, frame_sequence
    last_process_time = time.time()
    while camera.IsGrabbing():
        grab_result = camera.RetrieveResult(5000, pylon.TimeoutHandling_ThrowException)
        if grab_result.GrabSucceeded():
            frame = grab_result.GetArray()

            # --- Handle BayerBG12 → convert to 8-bit BGR ---
            if frame.dtype == 'uint16' and len(frame.shape) == 2:
                # Debayer BG to BGR (works directly on uint16)
                debayered = cv2.cvtColor(frame, cv2.COLOR_BAYER_BG2BGR)
                # Scale 12-bit (0-4095) to 8-bit (0-255) — accurate linear mapping
                frame = cv2.convertScaleAbs(debayered, alpha=(255.0 / 4095.0))
            else:
                print(f"Unexpected frame format: shape={frame.shape}, dtype={frame.dtype}")
                grab_result.Release()
                continue

            # Optional: Cap processing to ~25 FPS to avoid overload
            now = time.time()
            if now - last_process_time < 0.04:  # ~25 FPS
                grab_result.Release()
                continue
            last_process_time = now

            # Run YOLOE predict
            results = model.predict(frame, conf=0.25, iou=0.45)
            annotated_frame = results[0].plot()

            # Optional: Draw frame number for debugging latency
            global frame_sequence
            frame_sequence += 1
            cv2.putText(annotated_frame, f"Frame #{frame_sequence}", (10, 30),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)

            # Resize for smoother streaming
            annotated_frame = imutils.resize(annotated_frame, width=800)

            # Update global frame
            with frame_lock:
                latest_frame = annotated_frame.copy()

        grab_result.Release()

# --- 4. MJPEG Streaming Server ---
class MJPEGHandler(BaseHTTPRequestHandler):
    def do_GET(self):
        if self.path == '/video':
            self.send_response(200)
            self.send_header('Content-type', 'multipart/x-mixed-replace; boundary=--jpgboundary')
            # Disable caching — critical for live streams
            self.send_header('Cache-Control', 'no-cache, no-store, must-revalidate')
            self.send_header('Pragma', 'no-cache')
            self.send_header('Expires', '0')
            self.end_headers()

            try:
                last_frame_sent = 0
                while True:
                    with frame_lock:
                        if latest_frame is None:
                            time.sleep(0.01)
                            continue
                        frame = latest_frame.copy()

                    # Cap streaming to ~20 FPS to match client and avoid flooding
                    now = time.time()
                    if now - last_frame_sent < 0.05:  # 20 FPS
                        time.sleep(0.001)
                        continue
                    last_frame_sent = now

                    # Encode frame to JPEG
                    ret, jpeg = cv2.imencode('.jpg', frame, [int(cv2.IMWRITE_JPEG_QUALITY), 70])
                    if not ret:
                        continue

                    # Send frame
                    try:
                        self.wfile.write(b"--jpgboundary\r\n")
                        self.send_header('Content-type', 'image/jpeg')
                        self.send_header('Content-length', str(len(jpeg)))
                        self.end_headers()
                        self.wfile.write(jpeg.tobytes())
                        self.wfile.write(b'\r\n')
                        self.wfile.flush()  # ← Force immediate send, reduce buffering
                    except (BrokenPipeError, ConnectionResetError):
                        print("Client disconnected. Stopping stream.")
                        break  # Client gone — stop sending

            except Exception as e:
                print(f"Stream error: {e}")

        else:
            self.send_response(404)
            self.end_headers()

class ThreadedHTTPServer(ThreadingMixIn, HTTPServer):
    """Handle requests in a separate thread."""

# --- 5. Start Threads ---
if __name__ == "__main__":
    capture_thread = threading.Thread(target=capture_and_process, daemon=True)
    capture_thread.start()

    server = ThreadedHTTPServer(('0.0.0.0', 7860), MJPEGHandler)
    print("Server listening on http://192.168.1.99:7860/video")

    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\n🛑 Interrupted by user.")
    finally:
        print("🛑 Stopping camera and server...")
        if camera.IsGrabbing():
            camera.StopGrabbing()
        camera.Close()
        server.server_close()
        print("✅ Server stopped cleanly.")