#!/opt/homebrew/bin/python3.11
"""
整合 Face + ASRX 說話人分離（版本 3 - 修復 face_detected 檢查）
"""

import json
import argparse
from pathlib import Path
from typing import Dict, List


def load_json(path: str):
    """載入 JSON 文件"""
    with open(path, 'r', encoding='utf-8') as f:
        return json.load(f)


def match_face_with_speaker_v3(face_data: Dict, asrx_data: Dict, 
                                time_threshold: float = 3.0) -> List[Dict]:
    """
    匹配人臉與說話人（版本 3 - 修復版）
    
    修復：Face 數據沒有 face_detected 欄位，改用 faces 列表是否為空判斷
    """
    face_frames = face_data.get('frames', [])
    asrx_segments = asrx_data.get('segments', [])
    
    # 將 Face 幀按時間排序
    face_frames_sorted = sorted(face_frames, key=lambda x: x.get('timestamp', 0))
    
    print(f"  Face frames: {len(face_frames_sorted)}")
    print(f"  ASRX segments: {len(asrx_segments)}")
    
    # 匹配
    integrated = []
    
    for i, seg in enumerate(asrx_segments):
        start = seg['start']
        end = seg['end']
        speaker = seg['speaker']
        mid_time = (start + end) / 2
        
        # 找到時間範圍內的人臉
        faces_in_range = []
        for frame in face_frames_sorted:
            ts = frame.get('timestamp', 0)
            
            # 檢查是否在時間範圍內
            if start - time_threshold <= ts <= end + time_threshold:
                # 檢查是否有人臉（faces 列表不為空）
                faces = frame.get('faces', [])
                if faces and len(faces) > 0:
                    faces_in_range.append({
                        'timestamp': ts,
                        'faces': faces,
                        'distance_from_mid': abs(ts - mid_time)
                    })
        
        # 選擇最接近片段中間的人臉
        if faces_in_range:
            faces_in_range.sort(key=lambda x: x['distance_from_mid'])
            best_face = faces_in_range[0]
        else:
            best_face = None
        
        # 建立整合結果
        integrated.append({
            'start': start,
            'end': end,
            'duration': seg.get('duration', end - start),
            'speaker': speaker,
            'has_face': best_face is not None,
            'face_timestamp': best_face['timestamp'] if best_face else None,
            'face_location': best_face['faces'][0] if best_face and best_face['faces'] else None,
            'face_count_in_range': len(faces_in_range)
        })
        
        # 進度顯示
        if (i + 1) % 200 == 0:
            print(f"  Processed {i+1}/{len(asrx_segments)} segments...")
    
    return integrated


def analyze_speaker_face(integrated: List[Dict]):
    """分析說話人與人臉的對應"""
    speaker_stats = {}
    
    for item in integrated:
        speaker = item['speaker']
        if speaker not in speaker_stats:
            speaker_stats[speaker] = {
                'total_segments': 0,
                'with_face': 0,
                'without_face': 0,
                'total_duration': 0
            }
        
        speaker_stats[speaker]['total_segments'] += 1
        speaker_stats[speaker]['total_duration'] += item['duration']
        
        if item['has_face']:
            speaker_stats[speaker]['with_face'] += 1
        else:
            speaker_stats[speaker]['without_face'] += 1
    
    return speaker_stats


def main():
    parser = argparse.ArgumentParser(description='整合 Face + ASRX 說話人')
    parser.add_argument('face_json', help='Face 檢測結果 JSON')
    parser.add_argument('asrx_json', help='ASRX 說話人分離 JSON')
    parser.add_argument('-o', '--output', help='輸出整合結果 JSON')
    parser.add_argument('--threshold', type=float, default=3.0,
                        help='時間閾值（秒）')
    parser.add_argument('--stats', action='store_true', help='只显示統計')
    
    args = parser.parse_args()
    
    # 載入數據
    print(f"[Load] Face: {args.face_json}")
    face_data = load_json(args.face_json)
    
    print(f"[Load] ASRX: {args.asrx_json}")
    asrx_data = load_json(args.asrx_json)
    
    # 匹配
    print(f"\n[Match] Matching faces with speakers (threshold={args.threshold}s)...")
    integrated = match_face_with_speaker_v3(face_data, asrx_data, args.threshold)
    
    # 分析
    print(f"\n[Analyze] Analyzing speaker-face correspondence...")
    speaker_stats = analyze_speaker_face(integrated)
    
    # 顯示統計
    print(f"\n{'='*70}")
    print(f"說話人 - 人臉對應統計")
    print(f"{'='*70}")
    
    total_segments = len(integrated)
    total_with_face = sum(1 for item in integrated if item['has_face'])
    
    for speaker, stats in sorted(speaker_stats.items()):
        with_face_pct = stats['with_face'] / stats['total_segments'] * 100 if stats['total_segments'] > 0 else 0
        print(f"\n🔊 {speaker}:")
        print(f"  總片段：{stats['total_segments']}")
        print(f"  有人臉：{stats['with_face']} ({with_face_pct:.1f}%)")
        print(f"  無人臉：{stats['without_face']}")
        print(f"  總時長：{stats['total_duration']:.1f}s ({stats['total_duration']/60:.1f}分鐘)")
    
    print(f"\n{'='*70}")
    print(f"總計：{total_segments} 片段，{total_with_face} 片段有人臉 ({total_with_face/total_segments*100:.1f}%)")
    print(f"{'='*70}")
    
    # 保存結果
    if args.output:
        output_path = Path(args.output)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        
        result = {
            'face_source': str(args.face_json),
            'asrx_source': str(args.asrx_json),
            'time_threshold': args.threshold,
            'integrated_segments': integrated,
            'speaker_stats': speaker_stats
        }
        
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(result, f, indent=2, ensure_ascii=False)
        
        print(f"\n[Save] Results saved to: {output_path}")
    
    return integrated, speaker_stats


if __name__ == "__main__":
    main()
