import os
import cv2
import tempfile
import requests
import base64
import numpy as np
from PIL import Image
from io import BytesIO
from ultralytics import YOLO
import streamlit as st
import yt_dlp as youtube_dl

def load_yolov8_model(model_name='yolov8s.pt'):
    try:
        return YOLO(model_name)
    except Exception as e:
        st.error(f"Error loading model: {e}")
        return None

def detect_objects(image, model):
    try:
        results = model(image)
        for result in results[0].boxes:
            x1, y1, x2, y2 = map(int, result.xyxy[0])
            label = model.names[int(result.cls)]
            confidence = result.conf.item()
            cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
            label_text = f'{label} {confidence:.2f}'
            cv2.putText(image, label_text, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
        return image
    except Exception as e:
        st.error(f"Error during object detection: {e}")
        return None

def process_image_with_yolov8(model_name, image=None, url=None):
    model = load_yolov8_model(model_name)
    if model is None:
        return None

    if url:
        if url.startswith('data:image'):
            try:
                header, encoded = url.split(',', 1)
                image_data = base64.b64decode(encoded)
                image = Image.open(BytesIO(image_data))
            except Exception as e:
                st.error(f"Error decoding base64 image: {e}")
                return None
        else:
            try:
                response = requests.get(url)
                response.raise_for_status()
                image = Image.open(BytesIO(response.content))
            except Exception as e:
                st.error(f"Error loading image from URL: {e}")
                return None

    try:
        image = np.array(image)
        output_image = detect_objects(image, model)
        return output_image
    except Exception as e:
        st.error(f"Error processing image: {e}")
        return None

def download_youtube_video(youtube_url):
    try:
        temp_dir = tempfile.gettempdir()
        output_path = os.path.join(temp_dir, 'downloaded_video.mp4')
        ydl_opts = {
            'format': 'best',
            'outtmpl': output_path
        }
        with youtube_dl.YoutubeDL(ydl_opts) as ydl:
            ydl.download([youtube_url])
        return output_path
    except Exception as e:
        st.error(f"Failed to retrieve video from YouTube. Error: {e}")
        return None

def process_video(input_video_path, output_video_path, model):
    cap = cv2.VideoCapture(input_video_path)
    if not cap.isOpened():
        st.error(f"Error: Cannot open video file {input_video_path}")
        return None

    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS)
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')

    out = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        processed_frame = detect_objects(frame, model)
        out.write(processed_frame)

    cap.release()
    out.release()

    return output_video_path

st.title("YOLOv8 Object Detection on Images and Videos")

model_choice = st.selectbox("Select Model", ["yolov8s.pt", "yolov8m.pt"])

tabs = st.tabs(["Image Detection", "Video Detection"])

with tabs[0]:
    st.header("Image Detection")
    input_choice = st.radio("Select Input Method", ["Upload", "URL"])

    if input_choice == "Upload":
        uploaded_image = st.file_uploader("Upload Image", type=["jpg", "jpeg", "png"])
        if uploaded_image is not None:
            image = Image.open(uploaded_image)
            processed_image = process_image_with_yolov8(model_choice, image=image)
            if processed_image is not None:
                st.image(processed_image, caption="Processed Image", use_column_width=True)
    elif input_choice == "URL":
        image_url = st.text_input("Image URL")
        if image_url:
            processed_image = process_image_with_yolov8(model_choice, url=image_url)
            if processed_image is not None:
                st.image(processed_image, caption="Processed Image", use_column_width=True)

with tabs[1]:
    st.header("Video Detection")
    video_choice = st.radio("Select Input Method", ["Upload", "YouTube"])

    if video_choice == "Upload":
        uploaded_video = st.file_uploader("Upload Local Video", type=["mp4", "mov", "avi"])
        if uploaded_video is not None:
            input_video_path = os.path.join(tempfile.gettempdir(), uploaded_video.name)
            with open(input_video_path, "wb") as f:
                f.write(uploaded_video.read())
            model = load_yolov8_model(model_choice)
            output_video_path = os.path.join(tempfile.gettempdir(), "processed_video.mp4")
            processed_video = process_video(input_video_path, output_video_path, model)
            if processed_video is not None:
                st.video(processed_video)
    elif video_choice == "YouTube":
        video_url = st.text_input("YouTube Video URL")
        if video_url:
            input_video_path = download_youtube_video(video_url)
            if input_video_path:
                model = load_yolov8_model(model_choice)
                output_video_path = os.path.join(tempfile.gettempdir(), "processed_video.mp4")
                processed_video = process_video(input_video_path, output_video_path, model)
                if processed_video is not None:
                    st.video(processed_video)