#!/usr/bin/env python3
"""存量关键帧批量补红框（计算在 Edge/Oracle，NAS 只编排与存图）

流程: 遍历 frames/event_*/frame_*.jpg -> 分批(8张)上传 Edge /api/edge/mark_frames
      -> 用标记后的图覆盖原文件 -> 写 meta.json (frame_index -> face_count)
幂等: 已有 meta.json 的事件跳过；--force 强制重跑
备份: 首次覆盖前原文件备份到 frames_orig/event_*/
"""
import argparse
import base64
import json
import os
import shutil
import sys

import requests

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))
from fam_core.config_loader import load_config  # noqa: E402

FRAME_DIR = load_config().get('storage', {}).get(
    'frame_image_dir', '/volume1/web/sentinel-home-ai/fam-ui/static/frames')
MARK_URL = load_config().get('poller', {}).get(
    'results_url', 'http://129.146.203.203:5000/api/edge/results'
).rsplit('/', 1)[0] + '/mark_frames'
BATCH = 8


def mark_batch(images):
    """images: [(key, jpeg_bytes)] -> {key: (marked_bytes, faces)}"""
    payload = {
        'images': [
            {'key': k, 'data': base64.b64encode(b).decode('ascii')}
            for k, b in images
        ]
    }
    resp = requests.post(MARK_URL, json=payload, timeout=(30, 120))
    resp.raise_for_status()
    out = {}
    for r in resp.json().get('results', []):
        if r.get('data'):
            out[r['key']] = (base64.b64decode(r['data']), r.get('faces', 0))
        else:
            out[r['key']] = (None, 0)
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--force', action='store_true', help='忽略已有 meta.json 重跑')
    ap.add_argument('--event', type=int, help='只处理指定 event_id')
    args = ap.parse_args()

    orig_root = os.path.join(os.path.dirname(FRAME_DIR.rstrip('/')), 'frames_orig')
    stat = {'events': 0, 'frames': 0, 'faces': 0, 'skip': 0, 'fail': 0}

    for name in sorted(os.listdir(FRAME_DIR)):
        if not name.startswith('event_'):
            continue
        eid = int(name.split('_')[1])
        if args.event and eid != args.event:
            continue
        edir = os.path.join(FRAME_DIR, name)
        meta_path = os.path.join(edir, 'meta.json')
        frames = sorted(
            f for f in os.listdir(edir)
            if f.startswith('frame_') and f.endswith('.jpg'))
        if not frames:
            continue
        if os.path.exists(meta_path) and not args.force:
            stat['skip'] += 1
            continue

        # 备份原件（一次）
        bak_dir = os.path.join(orig_root, name)
        if not os.path.isdir(bak_dir):
            os.makedirs(bak_dir, exist_ok=True)
            for f in frames:
                shutil.copy2(os.path.join(edir, f), os.path.join(bak_dir, f))

        face_counts = {}
        for i in range(0, len(frames), BATCH):
            batch = []
            for f in frames[i:i + BATCH]:
                with open(os.path.join(edir, f), 'rb') as fh:
                    batch.append((f, fh.read()))
            try:
                marked = mark_batch(batch)
            except Exception as e:
                print(f'[event {eid}] 批次失败 (跳过 {len(batch)} 帧): {e}')
                stat['fail'] += len(batch)
                continue
            for f, _ in batch:
                data, faces = marked.get(f, (None, 0))
                if data is None:
                    stat['fail'] += 1
                    continue
                with open(os.path.join(edir, f), 'wb') as fh:
                    fh.write(data)
                idx = f[len('frame_'):-len('.jpg')]
                face_counts[idx] = faces
                stat['frames'] += 1
                stat['faces'] += faces

        if face_counts:
            with open(meta_path, 'w') as fh:
                json.dump(face_counts, fh)
        stat['events'] += 1
        print(f'[event {eid}] 标记 {len(face_counts)}/{len(frames)} 帧, '
              f'人脸合计 {sum(face_counts.values())}')

    print(f"\n完成: 事件 {stat['events']} | 帧标记 {stat['frames']} "
          f"(含人脸帧人脸数 {stat['faces']}) | 已标跳过 {stat['skip']} | 失败 {stat['fail']}")
    print(f'备份目录: {orig_root}')


if __name__ == '__main__':
    main()