"""
Step 1: 从视频目录中提取每个三级文件夹（采集时段）的代表帧
- 遍历所有 L1/L2/L3 结构，找最后一个 MP4，抽取第 0 帧
- 统计该 L3 下所有 MP4 的总时长（分钟）
- 输出图片到 /mnt/nas/Processing_data/Label_cites/ 保持相同目录结构
- 生成索引文件 sessions_index.json 供 step2 使用
- 每个 (L2, 日期) 组生成日代表图：取当天首/中/尾三个 L3 的图片

特殊结构处理：
  - ShenZhen: L3 下还有 L4 子文件夹，递归找 MP4
  - Haerbin: MP4 直接在 L2（无 L3），L2 本身作为唯一 session
  - Hongkong: L2 下有 'video' 子文件夹作为 pseudo-L3

运行：
  conda activate nds && python step1_extract_frames.py           # 断点续跑
  conda activate nds && python step1_extract_frames.py --reset  # 清空重新提取
"""

import os
import json
import re
import sys
from collections import defaultdict
from pathlib import Path
from concurrent.futures import ProcessPoolExecutor, as_completed

import cv2

NUM_WORKERS = os.cpu_count() or 4

VIDEO_ROOT = "/mnt/nas/ADSafety/ADSafety/Aerial_Raw_Videos/2025-Changchun-AerialVideo-V1/"
OUTPUT_ROOT = "/mnt/nas/Processing_data/Label_cites"
INDEX_FILE = "/home/stu6/Projects/LabelWork/sessions_index.json"
VIDEO_EXTS = {".mp4", ".MP4"}


# ─────────────────────────── 工具函数 ────────────────────────────

def find_mp4s(folder: Path) -> list:
    """递归找 folder 下所有 MP4 文件，按名称排序"""
    found = []
    for root, _, files in os.walk(folder):
        for f in sorted(files):
            if Path(f).suffix in VIDEO_EXTS:
                found.append(Path(root) / f)
    return sorted(found)


def get_video_duration_minutes(video_path: Path) -> float:
    """用 OpenCV 获取视频时长（分钟）"""
    try:
        cap = cv2.VideoCapture(str(video_path))
        fps = cap.get(cv2.CAP_PROP_FPS)
        frame_count = cap.get(cv2.CAP_PROP_FRAME_COUNT)
        cap.release()
        if fps > 0 and frame_count > 0:
            return frame_count / fps / 60.0
    except Exception:
        pass
    # 回退：文件大小估算（1GB ≈ 1分钟）
    try:
        return video_path.stat().st_size / (1024 ** 3)
    except Exception:
        return 0.0


def extract_first_frame(video_path: Path, output_jpg: Path) -> bool:
    """用 OpenCV 抽取视频第 0 帧为 JPG"""
    output_jpg.parent.mkdir(parents=True, exist_ok=True)
    if output_jpg.exists():
        return True
    try:
        cap = cv2.VideoCapture(str(video_path))
        ret, frame = cap.read()
        cap.release()
        if ret:
            cv2.imwrite(str(output_jpg), frame)
            return output_jpg.exists()
    except Exception as e:
        print(f"    [ERROR] OpenCV 失败 {video_path.name}: {e}")
    return False


def parse_collection_time(video_name: str) -> str:
    """从文件名解析采集时间，格式：DJI_20250711155546_0001_V → 2025年7月11日 15:55"""
    m = re.search(r"(\d{4})(\d{2})(\d{2})(\d{2})(\d{2})", video_name)
    if m:
        y, mo, d, h, mi = m.groups()
        return f"{y}年{int(mo)}月{int(d)}日 {int(h):02d}:{int(mi):02d}"
    return ""


def parse_date_str(name: str) -> str:
    """从名称中提取日期字符串 YYYYMMDD，用于按天分组"""
    m = re.search(r"(\d{8})", name)
    return m.group(1) if m else ""


# ──────────────────────────── 核心处理 ────────────────────────────

