#!/usr/bin/env python3
"""
視覺分片處理器 (Phase 2.2)

從 YOLO 結果生成視覺分片，支持多種分片策略：
1. 固定幀數分片
2. 基於物件相似度分片
3. 基於場景變化分片
"""

import json
import sys
import os
import argparse
from pathlib import Path
from typing import Dict, List, Any, Optional
import numpy as np
from datetime import datetime

# 添加父目錄到路徑以導入其他模組
sys.path.insert(0, str(Path(__file__).parent.parent))
from scripts.yolo_processor_contract_v1 import YOLOProcessor


class VisualChunkProcessor:
    """視覺分片處理器"""

    def __init__(self, video_path: str, yolo_result_path: Optional[str] = None):
        self.video_path = video_path
        self.yolo_result_path = yolo_result_path
        self.yolo_result = None

    def load_yolo_result(self):
        """加載 YOLO 結果"""
        if self.yolo_result_path and os.path.exists(self.yolo_result_path):
            with open(self.yolo_result_path, "r", encoding="utf-8") as f:
                self.yolo_result = json.load(f)
        else:
            # 如果沒有提供 YOLO 結果路徑，則運行 YOLO 檢測
            print(f"[VisualChunk] Running YOLO detection for: {self.video_path}")
            yolo_processor = YOLOProcessor(self.video_path)
            yolo_result = yolo_processor.process()
            self.yolo_result = yolo_processor.to_json_dict()

    def create_fixed_frame_chunks(
        self, frames_per_chunk: int = 30
    ) -> List[Dict[str, Any]]:
        """創建固定幀數分片

        Args:
            frames_per_chunk: 每個分片的幀數

        Returns:
            視覺分片列表
        """
        if not self.yolo_result:
            self.load_yolo_result()

        frames = self.yolo_result.get("frames", {})
        if not frames:
            return []

        # 將幀字典轉換為排序後的列表
        frame_list = []
        for frame_key, frame_data in frames.items():
            frame_list.append(
                {
                    "frame_number": int(frame_key),
                    "timestamp": frame_data.get("time_seconds", 0),
                    "objects": frame_data.get("detections", []),
                }
            )

        # 按幀號排序
        frame_list.sort(key=lambda x: x["frame_number"])

        chunks = []
        total_frames = len(frame_list)

        for start_idx in range(0, total_frames, frames_per_chunk):
            end_idx = min(start_idx + frames_per_chunk, total_frames)
            chunk_frames = frame_list[start_idx:end_idx]

            if not chunk_frames:
                continue

            # 計算分片統計
            chunk_stats = self._calculate_chunk_stats(chunk_frames)

            chunk = {
                "start_frame": chunk_frames[0]["frame_number"],
                "end_frame": chunk_frames[-1]["frame_number"] + 1,  # exclusive
                "frame_count": len(chunk_frames),
                "keyframe_objects": self._extract_keyframe_objects(chunk_frames),
                "dominant_objects": chunk_stats["dominant_objects"],
                "metadata": {
                    "object_count": chunk_stats["total_objects"],
                    "unique_classes": chunk_stats["unique_classes"],
                    "max_confidence": chunk_stats["max_confidence"],
                    "avg_confidence": chunk_stats["avg_confidence"],
                    "spatial_density": chunk_stats["spatial_density"],
                },
            }

            chunks.append(chunk)

        return chunks

    def create_similarity_based_chunks(
        self, similarity_threshold: float = 0.5, min_frames_per_chunk: int = 10
    ) -> List[Dict[str, Any]]:
        """基於物件相似度創建分片

        Args:
            similarity_threshold: 相似度閾值 (0-1)
            min_frames_per_chunk: 最小幀數

        Returns:
            視覺分片列表
        """
        if not self.yolo_result:
            self.load_yolo_result()

        frames = self.yolo_result.get("frames", {})
        if not frames:
            return []

        # 將幀字典轉換為排序後的列表
        frame_list = []
        for frame_key, frame_data in frames.items():
            frame_list.append(
                {
                    "frame_number": int(frame_key),
                    "timestamp": frame_data.get("time_seconds", 0),
                    "objects": frame_data.get("detections", []),
                }
            )

        # 按幀號排序
        frame_list.sort(key=lambda x: x["frame_number"])

        chunks = []
        current_chunk_frames = []
        current_start_frame = 0

        for i, frame in enumerate(frame_list):
            if not current_chunk_frames:
                current_chunk_frames.append(frame)
                current_start_frame = frame["frame_number"]
                continue

            # 計算相似度
            last_frame = current_chunk_frames[-1]
            similarity = self._calculate_frame_similarity(last_frame, frame)

            if similarity >= similarity_threshold:
                # 相似度高，加入當前分片
                current_chunk_frames.append(frame)
            else:
                # 相似度低，創建新分片
                if len(current_chunk_frames) >= min_frames_per_chunk:
                    chunk = self._create_chunk_from_frames(
                        current_chunk_frames,
                        current_start_frame,
                        frame_list[i - 1]["frame_number"] + 1,
                    )
                    chunks.append(chunk)

                # 開始新的分片
                current_chunk_frames = [frame]
                current_start_frame = frame["frame_number"]

        # 處理最後一個分片
        if len(current_chunk_frames) >= min_frames_per_chunk:
            chunk = self._create_chunk_from_frames(
                current_chunk_frames,
                current_start_frame,
                current_chunk_frames[-1]["frame_number"] + 1,
            )
            chunks.append(chunk)

        return chunks

    def _calculate_frame_similarity(self, frame1: Dict, frame2: Dict) -> float:
        """計算兩個幀之間的相似度（基於物件類別）"""
        objects1 = frame1.get("objects", [])
        objects2 = frame2.get("objects", [])

        if not objects1 and not objects2:
            return 1.0

        if not objects1 or not objects2:
            return 0.0

        # 提取物件類別
        classes1 = set(
            obj.get("class_name", "") for obj in objects1 if obj.get("class_name")
        )
        classes2 = set(
            obj.get("class_name", "") for obj in objects2 if obj.get("class_name")
        )

        # 計算 Jaccard 相似度
        intersection = classes1.intersection(classes2)
        union = classes1.union(classes2)

        if not union:
            return 0.0

        return len(intersection) / len(union)

    def _calculate_chunk_stats(self, frames: List[Dict]) -> Dict[str, Any]:
        """計算分片統計信息"""
        all_objects = []
        for frame in frames:
            all_objects.extend(frame.get("objects", []))

        # 總物件數
        total_objects = len(all_objects)

        # 唯一類別
        unique_classes = list(
            set(
                obj.get("class_name", "")
                for obj in all_objects
                if obj.get("class_name")
            )
        )

        # 信心值統計
        confidences = [obj.get("confidence", 0) for obj in all_objects]
        max_confidence = max(confidences) if confidences else 0
        avg_confidence = np.mean(confidences) if confidences else 0

        # 空間密度（每幀平均物件數）
        spatial_density = total_objects / len(frames) if frames else 0

        # 主要物件（出現在大多數幀中的物件）
        object_counts = {}
        for frame in frames:
            frame_classes = set(
                obj.get("class_name", "")
                for obj in frame.get("objects", [])
                if obj.get("class_name")
            )
            for class_name in frame_classes:
                object_counts[class_name] = object_counts.get(class_name, 0) + 1

        dominant_objects = [
            class_name
            for class_name, count in object_counts.items()
            if count / len(frames) > 0.5
        ]
        dominant_objects.sort()

        return {
            "total_objects": total_objects,
            "unique_classes": unique_classes,
            "max_confidence": float(max_confidence),
            "avg_confidence": float(avg_confidence),
            "spatial_density": float(spatial_density),
            "dominant_objects": dominant_objects,
        }

    def _extract_keyframe_objects(self, frames: List[Dict]) -> List[Dict[str, Any]]:
        """提取關鍵幀物件"""
        keyframe_objects = []

        # 簡化：每5幀取一個關鍵幀
        for i in range(0, len(frames), 5):
            if i < len(frames):
                frame = frames[i]
                objects = []

                for obj in frame.get("objects", []):
                    objects.append(
                        {
                            "class_name": obj.get("class_name", ""),
                            "class_id": obj.get("class_id", 0),
                            "confidence": float(obj.get("confidence", 0)),
                            "bbox": {
                                "x": obj.get("x1", 0),
                                "y": obj.get("y1", 0),
                                "width": obj.get("width", 0),
                                "height": obj.get("height", 0),
                            }
                            if "x1" in obj
                            else None,
                            "occurrence": 1,
                        }
                    )

                keyframe_objects.append(
                    {
                        "timestamp": float(frame.get("timestamp", 0)),
                        "frame_number": frame.get("frame_number", 0),
                        "objects": objects,
                    }
                )

        return keyframe_objects

    def _create_chunk_from_frames(
        self, frames: List[Dict], start_frame: int, end_frame: int
    ) -> Dict[str, Any]:
        """從幀列表創建分片"""
        chunk_stats = self._calculate_chunk_stats(frames)

        return {
            "start_frame": start_frame,
            "end_frame": end_frame,  # exclusive
            "frame_count": len(frames),
            "keyframe_objects": self._extract_keyframe_objects(frames),
            "dominant_objects": chunk_stats["dominant_objects"],
            "object_relationships": [],  # 可選：後期添加關係檢測
            "scene_description": None,  # 可選：後期添加 LLM 生成的場景描述
            "metadata": {
                "object_count": chunk_stats["total_objects"],
                "unique_classes": chunk_stats["unique_classes"],
                "max_confidence": chunk_stats["max_confidence"],
                "avg_confidence": chunk_stats["avg_confidence"],
                "spatial_density": chunk_stats["spatial_density"],
            },
        }

    def process(self, strategy: str = "fixed", **kwargs) -> Dict[str, Any]:
        """處理視覺分片生成

        Args:
            strategy: 分片策略 ("fixed" 或 "similarity")
            **kwargs: 策略參數

        Returns:
            處理結果
        """
        if not self.yolo_result:
            self.load_yolo_result()

        start_time = datetime.now()

        if strategy == "fixed":
            frames_per_chunk = kwargs.get("frames_per_chunk", 30)
            chunks = self.create_fixed_frame_chunks(frames_per_chunk)
        elif strategy == "similarity":
            similarity_threshold = kwargs.get("similarity_threshold", 0.5)
            min_frames = kwargs.get("min_frames_per_chunk", 10)
            chunks = self.create_similarity_based_chunks(
                similarity_threshold, min_frames
            )
        else:
            raise ValueError(f"Unknown strategy: {strategy}")

        # 計算總統計
        total_frames = sum(chunk["frame_count"] for chunk in chunks)
        total_objects = sum(chunk["metadata"]["object_count"] for chunk in chunks)

        # 收集所有唯一類別
        all_unique_classes = set()
        for chunk in chunks:
            all_unique_classes.update(chunk["metadata"]["unique_classes"])

        processing_time = (datetime.now() - start_time).total_seconds()

        result = {
            "metadata": {
                "video_path": self.video_path,
                "processing_time": processing_time,
                "strategy": strategy,
                "parameters": kwargs,
                "processed_at": datetime.now().isoformat(),
            },
            "chunk_count": len(chunks),
            "total_frames": total_frames,
            "total_objects": total_objects,
            "unique_classes": len(all_unique_classes),
            "chunks": chunks,
        }

        return result


