import os
import sys
import tempfile
import cv2
import requests
from ultralytics import YOLO
import streamlit as st

# Set page configuration
st.set_page_config(
    page_title="People Tracking with YOLO11-pose",
    page_icon="👥",
    layout="wide",
    initial_sidebar_state="expanded"
)
st.title("People Tracking with YOLO11-pose")

# Sidebar: Input method and settings
st.sidebar.header("Input Settings")
uploaded_file = st.sidebar.file_uploader("Upload Image/Video", type=["jpg", "jpeg", "png", "bmp", "webp", "mp4"])
youtube_link = st.sidebar.text_input("YouTube Link (optional)", "")
image_url = st.sidebar.text_input("Image URL (optional)", "")
sensitivity = st.sidebar.slider("Sensitivity (Confidence Threshold)", 0.0, 1.0, 0.2, step=0.01)
process_button = st.sidebar.button("Process Input")

# Define the video extensions for later use
video_exts = [".mp4", ".mov", ".avi", ".webm"]

def process_input(uploaded_file, youtube_link, image_url, sensitivity):
    input_path = None
    temp_files = []

    # Input priority: YouTube link > Image URL > Uploaded file.
    if youtube_link and youtube_link.strip():
        try:
            from pytubefix import YouTube
            yt = YouTube(youtube_link)
            stream = yt.streams.filter(file_extension='mp4', progressive=True).order_by("resolution").desc().first()
            if not stream:
                return None, None, None, "No suitable mp4 stream found."
            temp_path = os.path.join(tempfile.gettempdir(), f"yt_{os.urandom(8).hex()}.mp4")
            stream.download(output_path=tempfile.gettempdir(), filename=os.path.basename(temp_path))
            input_path = temp_path
            temp_files.append(input_path)
        except Exception as e:
            return None, None, None, f"Error downloading YouTube video: {str(e)}"
    elif image_url and image_url.strip():
        try:
            response = requests.get(image_url, stream=True, timeout=10)
            response.raise_for_status()
            temp_path = os.path.join(tempfile.gettempdir(), f"img_{os.urandom(8).hex()}.jpg")
            with open(temp_path, "wb") as f:
                f.write(response.content)
            input_path = temp_path
            temp_files.append(input_path)
        except Exception as e:
            return None, None, None, f"Error downloading image: {str(e)}"
    elif uploaded_file is not None:
        # Save the uploaded file to a temporary file
        ext = os.path.splitext(uploaded_file.name)[1]
        with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as tmp:
            tmp.write(uploaded_file.read())
            input_path = tmp.name
            temp_files.append(input_path)
    else:
        return None, None, None, "Please provide an input."

    ext = os.path.splitext(input_path)[1].lower()
    output_path = None

    # Load the YOLO model (ensure the model file is available in your repository)
    model = YOLO("yolo11n-pose.pt")

    try:
        if ext in video_exts:
            # Video processing
            cap = cv2.VideoCapture(input_path)
            if not cap.isOpened():
                return None, None, None, f"Cannot open video file: {input_path}"
            
            fps = cap.get(cv2.CAP_PROP_FPS)
            width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            
            if fps <= 0 or width <= 0 or height <= 0:
                return None, None, None, "Invalid video properties detected."
            
            output_path = os.path.join(tempfile.gettempdir(), f"out_{os.urandom(8).hex()}.mp4")
            # Use 'mp4v' as codec
            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
            out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
            
            if not out.isOpened():
                return None, None, None, "Video processing failed: No suitable encoder available."
            
            processed_frames = 0
            while True:
                ret, frame = cap.read()
                if not ret:
                    break
                    
                # Process frame: convert to RGB, run YOLO, then annotate and convert back to BGR.
                frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                results = model.predict(source=frame_rgb, conf=sensitivity)[0]
                annotated_frame = results.plot()
                annotated_frame_bgr = cv2.cvtColor(annotated_frame, cv2.COLOR_RGB2BGR)
                
                out.write(annotated_frame_bgr)
                processed_frames += 1
            
            cap.release()
            out.release()
            temp_files.append(output_path)
            
            if processed_frames == 0:
                return None, None, None, "No frames processed from video."
                
            if not os.path.exists(output_path) or os.path.getsize(output_path) < 1024:
                return None, None, None, f"Output video created but too small ({os.path.getsize(output_path)} bytes) - processing failed."
                
            return output_path, None, output_path, f"Video processed successfully! ({processed_frames}/{frame_count} frames)"

        else:
            # Image processing
            results = model.predict(source=input_path, conf=sensitivity)[0]
            annotated = results.plot()
            output_path = os.path.join(tempfile.gettempdir(), f"out_{os.urandom(8).hex()}.jpg")
            cv2.imwrite(output_path, annotated)
            temp_files.append(output_path)
            return output_path, output_path, None, "Image processed successfully!"

    except Exception as e:
        return None, None, None, f"Processing error: {str(e)}"
    
    finally:
        # Clean up temporary files except the final output
        for f in temp_files[:-1]:
            if f and os.path.exists(f):
                try:
                    os.remove(f)
                except:
                    pass

# When the user clicks "Process Input"
if process_button:
    out_file, out_img, out_vid, status = process_input(uploaded_file, youtube_link, image_url, sensitivity)
    st.write(status)
    if out_img:
        st.image(out_img, caption="Annotated Output (Image)", use_column_width=True)
    if out_vid:
        st.video(out_vid)
    if out_file:
        with open(out_file, "rb") as f:
            st.download_button(
                label="Download Annotated Output",
                data=f,
                file_name=os.path.basename(out_file),
                mime="video/mp4" if os.path.splitext(out_file)[1].lower() in video_exts else "image/jpeg"
            )