#!/usr/bin/env python3
"""
extract_face_crops.py - 批量提取 face crops

Usage:
    python3 scripts/extract_face_crops.py --uuid <file_uuid>
    python3 scripts/extract_face_crops.py --uuid <file_uuid> --video <video_path>

儲存位置: {OUTPUT_DIR}/.faces/{file_uuid}/{trace_id}/{frame}.jpg

條件:
    - trace_id != None and trace_id != 0
    - landmarks.left_eye or landmarks.right_eye

品檢:
    - file_size > 500 bytes
    - mean_brightness > 5
    - std_deviation > 10

Retry: 最多 3 次
"""

import argparse
import json
import subprocess
import os
import sys
from pathlib import Path
from datetime import datetime
from typing import Dict, List, Optional, Tuple, Set
from concurrent.futures import ThreadPoolExecutor, as_completed
import threading

# Constants
MAX_RETRIES = 3
MIN_FILE_SIZE = 500
MIN_BRIGHTNESS = 5
MIN_STD_DEV = 10
FFMPEG_TIMEOUT = 30
MAX_WORKERS = 8  # Parallel threads for ffmpeg


class FaceCropExtractor:
    def __init__(self, output_dir: str):
        self.output_dir = Path(output_dir)
        self.faces_dir = self.output_dir / ".faces"
        self.faces_dir.mkdir(parents=True, exist_ok=True)
        self.stats = {"total_faces": 0, "qualified": 0, "successful": 0, "failed": 0, "skipped": 0, "low_confidence": 0, "too_small": 0}
        self.stats_lock = threading.Lock()
    
    def process_video(self, uuid: str, video_path: str) -> dict:
        """處理單一影片"""
        face_json = self.output_dir / f"{uuid}.face.json"
        traced_json = self.output_dir / f"{uuid}.face_traced.json"
        
        if not face_json.exists():
            print(f"[ERROR] face.json not found: {uuid}")
            return {"error": "face.json not found"}
        
        if not os.path.exists(video_path):
            print(f"[ERROR] Video not found: {video_path}")
            return {"error": "video not found"}
        
        # Load face.json (landmarks)
        print(f"[LOAD] Reading {face_json}")
        with open(face_json) as f:
            face_data = json.load(f)
        
        # Load face_traced.json if exists (trace_id)
        traced_data = {}
        if traced_json.exists():
            print(f"[LOAD] Reading {traced_json}")
            with open(traced_json) as f:
                traced_data = json.load(f)
        
        # Build lookup: (frame, x, y) -> trace_id from traced_data
        trace_lookup: Dict[Tuple[int, int, int], int] = {}
        frames = traced_data.get("frames", {})
        if isinstance(frames, dict):
            for fnum, frm in frames.items():
                faces = frm.get("faces", [])
                if faces is None:
                    continue
                for face in faces:
                    if face is None:
                        continue
                    trace_id = face.get("trace_id")
                    if trace_id and trace_id != 0:
                        x = face.get("x", 0)
                        y = face.get("y", 0)
                        key = (int(fnum), x, y)
                        trace_lookup[key] = trace_id
        
        # Create output directory
        uuid_dir = self.faces_dir / uuid
        uuid_dir.mkdir(parents=True, exist_ok=True)
        
        results = {"successful": [], "failed": []}
        processed: Set[Tuple[int, int]] = set()  # (trace_id, frame)
        trace_counts: Dict[int, int] = {}  # trace_id -> count
        
        # Process faces from face.json
        frames = face_data.get("frames", {})
        if isinstance(frames, dict):
            frame_items = frames.items()
        elif isinstance(frames, list):
            frame_items = [(frm.get("frame"), frm) for frm in frames]
        else:
            frame_items = []
        
        # Collect extraction tasks
        tasks = []
        for fnum, frm in frame_items:
            if fnum is None:
                continue
            
            faces = frm.get("faces", [])
            if faces is None:
                continue
            
            for face in faces:
                if face is None:
                    continue
                
                self.stats["total_faces"] += 1
                
                bb = face.get("bbox", face)
                x = bb.get("x", 0) if isinstance(bb, dict) else 0
                y = bb.get("y", 0) if isinstance(bb, dict) else 0
                w = bb.get("width", 0) if isinstance(bb, dict) else 0
                h = bb.get("height", 0) if isinstance(bb, dict) else 0
                confidence = face.get("confidence", 0.0)
                
                # Quality filtering: confidence + size
                MIN_CONFIDENCE = 0.6
                MIN_SIZE = 20
                
                if confidence < MIN_CONFIDENCE:
                    self.stats["low_confidence"] += 1
                    continue
                if w < MIN_SIZE or h < MIN_SIZE:
                    self.stats["too_small"] += 1
                    continue
                
                # Check: has eyes
                lm = face.get("landmarks")
                has_eyes = False
                if lm:
                    if isinstance(lm, dict):
                        has_eyes = lm.get("left_eye") or lm.get("right_eye")
                    elif isinstance(lm, list) and len(lm) >= 2:
                        has_eyes = True
                
                if not has_eyes:
                    self.stats["skipped"] += 1
                    continue
                
                self.stats["qualified"] += 1
                
                # Check: in trace
                key = (int(fnum), x, y)
                trace_id = trace_lookup.get(key) or face.get("trace_id")
                
                # Determine output directory
                if trace_id and trace_id != 0:
                    output_dir = uuid_dir / str(trace_id)
                    crop_key = (trace_id, int(fnum))
                else:
                    # No trace_id → unbound directory
                    output_dir = uuid_dir / "unbound"
                    crop_key = ("unbound", int(fnum), x, y)  # unique key for unbound
                
                if crop_key in processed:
                    continue
                processed.add(crop_key)
                
                output_dir.mkdir(parents=True, exist_ok=True)
                output_path = output_dir / f"{fnum}.jpg"
                
                tasks.append({
                    "trace_id": trace_id or "unbound",
                    "frame": int(fnum),
                    "x": x, "y": y, "w": w, "h": h,
                    "output_path": output_path
                })
        
        # Parallel extraction
        print(f"[EXTRACT] Processing {len(tasks)} faces with {MAX_WORKERS} threads...")
        
        with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
            futures = {
                executor.submit(
                    self.extract_face, video_path, t["frame"], 
                    t["x"], t["y"], t["w"], t["h"], t["output_path"]
                ): t for t in tasks
            }
            
            for i, future in enumerate(as_completed(futures)):
                t = futures[future]
                result = future.result()
                
                if result["success"]:
                    self.stats["successful"] += 1
                    results["successful"].append({
                        "trace_id": t["trace_id"],
                        "frame": t["frame"],
                        "path": str(t["output_path"])
                    })
                    trace_counts[t["trace_id"]] = trace_counts.get(t["trace_id"], 0) + 1
                else:
                    self.stats["failed"] += 1
                    results["failed"].append({
                        "trace_id": t["trace_id"],
                        "frame": t["frame"],
                        "bbox": {"x": t["x"], "y": t["y"], "w": t["w"], "h": t["h"]},
                        "reason": result.get("reason", "unknown")
                    })
                
                # Progress every 1000
                if (i + 1) % 1000 == 0:
                    print(f"  Progress: {i+1}/{len(tasks)} ({self.stats['successful']} OK, {self.stats['failed']} fail)")
        
        # Write summary
        self.write_summary(uuid, trace_counts, results)
        
        return results
    
    def extract_face(self, video_path: str, frame: int, x: int, y: int, 
                     w: int, h: int, output_path: Path) -> dict:
        """提取 face crop（含 retry，使用 -ss 快速 seek）"""
        for attempt in range(MAX_RETRIES):
            try:
                ts = frame / 24.0  # FPS is always 24 for this video
                cmd = [
                    "ffmpeg", "-y", "-ss", f"{ts:.3f}",
                    "-i", video_path,
                    "-vf", f"crop={w}:{h}:{x}:{y}",
                    "-frames:v", "1",
                    "-q:v", "2",  # 高品質 JPEG
                    str(output_path)
                ]
                
                proc = subprocess.run(
                    cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                    timeout=FFMPEG_TIMEOUT
                )
                
                if proc.returncode != 0:
                    if attempt < MAX_RETRIES - 1:
                        continue
                    return {"success": False, "reason": "ffmpeg_failed"}
                
                # Quality check
                quality = self.check_quality(output_path)
                if quality["ok"]:
                    return {"success": True, "path": str(output_path)}
                
                # Quality failed, retry
                if attempt < MAX_RETRIES - 1:
                    # Remove bad file
                    if output_path.exists():
                        output_path.unlink()
                    continue
                
                return {"success": False, "reason": quality.get("reason", "quality_failed")}
            
            except subprocess.TimeoutExpired:
                if attempt < MAX_RETRIES - 1:
                    continue
                return {"success": False, "reason": "timeout"}
            except Exception as e:
                return {"success": False, "reason": str(e)}
        
        return {"success": False, "reason": "max_retries"}
    
    def check_quality(self, path: Path) -> dict:
        """品檢"""
        if not path.exists():
            return {"ok": False, "reason": "file_not_exist"}
        
        file_size = path.stat().st_size
        if file_size < MIN_FILE_SIZE:
            return {"ok": False, "reason": f"empty_file ({file_size}B)"}
        
        try:
            from PIL import Image
            import numpy as np
            
            img = Image.open(path)
            arr = np.array(img.convert('RGB'))
            
            mean_brightness = arr.mean()
            if mean_brightness < MIN_BRIGHTNESS:
                return {"ok": False, "reason": f"black_frame (mean={mean_brightness:.1f})"}
            
            std_dev = arr.std()
            if std_dev < MIN_STD_DEV:
                return {"ok": False, "reason": f"low_contrast (std={std_dev:.1f})"}
            
            return {"ok": True}
        
        except ImportError:
            # PIL not available, skip advanced quality check
            return {"ok": True}
        except Exception as e:
            return {"ok": False, "reason": str(e)}
    
    def write_summary(self, uuid: str, trace_counts: Dict[int, int], results: dict):
        """寫摘要報告"""
        summary_path = self.faces_dir / uuid / "_summary.json"
        
        summary = {
            "file_uuid": uuid,
            "timestamp": datetime.now().isoformat(),
            "stats": self.stats,
            "trace_counts": trace_counts,
            "total_traces": len(trace_counts),
            "failed_count": len(results["failed"]),
            "failed_faces": results["failed"] if results["failed"] else None
        }
        
        with open(summary_path, "w") as f:
            json.dump(summary, f, indent=2, ensure_ascii=False)
        
        print(f"\n[SUMMARY] Written to {summary_path}")
    
    def print_stats(self):
        """印統計"""
        print(f"\n=== Statistics ===")
        print(f"Total faces scanned: {self.stats['total_faces']}")
        print(f"Filtered (low confidence < 0.6): {self.stats['low_confidence']}")
        print(f"Filtered (too small < 20px): {self.stats['too_small']}")
        print(f"Qualified (trace_id + eyes): {self.stats['qualified']}")
        print(f"Successfully extracted: {self.stats['successful']}")
        print(f"Failed: {self.stats['failed']}")
        print(f"Skipped (no trace/eyes): {self.stats['skipped']}")