def main():
    parser = argparse.ArgumentParser(description="視覺分片處理器")
    parser.add_argument("video_path", help="視頻文件路徑")
    parser.add_argument("output_path", help="輸出文件路徑")
    parser.add_argument("--yolo-result", help="YOLO 結果文件路徑（可選）")
    parser.add_argument("--uuid", help="檔案 UUID（由 executor 傳入）")
    parser.add_argument(
        "--strategy", choices=["fixed", "similarity"], default="fixed", help="分片策略"
    )
    parser.add_argument(
        "--frames-per-chunk", type=int, default=30, help="固定幀數策略：每個分片的幀數"
    )
    parser.add_argument(
        "--similarity-threshold", type=float, default=0.5, help="相似度策略：相似度閾值"
    )
    parser.add_argument(
        "--min-frames-per-chunk", type=int, default=10, help="相似度策略：最小幀數"
    )

    args = parser.parse_args()

    print(f"[VisualChunk] Starting processing: {args.video_path}")
    print(f"[VisualChunk] Strategy: {args.strategy}")

    try:
        processor = VisualChunkProcessor(args.video_path, args.yolo_result)

        if args.strategy == "fixed":
            result = processor.process(
                strategy="fixed", frames_per_chunk=args.frames_per_chunk
            )
        else:
            result = processor.process(
                strategy="similarity",
                similarity_threshold=args.similarity_threshold,
                min_frames_per_chunk=args.min_frames_per_chunk,
            )

        # 保存結果
        with open(args.output_path, "w", encoding="utf-8") as f:
            json.dump(result, f, ensure_ascii=False, indent=2)

        print("[VisualChunk] Processing completed")
        print(f"[VisualChunk] Generated {result['chunk_count']} visual chunks")
        print(f"[VisualChunk] Total frames: {result['total_frames']}")
        print(f"[VisualChunk] Total objects: {result['total_objects']}")
        print(f"[VisualChunk] Unique classes: {result['unique_classes']}")
        print(f"[VisualChunk] Result saved to: {args.output_path}")
    except Exception as e:
        print(f"[VisualChunk] Error: {e}", file=sys.stderr)
        result = {
            "chunk_count": 0,
            "total_frames": 0,
            "total_objects": 0,
            "unique_classes": 0,
            "chunks": [],
        }
        with open(args.output_path, "w", encoding="utf-8") as f:
            json.dump(result, f, ensure_ascii=False, indent=2)
        print(f"[VisualChunk] Fallback: empty result saved to {args.output_path}")


if __name__ == "__main__":
    main()