def process_session(l1_name: str, l2_name: str, l3_name: str,
                    session_folder: Path, relative_path: str):
    """
    处理一个 session：找最后一个 MP4、提取其第 0 帧、算总时长
    返回 session 信息字典，失败返回 None
    """
    mp4s = find_mp4s(session_folder)
    if not mp4s:
        print(f"    [SKIP] 无 MP4: {session_folder}")
        return None

    last_mp4 = mp4s[-1]   # 取最后一个视频
    video_name = last_mp4.name

    output_jpg = Path(OUTPUT_ROOT) / relative_path / (last_mp4.stem + ".jpg")

    ok = extract_first_frame(last_mp4, output_jpg)
    if not ok:
        print(f"    [ERROR] 帧提取失败: {relative_path}")
        return None

    total_duration = sum(get_video_duration_minutes(mp4) for mp4 in mp4s)
    collection_time = parse_collection_time(video_name)

    return {
        "l1": l1_name,
        "l2": l2_name,
        "l3": l3_name,
        "folder_path": relative_path,
        "video_name": video_name,
        "location_name": l2_name,
        "collection_time": collection_time,
        "duration": round(total_duration, 4),
        "image_path": str(output_jpg),
        "mp4_count": len(mp4s),
    }


def build_day_summaries(sessions: list) -> list:
    """
    对每个 (l1, l2, 日期) 分组，选出首/中/尾三个 session 作为当天代表图。
    日期从 l3 名称解析，回退到 video_name。
    """
    groups = defaultdict(list)
    for s in sessions:
        date = parse_date_str(s["l3"]) or parse_date_str(s["video_name"])
        groups[(s["l1"], s["l2"], date)].append(s)

    summaries = []
    for (l1, l2, date), group in sorted(groups.items()):
        group = sorted(group, key=lambda s: s["l3"])
        n = len(group)

        # 选首/中/尾，去重
        candidates = [group[0], group[n // 2], group[-1]]
        seen, picks = set(), []
        for p in candidates:
            if p["folder_path"] not in seen:
                seen.add(p["folder_path"])
                picks.append(p)

        date_display = ""
        if len(date) == 8:
            date_display = f"{date[:4]}年{int(date[4:6])}月{int(date[6:8])}日"

        total_duration = round(sum(s["duration"] for s in group), 4)

        representatives = [
            {
                "folder_path": p["folder_path"],
                "image_path": p["image_path"],
                "collection_time": p["collection_time"],
                "duration": p["duration"],
            }
            for p in picks
        ]

        summaries.append({
            "l1": l1,
            "l2": l2,
            "date": date,
            "date_display": date_display,
            "session_count": n,
            "total_duration": total_duration,
            "representatives": representatives,
        })

    return summaries


# ────────────────────────── 目录扫描 ──────────────────────────────

def collect_tasks(video_root: Path, done_paths: set) -> list:
    """
    Phase 1：单线程遍历目录，收集所有待处理任务。
    返回 list of (l1_name, l2_name, l3_name, session_folder, rel)
    """
    tasks = []
    for l1_dir in sorted(video_root.iterdir()):
        if not l1_dir.is_dir():
            continue
        l1_name = l1_dir.name

        for l2_dir in sorted(l1_dir.iterdir()):
            if not l2_dir.is_dir():
                continue
            l2_name = l2_dir.name

            # ── 特殊情况1：Haerbin —— MP4 直接在 L2，无 L3 ──
            mp4s_in_l2 = [f for f in l2_dir.iterdir()
                          if f.is_file() and f.suffix in VIDEO_EXTS]
            if mp4s_in_l2:
                rel = f"{l1_name}/{l2_name}"
                if rel not in done_paths:
                    tasks.append((l1_name, l2_name, l2_name, l2_dir, rel))
                continue

            # ── 特殊情况2：Hongkong —— L2 下有 'video' 子文件夹 ──
            video_subdir = l2_dir / "video"
            if video_subdir.is_dir():
                rel = f"{l1_name}/{l2_name}/video"
                if rel not in done_paths:
                    tasks.append((l1_name, l2_name, "video", video_subdir, rel))
                continue

            # ── 标准情况：L3 子文件夹 ──
            for l3_dir in sorted(d for d in l2_dir.iterdir() if d.is_dir()):
                l3_name = l3_dir.name
                rel = f"{l1_name}/{l2_name}/{l3_name}"
                if rel not in done_paths:
                    tasks.append((l1_name, l2_name, l3_name, l3_dir, rel))

    return tasks


def _worker(args):
    """进程池工作函数"""
    l1_name, l2_name, l3_name, session_folder, rel = args
    return process_session(l1_name, l2_name, l3_name, session_folder, rel)


# ──────────────────────────── 主流程 ──────────────────────────────

def main():
    reset = "--reset" in sys.argv
    video_root = Path(VIDEO_ROOT)

    if reset:
        print("[RESET] 删除已有图片和索引...")
        if Path(OUTPUT_ROOT).exists():
            jpgs = list(Path(OUTPUT_ROOT).rglob("*.jpg")) + list(Path(OUTPUT_ROOT).rglob("*.JPG"))
            for jpg in jpgs:
                jpg.unlink(missing_ok=True)
            print(f"  已删除 {len(jpgs)} 张图片")
        if Path(INDEX_FILE).exists():
            Path(INDEX_FILE).unlink()
            print(f"  已删除 {INDEX_FILE}")

    # 加载已有索引（支持断点续跑）
    sessions = []
    done_paths = set()
    if Path(INDEX_FILE).exists():
        with open(INDEX_FILE, encoding="utf-8") as f:
            existing = json.load(f)
        sessions = existing.get("sessions", existing) if isinstance(existing, dict) else existing
        done_paths = {s["folder_path"] for s in sessions}
        print(f"[断点续跑] 已有 {len(done_paths)} 条记录")

    # Phase 1：收集任务
    print("Phase 1: 扫描目录结构...")
    tasks = collect_tasks(video_root, done_paths)
    print(f"  待处理 session 数: {len(tasks)}，使用 {NUM_WORKERS} 个进程")

    if tasks:
        # Phase 2：并行提取帧
        print("Phase 2: 并行提取帧...")
        SAVE_INTERVAL = 20
        completed = 0

        with ProcessPoolExecutor(max_workers=NUM_WORKERS) as executor:
            futures = {executor.submit(_worker, t): t[4] for t in tasks}
            for future in as_completed(futures):
                rel = futures[future]
                try:
                    session = future.result()
                except Exception as e:
                    print(f"  [ERROR] {rel}: {e}")
                    session = None

                if session:
                    sessions.append(session)
                    done_paths.add(rel)
                    print(f"  [OK] {rel}")

                completed += 1
                if completed % SAVE_INTERVAL == 0:
                    _save_index(sessions)
                    print(f"  [保存] {completed}/{len(tasks)} 完成")

    # Phase 3：生成日代表图索引
    print("Phase 3: 生成日代表图索引...")
    day_summaries = build_day_summaries(sessions)
    print(f"  共 {len(day_summaries)} 个 (场景, 日期) 组合")

    _save_index(sessions, day_summaries)

    # 统计汇报
    print(f"\n✅ 完成: {len(sessions)} 个 sessions，{len(day_summaries)} 个日期组")
    by_l1 = defaultdict(lambda: {"count": 0, "duration": 0.0})
    for s in sessions:
        by_l1[s["l1"]]["count"] += 1
        by_l1[s["l1"]]["duration"] += s["duration"]
    for k in sorted(by_l1):
        print(f"  {k}: {by_l1[k]['count']} sessions, {by_l1[k]['duration']:.1f} 分钟")


def _save_index(sessions: list, day_summaries: list = None):
    data = {"sessions": sessions}
    if day_summaries is not None:
        data["day_summaries"] = day_summaries
    elif Path(INDEX_FILE).exists():
        # 保留已有的 day_summaries
        with open(INDEX_FILE, encoding="utf-8") as f:
            old = json.load(f)
        if isinstance(old, dict) and "day_summaries" in old:
            data["day_summaries"] = old["day_summaries"]
    with open(INDEX_FILE, "w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)


if __name__ == "__main__":
    main()