import os
# Disable CUDA before torch import to avoid encoding errors in container
#os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")

from typing import Dict, List, Optional

import cv2
import numpy as np
import requests
from celery import Celery

from app.model_registry import MODEL_MAP, iter_models


#REDIS_URL = os.getenv("REDIS_URL", "redis://localhost:6379/0")
REDIS_URL = os.getenv("CELERY_BROKER_URL", "redis://localhost:6379/0")

celery_app = Celery(
    "reservoir_tasks",
    broker=REDIS_URL,
    backend=REDIS_URL,
)


def check_roi(point, roi_polygon: List[List[int]]) -> bool:
    polygon = np.array(roi_polygon, dtype=np.int32)
    return cv2.pointPolygonTest(polygon, point, False) >= 0


def collect_scene_alerts(
    frame,
    roi_polygon: Optional[List[List[int]]] = None,
    model_type: Optional[str] = None,
) -> Dict[str, List[dict]]:
    scene_alerts: Dict[str, List[dict]] = {}

    for current_model_type, model in iter_models(model_type):
        results = model.predict(frame, conf=0.4, verbose=False)
        alerts = []

        for result in results:
            boxes = result.boxes
            if boxes is None:
                continue

            for box in boxes:
                cls = int(box.cls[0])
                label = model.names[cls]
                bbox = box.xyxy[0].cpu().numpy()
                center_point = ((bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2)

                is_in_roi = True
                if roi_polygon:
                    is_in_roi = check_roi(center_point, roi_polygon)

                if is_in_roi:
                    alerts.append(
                        {
                            "label": label,
                            "confidence": float(box.conf[0]),
                            "bbox": bbox.tolist(),
                        }
                    )

        if alerts:
            scene_alerts[current_model_type] = alerts

    return scene_alerts


@celery_app.task(name="analyze_video_stream", bind=True)
def analyze_video_stream(
    self,
    stream_url: str,
    webhook_url: str,
    model_type: Optional[str] = None,
    roi_polygon: Optional[List[List[int]]] = None,
):
    if model_type is not None and model_type not in MODEL_MAP:
        return {"status": "error", "message": f"unsupported model type: {model_type}"}

    cap = cv2.VideoCapture(stream_url)
    if not cap.isOpened():
        return {"status": "error", "message": f"failed to open stream: {stream_url}"}

    frame_skip = 10
    count = 0
    detected_scenes = set()

    # 获取视频信息
    fps = cap.get(cv2.CAP_PROP_FPS)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) if stream_url.endswith('.mp4') else None
    
    # 更新任务状态
    self.update_state(state='PROCESSING', meta={'progress': 0, 'message': 'Starting stream analysis'})

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

        if count % frame_skip == 0:
            scene_alerts = collect_scene_alerts(frame, roi_polygon=roi_polygon, model_type=model_type)

            for current_model_type, alerts in scene_alerts.items():
                detected_scenes.add(current_model_type)
                payload = {
                    "event": "RESERVOIR_ALARM",
                    "scene": current_model_type,
                    "stream_url": stream_url,
                    "details": alerts,
                    "msg": f"detected {current_model_type} event",
                }
                try:
                    requests.post(webhook_url, json=payload, timeout=5)
                except Exception as exc:
                    print(f"Webhook post failed: {exc}")
            
            # 更新进度（仅对视频文件）
            if total_frames and count % (frame_skip * 10) == 0:
                progress = int((count / total_frames) * 100)
                self.update_state(
                    state='PROCESSING',
                    meta={
                        'progress': progress,
                        'frame': count,
                        'total_frames': total_frames,
                        'message': f'Analyzing frame {count}/{total_frames}'
                    }
                )

        count += 1

    cap.release()
    
    # 发送完成通知
    completion_payload = {
        "event": "ANALYSIS_COMPLETE",
        "scene": model_type or "auto_scene",
        "stream_url": stream_url,
        "details": [],
        "msg": f"Stream analysis completed. Processed {count} frames"
    }
    try:
        requests.post(webhook_url, json=completion_payload, timeout=5)
    except Exception as exc:
        print(f"Completion webhook failed: {exc}")
    
    return {
        "status": "completed",
        "processed_frames": count,
        "scene": model_type or "auto_scene",
        "detected_scenes": sorted(detected_scenes),
    }


