import cv2
import numpy as np
import subprocess
import yt_dlp
from ultralytics import YOLO
from flask import Flask, Response

app = Flask(__name__)

# currently cuda not working with torch version, use cpu

YOUTUBE_URL = "https://www.youtube.com/watch?v=i3w7qZVSAsY"  # Example stream
CONFIDENCE_THRESHOLD = 0.25
MODEL_PATH = "yolov8n.pt"  # Using YOLOv8 nano for speed

# Load YOLOv8 model
print("Loading YOLOv8 model...")
model = YOLO(MODEL_PATH)
model.to("cpu")  # set for cpu
print("YOLOv8 loaded successfully!")

def get_stream_url():
    """Fetch fresh 720p YouTube stream URL."""
    ydl_opts = {'quiet': True, 'format': 'bestvideo[height=720]'}
    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        info_dict = ydl.extract_info(YOUTUBE_URL, download=False)
        return info_dict.get("url", None)

def generate_frames():
    """Capture video frames, apply object detection, and anonymize persons."""
    stream_url = get_stream_url()
    if not stream_url:
        print("Failed to fetch stream URL!")
        return

    print("Starting FFmpeg stream...")
    ffmpeg_process = subprocess.Popen([
        "ffmpeg", "-re", "-i", stream_url, "-r", "10",  # Set frame rate to 10 FPS
        "-fflags", "nobuffer", "-flags", "low_delay",
        "-f", "rawvideo", "-pix_fmt", "bgr24", "pipe:1"
    ], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, bufsize=10**8)

    while True:
        raw_frame = ffmpeg_process.stdout.read(1280 * 720 * 3)  # Read raw BGR frame
        if not raw_frame:
            print("No frame received!")
            break

        frame = np.frombuffer(raw_frame, np.uint8).reshape((720, 1280, 3)).copy()  # Make writable


        # Run YOLOv8 object detection
        results = model(frame)[0]  # Get first result

        for box in results.boxes:
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            confidence = box.conf[0].item()
            class_id = int(box.cls[0].item())
            label = model.names[class_id]

            if confidence > CONFIDENCE_THRESHOLD and label == "person":
                roi = frame[y1:y2, x1:x2]
                avg_color = np.mean(roi, axis=(0, 1), dtype=int)  # Compute avg color
                frame[y1:y2, x1:x2] = avg_color  # Fill with avg color

        # Encode and yield the frame as JPEG
        _, buffer = cv2.imencode('.jpg', frame)
        yield (b'--frame\r\n'
               b'Content-Type: image/jpeg\r\n\r\n' + buffer.tobytes() + b'\r\n')

@app.route('/video')
def video_feed():
    """Stream processed video frames."""
    return Response(generate_frames(), mimetype='multipart/x-mixed-replace; boundary=frame')

if __name__ == '__main__':
    print("Running at http://localhost:5000/video")
    app.run(host='0.0.0.0', port=5000, debug=True, threaded=True)