"""
按 group_id 提取真值产物（标注 JSON + 标定文件，可选保存图片）。

用法：
    python extract_gt_by_group_id.py \
        --group_ids <uuid1> <uuid2> ... \
        --output /data/output_dir \
        [--car_name G1M3_0630] \
        [--product_type 2d-3d-association] \
        [--camera_id camera4] \
        [--image_suffix png] \
        [--annotation_stride 2] \
        [--skip_images] \
        [--workers 4]

或者通过文本文件批量输入（每行: group_uuid [date] [car_name]）：
    python extract_gt_by_group_id.py \
        --uuid_file /data/group_ids.txt \
        --output /data/output_dir
"""

import argparse
import io
import os
import tarfile
from datetime import datetime, timezone, timedelta
from math import ceil
from multiprocessing import Pool
from typing import Dict, List, Optional, Tuple

import numpy as np
import pandas as pd

from dotenv import load_dotenv
load_dotenv()

from pdcl_dss import Group

os.environ['STS_UID'] = 'dis-uploader'
os.environ['STS_SECRET_KEY'] = '277310cc09724d315514a79701fecb0f'

BEIJING_TZ = timezone(timedelta(hours=8))
DEFAULT_CAR_NAME = 'unknown'
DEFAULT_DATE_NAME = '00000000000000'
DEFAULT_IMAGE_SUFFIX = 'png'
DEFAULT_ANNOTATION_STRIDE = 1
SUPPORTED_IMAGE_SUFFIXES = {'png', 'jpg', 'jpeg'}


# ---------------------------------------------------------------------------
# 视频解码工具
# ---------------------------------------------------------------------------

def _plane_to_ndarray(plane) -> np.ndarray:
    stride = plane.line_size
    height = plane.height
    width = plane.width
    arr = np.frombuffer(plane, dtype=np.uint8)
    if stride == width:
        return arr.reshape(height, width)
    return arr.reshape(height, stride)[:, :width]