@celery_app.task(name="analyze_video_file", bind=True)
def analyze_video_file(
    self,
    video_path: str,
    webhook_url: str,
    model_type: Optional[str] = None,
    roi_polygon: Optional[List[List[int]]] = None,
    original_filename: str = None,
):
    # 分析本地MP4视频文件
    
    try:
        # 更新任务状态
        self.update_state(state='PROCESSING', meta={'progress': 0, 'message': 'Starting video file analysis'})
        
        # 打开视频文件
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            raise Exception(f"Failed to open video file: {video_path}")
        
        # 获取视频信息
        fps = cap.get(cv2.CAP_PROP_FPS)
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        
        # 采样间隔（每秒处理2-3帧）
        sample_interval = max(1, int(fps / 2))  # 每秒处理2帧
        frame_count = 0
        processed_frames = 0
        detected_scenes = set()
        
        # 检测结果去重（避免短时间内重复告警）
        last_alarm_time = {}
        alarm_cooldown = 30  # 同类型告警冷却时间（秒）
        
        import time
        
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            
            frame_count += 1
            
            # 按照采样率处理
            if frame_count % sample_interval != 0:
                continue
            
            processed_frames += 1
            progress = int((frame_count / total_frames) * 100)
            
            # 更新进度（每10%更新一次）
            if progress % 10 == 0 and progress != self.request.meta.get('progress', 0):
                self.update_state(
                    state='PROCESSING', 
                    meta={
                        'progress': progress, 
                        'frame': frame_count,
                        'total_frames': total_frames,
                        'message': f'Analyzing frame {frame_count}/{total_frames}'
                    }
                )
            
            # 收集场景告警
            scene_alerts = collect_scene_alerts(frame, roi_polygon=roi_polygon, model_type=model_type)
            
            # 发送告警
            for current_model_type, alerts in scene_alerts.items():
                detected_scenes.add(current_model_type)
                current_time = time.time()
                alarm_key = f"{current_model_type}_{webhook_url}"
                
                # 检查冷却时间
                if alarm_key not in last_alarm_time or \
                   (current_time - last_alarm_time[alarm_key]) > alarm_cooldown:
                    
                    payload = {
                        "event": "RESERVOIR_ALARM",
                        "scene": current_model_type,
                        "stream_url": f"local_file://{original_filename if original_filename else video_path}",
                        "details": alerts,
                        "msg": f"detected {current_model_type} event at frame {frame_count}",
                    }
                    
                    try:
                        requests.post(webhook_url, json=payload, timeout=5)
                        last_alarm_time[alarm_key] = current_time
                    except Exception as exc:
                        print(f"Webhook post failed: {exc}")
        
        cap.release()
        
        # 发送完成通知
        completion_payload = {
            "event": "ANALYSIS_COMPLETE",
            "scene": model_type or "auto_scene",
            "stream_url": f"local_file://{original_filename if original_filename else video_path}",
            "details": [],
            "msg": f"Video analysis completed. Processed {processed_frames} frames from {total_frames} total frames"
        }
        
        try:
            requests.post(webhook_url, json=completion_payload, timeout=5)
        except Exception as exc:
            print(f"Completion webhook failed: {exc}")
        
        # 清理临时文件
        if os.path.exists(video_path):
            os.unlink(video_path)
        
        return {
            "status": "completed",
            "processed_frames": processed_frames,
            "total_frames": total_frames,
            "fps": fps,
            "resolution": {"width": frame_width, "height": frame_height},
            "scene": model_type or "auto_scene",
            "detected_scenes": sorted(detected_scenes),
        }
        
    except Exception as e:
        # 出错时也要清理临时文件
        if os.path.exists(video_path):
            try:
                os.unlink(video_path)
            except:
                pass
        raise e