def main():
    parser = argparse.ArgumentParser(
        description="Extract face crops from videos",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__
    )
    
    parser.add_argument("--uuid", type=str, required=True,
                        help="File UUID to process")
    parser.add_argument("--video", type=str,
                        help="Video file path (optional, will check DB if not provided)")
    parser.add_argument("--output-dir", type=str,
                        default="/Users/accusys/momentry/output_dev",
                        help="Output directory (default: output_dev)")
    
    args = parser.parse_args()
    
    # Get video path
    video_path = args.video
    if not video_path:
        # Query from DB
        video_path = query_video_path_from_db(args.uuid)
        if not video_path:
            print(f"[ERROR] Video path not found for UUID: {args.uuid}")
            sys.exit(1)
    
    print(f"=== Face Crop Extraction ===")
    print(f"UUID: {args.uuid}")
    print(f"Video: {video_path}")
    print(f"Output: {args.output_dir}/.faces/{args.uuid}/")
    print()
    
    extractor = FaceCropExtractor(args.output_dir)
    results = extractor.process_video(args.uuid, video_path)
    extractor.print_stats()


def query_video_path_from_db(uuid: str) -> Optional[str]:
    """從 PostgreSQL 取得影片路徑"""
    psql_path = "/opt/homebrew/Cellar/libpq/18.3/bin/psql"
    
    if not os.path.exists(psql_path):
        return None
    
        cmd = [
            psql_path, "-U", "accusys", "-d", "momentry", "-t", "-A",
            "-c", f"SELECT file_path FROM public.videos WHERE file_uuid = '{uuid}' LIMIT 1"
        ]
    
    try:
        proc = subprocess.run(cmd, capture_output=True, text=True, timeout=5)
        path = proc.stdout.strip()
        return path if path else None
    except Exception:
        return None


if __name__ == "__main__":
    main()