def _yuvj420p_to_bgr(frame) -> np.ndarray:
    """将 av YUVJ420P 帧转换为 BGR ndarray。"""
    import cv2

    y  = _plane_to_ndarray(frame.planes[0])
    u  = _plane_to_ndarray(frame.planes[1])
    v  = _plane_to_ndarray(frame.planes[2])
    h, w = y.shape
    uv = np.zeros((h // 2, w), dtype=np.uint8)
    uv[:, 0::2] = u
    uv[:, 1::2] = v
    yuv = np.concatenate([y, uv], axis=0)
    return cv2.cvtColor(yuv, cv2.COLOR_YUV2BGR_NV12)


def _h265_bytes_to_bgr(frame_bytes: bytes) -> np.ndarray:
    import av

    container = av.open(io.BytesIO(frame_bytes))
    for av_frame in container.decode(video=0):
        return _yuvj420p_to_bgr(av_frame)
    raise ValueError("H.265 解码失败")


# ---------------------------------------------------------------------------
# PDCL 数据读取
# ---------------------------------------------------------------------------

def get_frame_timestamps(group_uuid: str) -> pd.DataFrame:
    """从 clean_data_archive.tar 中读取 frame_timestamps.csv。"""
    with Group(group_uuid) as group:
        with group.open('clean_data_archive.tar', mode='rb') as f:
            tar = tarfile.open(fileobj=f)
            csv_obj = tar.extractfile(tar.getmember('frame_timestamps.csv'))
            return pd.read_csv(csv_obj)


def get_group_meta(group_uuid: str) -> Dict:
    """读取 group 元数据。"""
    with Group(group_uuid) as group:
        return dict(group.meta)


def get_clip_paths(group_uuid: str) -> List[str]:
    """获取 group 下所有 clip 的本地缓存路径。"""
    from pdcl_dss import Clip

    paths = []
    with Group(group_uuid) as group:
        for clip_ukey in group.list_clip_ukeys():
            clip = Clip(clip_ukey)
            files, _ = clip.list_files()
            paths.append(clip.get_cache_path(files[0]))
    return paths


def _find_index_by_timestamp(df: pd.DataFrame, camera_id: str, ts_ns: int) -> int:
    return df[abs(df[camera_id] - ts_ns) < 1e-3].index[0]


def _build_timestamp_index(df: pd.DataFrame, camera_id: str) -> Dict[int, int]:
    if camera_id not in df.columns:
        return {}

    timestamp_index: Dict[int, int] = {}
    for json_index, value in df[camera_id].items():
        if pd.isna(value):
            continue
        timestamp_index[int(value)] = int(json_index)
    return timestamp_index


def _find_index_by_timestamp_fast(
    timestamp_index: Dict[int, int],
    df: pd.DataFrame,
    camera_id: str,
    ts_ns: int,
) -> int:
    json_index = timestamp_index.get(int(ts_ns))
    if json_index is not None:
        return json_index
    return _find_index_by_timestamp(df, camera_id, ts_ns)


def _is_missing_car_name(car_name: str) -> bool:
    return not car_name or car_name == DEFAULT_CAR_NAME


def _is_missing_date_name(date_name: str) -> bool:
    return not date_name or date_name == DEFAULT_DATE_NAME


def _format_ms_to_date_name(timestamp_ms: int) -> str:
    dt = datetime.fromtimestamp(timestamp_ms / 1000, tz=timezone.utc).astimezone(BEIJING_TZ)
    return dt.strftime('%Y%m%d%H%M%S')


def _get_frame_id_column(camera_id: str) -> str:
    return f'{camera_id}_frame_id'


def _is_target_camera_annotation_member(member_name: str, camera_id: str) -> bool:
    normalized_name = member_name.replace('\\', '/')
    target_fragment = f'/{camera_id}/camera_radar_lidar_jsons/'
    return normalized_name.endswith('.json') and target_fragment in f'/{normalized_name}'


def normalize_image_suffix(image_suffix: str) -> str:
    normalized = image_suffix.strip().lower().lstrip('.')
    if normalized not in SUPPORTED_IMAGE_SUFFIXES:
        supported = ', '.join(sorted(SUPPORTED_IMAGE_SUFFIXES))
        raise ValueError(f'不支持的 image_suffix: {image_suffix}. 支持: {supported}')
    return normalized


def normalize_annotation_stride(annotation_stride: int) -> int:
    if annotation_stride < 1:
        raise ValueError(f'annotation_stride 必须 >= 1，当前为: {annotation_stride}')
    return annotation_stride


def _get_frame_timestamp_ns(df: pd.DataFrame, camera_id: str, json_index: int) -> Optional[int]:
    try:
        value = df.at[json_index, camera_id]
    except (KeyError, ValueError):
        return None

    if pd.isna(value):
        return None

    return int(value)


def _build_frame_basename(
    car_name: str,
    date_name: str,
    group_uuid: str,
    json_index: int,
    camera_id: str,
    df: pd.DataFrame,
) -> str:
    parts = [car_name, date_name, group_uuid, f'{json_index:06d}']

    frame_id_col = _get_frame_id_column(camera_id)
    if frame_id_col in df.columns:
        frame_id = df.at[json_index, frame_id_col]
        if not pd.isna(frame_id):
            parts.append(str(int(frame_id)))

    frame_timestamp_ns = _get_frame_timestamp_ns(df, camera_id, json_index)
    if frame_timestamp_ns is not None:
        parts.append(str(frame_timestamp_ns))

    return '_'.join(parts)


def _infer_car_name(group_meta: Dict) -> Tuple[str, str]:
    for key in ('plate_number', 'vehicle_name', 'car_name', 'plateNumber', 'vehicleName', 'carName'):
        value = str(group_meta.get(key, '')).strip()
        if value:
            return value, key

    project_name = str(group_meta.get('project_name', '')).strip()
    if project_name:
        return project_name, 'project_name'

    return DEFAULT_CAR_NAME, 'default'


def _infer_date_name(group_meta: Dict, df: pd.DataFrame, camera_id: str) -> Tuple[str, str]:
    collection_time = group_meta.get('collection_time')
    if isinstance(collection_time, (int, float)) and collection_time > 0:
        return _format_ms_to_date_name(int(collection_time)), 'collection_time'

    candidate_columns = [camera_id] if camera_id in df.columns else []
    candidate_columns.extend(col for col in df.columns if col not in candidate_columns)

    for col in candidate_columns:
        col_name = str(col)
        if col_name.startswith('Unnamed:'):
            continue
        if any(token in col_name for token in ('frame_id', 'clip_uuid', 'point_count', 'speed')):
            continue

        values = pd.to_numeric(df[col], errors='coerce').dropna()
        if values.empty:
            continue

        timestamp_ns = int(values.iloc[0])
        dt = datetime.fromtimestamp(timestamp_ns / 1e9, tz=timezone.utc).astimezone(BEIJING_TZ)
        return dt.strftime('%Y%m%d%H%M%S'), f'frame_timestamps:{col_name}'

    return DEFAULT_DATE_NAME, 'default'


def resolve_group_identifiers(
    group_uuid: str,
    car_name: str,
    date_name: str,
    df: pd.DataFrame,
    camera_id: str,
) -> Tuple[str, str, List[str]]:
    """
    仅在 car_name/date 为缺省值时，尝试从 group 元数据和时间戳中自动补全。
    """
    resolved_car_name = car_name
    resolved_date_name = date_name
    messages: List[str] = []

    need_car_name = _is_missing_car_name(car_name)
    need_date_name = _is_missing_date_name(date_name)
    if not need_car_name and not need_date_name:
        return resolved_car_name, resolved_date_name, messages

    group_meta = get_group_meta(group_uuid)

    if need_car_name:
        resolved_car_name, source = _infer_car_name(group_meta)
        messages.append(f'[{group_uuid}] 自动补全 car_name={resolved_car_name} (source={source})')

    if need_date_name:
        resolved_date_name, source = _infer_date_name(group_meta, df, camera_id)
        messages.append(f'[{group_uuid}] 自动补全 date={resolved_date_name} (source={source})')

    return resolved_car_name, resolved_date_name, messages


# ---------------------------------------------------------------------------
# 图片保存
# ---------------------------------------------------------------------------

def save_images(
    group_uuid: str,
    car_name: str,
    date_name: str,
    valid_ids: Optional[List[int]],
    clip_paths: List[str],
    camera_id: str,
    df: pd.DataFrame,
    timestamp_index: Dict[int, int],
    output_path: str,
    image_suffix: str,
) -> int:
    """
    解析视频流并保存为指定后缀图片。

    Args:
        valid_ids: 要保存的 json_index 列表；若为 None 则保存所有帧。

    Returns:
        保存的图片数量。
    """
    import cv2
    from pdcl_pyclip.decoder_struct import StructDecoder
    from pdcl_pyclip.msg_camera import VideoMessage
    from pdcl_pyclip.reader import ClipReader

    image_dir = os.path.join(output_path, 'images')
    os.makedirs(image_dir, exist_ok=True)

    valid_set = set(valid_ids) if valid_ids is not None else None
    struct_decoder = StructDecoder()
    saved = 0

    for clip_path in clip_paths:
        reader = ClipReader(clip_path)
        for schema, channel, msg in reader.iter_messages(topics=[camera_id]):
            if schema.encoding != 'struct':
                continue
            data = struct_decoder.decode(schema, channel, msg)
            if not isinstance(data, VideoMessage):
                continue
            json_index = _find_index_by_timestamp_fast(timestamp_index, df, camera_id, msg.log_time)
            if valid_set is not None and json_index not in valid_set:
                continue
            frame_timestamp_ns = _get_frame_timestamp_ns(df, camera_id, json_index)
            if frame_timestamp_ns is not None and msg.log_time != frame_timestamp_ns:
                print(
                    f"    [WARN] 帧 {json_index} 时间戳不一致: "
                    f"msg.log_time={msg.log_time}, df[{camera_id}]={frame_timestamp_ns}"
                )
            try:
                img = _h265_bytes_to_bgr(data.payload)
            except Exception as e:
                print(f"    [WARN] 帧 {json_index} 解码失败: {e}")
                continue
            fname = (
                f'{_build_frame_basename(car_name, date_name, group_uuid, json_index, camera_id, df)}'
                f'.{image_suffix}'
            )
            if not cv2.imwrite(os.path.join(image_dir, fname), img):
                print(f"    [WARN] 帧 {json_index} 写图失败: {fname}")
                continue
            saved += 1

    return saved


# ---------------------------------------------------------------------------
# 标注 / 标定文件解包
# ---------------------------------------------------------------------------

def extract_annotations_and_calib(
    group_uuid: str,
    car_name: str,
    date_name: str,
    product_type: str,
    camera_id: str,
    df: pd.DataFrame,
    output_path: str,
    annotation_stride: int = DEFAULT_ANNOTATION_STRIDE,
) -> List[int]:
    """
    从真值产物 tar 包中提取标注 JSON 和标定文件。

    Returns:
        实际保存的有效帧 json_index 列表（仅目标 camera_id 下的帧）。
    """
    tar_filename = f'{product_type}_archive.tar'
    valid_ids: List[int] = []
    candidate_annotation_count = 0
    saved_annotation_names = set()

    with Group(group_uuid) as group:
        if group.get_product(product_type) is None:
            print(f"  [WARN] Group {group_uuid} 不含产物 {product_type}，跳过")
            return valid_ids

        calib_dir = os.path.join(output_path, 'calib')
        annotations_dir = os.path.join(output_path, 'annotations')
        os.makedirs(calib_dir, exist_ok=True)
        os.makedirs(annotations_dir, exist_ok=True)

        with group.open(tar_filename, mode='rb') as f:
            with tarfile.open(fileobj=f, mode='r|*') as tar:
                for member in tar:
                    if 'depth_map_fg' in member.name:
                        continue

                    file_obj = tar.extractfile(member)
                    if file_obj is None:
                        continue

                    # ── 标定文件 ─────────────────────────────────────────
                    if 'calib' in member.name:
                        calib_fname = os.path.basename(member.name)
                        with open(os.path.join(calib_dir, calib_fname), 'wb') as fw:
                            fw.write(file_obj.read())
                        continue

                    # ── 标注 JSON ─────────────────────────────────────────
                    if _is_target_camera_annotation_member(member.name, camera_id):
                        json_fname = os.path.basename(member.name)
                        try:
                            json_index = int(json_fname[:-5])
                        except ValueError:
                            print(f"  [WARN] 无法解析帧号: {member.name}")
                            continue

                        candidate_annotation_count += 1
                        if (candidate_annotation_count - 1) % annotation_stride != 0:
                            continue

                        try:
                            new_fname = (
                                f'{_build_frame_basename(car_name, date_name, group_uuid, json_index, camera_id, df)}.json'
                            )
                        except (KeyError, IndexError, ValueError):
                            new_fname = f'{car_name}_{date_name}_{group_uuid}_{json_index:06d}.json'

                        if new_fname in saved_annotation_names:
                            print(f"  [WARN] 重复标注输出名，跳过覆盖: {new_fname} <- {member.name}")
                            continue

                        with open(os.path.join(annotations_dir, new_fname), 'wb') as fw:
                            fw.write(file_obj.read())
                        saved_annotation_names.add(new_fname)
                        valid_ids.append(json_index)

    print(
        f"  [INFO] annotations({camera_id}) 保存 {len(valid_ids)} / "
        f"{candidate_annotation_count} (annotation_stride={annotation_stride})"
    )
    return valid_ids


# ---------------------------------------------------------------------------
# 单个 group 的完整处理流程
# ---------------------------------------------------------------------------

def process_group(
    group_uuid: str,
    car_name: str,
    date_name: str,
    output_root: str,
    product_type: str = '2d-3d-association',
    camera_id: str = 'camera4',
    image_suffix: str = DEFAULT_IMAGE_SUFFIX,
    annotation_stride: int = DEFAULT_ANNOTATION_STRIDE,
    skip_images: bool = False,
) -> Tuple[str, str]:
    """
    处理单个 group：提取标注、标定文件，并按需保存对应图片。

    Returns:
        (group_uuid, status_message)
    """
    output_path = os.path.join(output_root, group_uuid)

    try:
        df = get_frame_timestamps(group_uuid)
        car_name, date_name, messages = resolve_group_identifiers(
            group_uuid, car_name, date_name, df, camera_id
        )
        print(
            f"[{group_uuid}] 开始处理 "
            f"(car={car_name}, date={date_name}, "
            f"annotation_stride={annotation_stride}, skip_images={skip_images})"
        )
        for message in messages:
            print(message)

        # 1. 提取标注 + 标定，获取有效帧列表
        valid_ids = extract_annotations_and_calib(
            group_uuid,
            car_name,
            date_name,
            product_type,
            camera_id,
            df,
            output_path,
            annotation_stride,
        )

        if not valid_ids:
            print(f"[{group_uuid}] 无有效标注帧，跳过后续处理")
            return group_uuid, 'success (no valid frames)'

        valid_ids.sort()
        if skip_images:
            print(f"[{group_uuid}] 共 {len(valid_ids)} 个有效帧，按要求跳过图片保存")
            return group_uuid, f'success ({len(valid_ids)} annotations only)'

        print(f"[{group_uuid}] 共 {len(valid_ids)} 个有效帧，开始保存图片")

        # 2. 保存对应帧图片
        timestamp_index = _build_timestamp_index(df, camera_id)
        clip_paths = get_clip_paths(group_uuid)
        n_saved = save_images(
            group_uuid, car_name, date_name,
            valid_ids, clip_paths, camera_id, df, timestamp_index, output_path, image_suffix
        )
        print(f"[{group_uuid}] 完成，保存图片 {n_saved} 张")
        return group_uuid, f'success ({n_saved} images)'

    except Exception as e:
        import traceback
        print(f"[{group_uuid}] 处理失败: {e}\n{traceback.format_exc()}")
        return group_uuid, f'failed: {e}'


# ---------------------------------------------------------------------------
# 多进程入口
# ---------------------------------------------------------------------------

def _worker(args):
    return process_group(*args)


def run(
    tasks: List[Tuple[str, str, str]],   # [(group_uuid, car_name, date_name), ...]
    output_root: str,
    product_type: str,
    camera_id: str,
    image_suffix: str,
    annotation_stride: int,
    skip_images: bool,
    workers: int,
):
    if not tasks:
        print('没有需要处理的 group')
        return

    full_tasks = [
        (
            uuid,
            car,
            date,
            output_root,
            product_type,
            camera_id,
            image_suffix,
            annotation_stride,
            skip_images,
        )
        for uuid, car, date in tasks
    ]

    actual_workers = max(1, min(workers, len(full_tasks)))

    if actual_workers == 1:
        results = [_worker(t) for t in full_tasks]
    else:
        chunksize = max(1, ceil(len(full_tasks) / (actual_workers * 4)))
        with Pool(processes=actual_workers) as pool:
            results = list(pool.imap_unordered(_worker, full_tasks, chunksize=chunksize))

    success = sum(1 for _, r in results if r.startswith('success'))
    failed  = sum(1 for _, r in results if r.startswith('failed'))
    print(f"\n完成！总计 {len(results)}  成功 {success}  失败 {failed}  workers {actual_workers}")
    for uuid, res in results:
        if res.startswith('failed'):
            print(f"  FAILED {uuid}: {res}")


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def parse_uuid_file(filepath: str) -> List[Tuple[str, str, str]]:
    """
    解析 UUID 文本文件，每行格式：
        group_uuid [date] [car_name]
    缺省 date/car_name 时会在处理时自动尝试补全。
    """
    tasks = []
    seen = set()
    duplicate_count = 0
    with open(filepath) as f:
        for line in f:
            parts = line.strip().split()
            if not parts:
                continue
            uuid      = parts[0]
            date      = parts[1] if len(parts) > 1 else DEFAULT_DATE_NAME
            car_name  = parts[2] if len(parts) > 2 else DEFAULT_CAR_NAME
            task = (uuid, car_name, date)
            if task in seen:
                duplicate_count += 1
                continue
            seen.add(task)
            tasks.append(task)
    if duplicate_count:
        print(f"[uuid_file] 跳过重复 group_uuid {duplicate_count} 条")
    return tasks


def main():
    parser = argparse.ArgumentParser(description='按 group_id 提取 PDCL 真值产物')
    src = parser.add_mutually_exclusive_group(required=True)
    src.add_argument('--group_ids', nargs='+', metavar='UUID',
                     help='直接指定一个或多个 group UUID')
    src.add_argument('--uuid_file', metavar='FILE',
                     help='UUID 列表文件（每行: uuid [date] [car_name]）')

    parser.add_argument('--output',       required=True, help='输出根目录')
    parser.add_argument('--car_name',     default=DEFAULT_CAR_NAME,
                        help='车辆标识；默认会尝试从 group 元数据自动推断')
    parser.add_argument('--date',         default=DEFAULT_DATE_NAME,
                        help='日期字符串；默认会尝试从 collection_time 或时间戳自动推断')
    parser.add_argument('--product_type', default='2d-3d-association',
                        help='真值产物类型（默认: 2d-3d-association）')
    parser.add_argument('--camera_id',    default='camera4', help='相机 topic 名称')
    parser.add_argument('--image_suffix', default=DEFAULT_IMAGE_SUFFIX,
                        help='输出图像后缀，支持 png/jpg/jpeg（默认: png）')
    parser.add_argument('--annotation_stride', type=int, default=DEFAULT_ANNOTATION_STRIDE,
                        help='标注保存抽样步长；1 为全保存，2 为 2 抽 1，同时图片也会与保留标注对齐')
    parser.add_argument('--skip_images',  action='store_true',
                        help='只提取 annotations/calib，不解码或保存图片帧')
    parser.add_argument('--workers',      type=int, default=4, help='并行进程数')

    args = parser.parse_args()
    image_suffix = normalize_image_suffix(args.image_suffix)
    annotation_stride = normalize_annotation_stride(args.annotation_stride)

    if args.uuid_file:
        tasks = parse_uuid_file(args.uuid_file)
    else:
        tasks = [(uuid, args.car_name, args.date) for uuid in args.group_ids]

    os.makedirs(args.output, exist_ok=True)
    run(
        tasks,
        args.output,
        args.product_type,
        args.camera_id,
        image_suffix,
        annotation_stride,
        args.skip_images,
        args.workers,
    )


if __name__ == '__main__':
    main()
