import os
import cv2
import json
import requests
from celery import Celery
from ultralytics import YOLO
from shapely.geometry import Point, Polygon

# 初始化 Celery
redis_url = os.getenv("REDIS_URL", "redis://reservoir_ai-redis-ai-1:6379/0")
app = Celery("reservoir_tasks", broker=redis_url, backend=redis_url)

# 加载模型 - 使用绝对路径确保容器内加载成功
model_path = os.path.join(os.path.dirname(__file__), 'yolov8n.pt')
model = YOLO(model_path)

# 水库场景类别映射 (COCO数据集索引)
# 0: person (非法游泳/捕鱼)
# 8: boat (非法船只)
# 14: bird (水面漂浮物初筛)
# 15: cat, 16: dog (岸边动物干扰排除)
RESERVOIR_CLASSES = [0, 8, 14] 

@app.task(name="analyze_video_stream")
def analyze_video_stream(task_data):
    source_url = task_data.get("source_url")
    webhook_url = task_data.get("webhook_url")
    scene_type = task_data.get("scene_type")
    roi_points = task_data.get("roi", [])

    # 建立电子围栏多边形
    roi_polygon = None
    if len(roi_points) >= 3:
        roi_polygon = Polygon(roi_points)

    cap = cv2.VideoCapture(source_url)
    if not cap.isOpened():
        return {"status": "error", "message": f"无法打开视频流: {source_url}"}

    frame_count = 0
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break

        # 抽帧处理：每5帧处理一次，降低CPU压力
        frame_count += 1
        if frame_count % 5 != 0:
            continue

        # 推理推理：设置置信度0.4，只看特定类别
        results = model.track(
            frame,
            persist=True,
            conf=0.4,
            classes=RESERVOIR_CLASSES,
            verbose=False
        )

        for r in results:
            if r.boxes:
                boxes = r.boxes
                for box in boxes:
                    # 获取坐标
                    x1, y1, x2, y2 = box.xyxy[0].tolist()
                    conf = float(box.conf[0])
                    cls = int(box.cls[0])
                    label = model.names[cls]

                    # 计算目标中心点
                    center_point = Point((x1 + x2) / 2, (y1 + y2) / 2)

                    # 电子围栏判定：如果划定了ROI，则目标中心必须在ROI内
                    is_in_roi = True
                    if roi_polygon:
                        is_in_roi = roi_polygon.contains(center_point)

                    if is_in_roi:
                        # 构造告警信息
                        payload = {
                            "event": "illegal_activity_detected",
                            "scene": scene_type,
                            "label": label,
                            "confidence": round(conf, 2),
                            "bbox": [x1, y1, x2, y2],
                            "timestamp": frame_count
                        }

                        # 发送实时告警
                        try:
                            requests.post(webhook_url, json=payload, timeout=2)
                        except Exception as e:
                            print(f"Webhook推送失败: {e}")

    cap.release()
    return {"status": "completed", "processed_frames": frame_count}