import cv2
import gradio as gr
import numpy as np
import spaces
from ultralytics import YOLO

# تحميل نموذج YOLOv8
model = YOLO("yolov8n.pt")

@spaces.GPU  # اتركه إذا كنت في Hugging Face، أو احذفه إذا كنت في Colab
def process_gate_video(video_path):
    if video_path is None:
        return None, "الرجاء رفع ملف فيديو صالح أولاً."

    cap = cv2.VideoCapture(video_path)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    if fps <= 0:
        fps = 25

    output_path = "output_counted_gate.avi"
    fourcc = cv2.VideoWriter_fourcc(*'XVID')
    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

    people_count = 0
    tracked_ids = set()
    line_y = int(height * 0.5)

    while cap.isOpened():
        success, frame = cap.read()
        if not success:
            break

        results = model.track(frame, persist=True, classes=[0], verbose=False)

        if results[0].boxes.id is not None:
            boxes = results[0].boxes.xyxy.cpu().numpy()
            track_ids = results[0].boxes.id.int().cpu().numpy()

            for box, track_id in zip(boxes, track_ids):
                x1, y1, x2, y2 = map(int, box)
                center_y = int(y2)

                cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                cv2.putText(frame, f"ID: {track_id}", (x1, y1 - 10), 
                            cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)

                if line_y - 25 < center_y < line_y + 25:
                    if track_id not in tracked_ids:
                        people_count += 1
                        tracked_ids.add(track_id)

        cv2.line(frame, (0, line_y), (width, line_y), (0, 0, 255), 3)
        cv2.putText(frame, f"Total Count: {people_count}", (30, 60), 
                    cv2.FONT_HERSHEY_SIMPLEX, 1.3, (0, 0, 255), 3)

        out.write(frame)

    cap.release()
    out.release()
    
    result_message = f"تم الانتهاء من المعالجة بنجاح. إجمالي الأشخاص العابرين: {people_count}"
    return output_path, result_message

# بناء واجهة Gradio
with gr.Blocks() as demo:
    gr.Markdown("# 🚪 نظام عد الأشخاص الذكي عبر البوابات")
    
    with gr.Row():
        with gr.Column():
            video_input = gr.Video(label="فيديو البوابة المدخل")
            submit_btn = gr.Button("بدء العد والمعالجة", variant="primary")
        
        with gr.Column():
            video_output = gr.Video(label="الفيديو بعد المعالجة والعد")
            text_output = gr.Textbox(label="ملخص النتيجة")

    submit_btn.click(
        fn=process_gate_video,
        inputs=video_input,
        outputs=[video_output, text_output]
    )

# تم إزالة debug=True و share=True لتجنب أي تكرار أو ومضات في المتصفح
if __name__ == "__main__":
    demo.launch()