#!/usr/bin/env python3
# coding: utf-8

import argparse
import json
from pathlib import Path
from typing import Any
from uuid import NAMESPACE_URL, uuid5


DEFAULT_RESULTS_DIR = "/data1/dongying/Mono3d/D4Q2/feishu_project/0416_chenyile/results"
DEFAULT_EXAMPLE_JSON = "/data1/dongying/Mono3d/D4Q2/feishu_project/0416_dongying/example_format.json"
DEFAULT_VIS_DIR = "/data1/dongying/Mono3d/D4Q2/feishu_project/0416_chenyile/vis"
DEFAULT_OUTPUT_JSON = "/data1/dongying/Mono3d/D4Q2/feishu_project/0416_chenyile/results_example_format.json"

DEFAULT_USER_ID = 1698
DEFAULT_SOURCE = "default"
DEFAULT_IMAGE_SIZE = [0, 0]
IMAGE_SUFFIXES = (".jpg", ".jpeg", ".png", ".bmp", ".webp")

LABEL_OVERRIDES = {
    "pedestrian": (["pedestrian", "pedestrian"], ["行人", "行人"]),
    "face": (["sensitive_area", "face"], ["敏感区域", "人脸"]),
    "suv": (["vehicle", "suv"], ["机动车", "SUV"]),
    "car": (["vehicle", "car"], ["机动车", "小轿车"]),
    "truck": (["vehicle", "truck"], ["机动车", "卡车"]),
    "traffic_sign": (["traffic_sign", "traffic_sign_unknown"], ["交通标志", "未知交通标志"]),
    "plate": (["sensitive_area", "plate"], ["敏感区域", "车牌"]),
    "wheel": (["vehicle", "wheel"], ["机动车", "车轮"]),
    "traffic_cone": (["road_barrier", "traffic_cone"], ["路障", "交通锥"]),
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="将结果目录重排为 example_format.json 同款结构。")
    parser.add_argument("results_dir", nargs="?", default=DEFAULT_RESULTS_DIR, help="输入结果目录。")
    parser.add_argument("output_json", nargs="?", default=DEFAULT_OUTPUT_JSON, help="输出 JSON 文件。")
    parser.add_argument("--example-json", default=DEFAULT_EXAMPLE_JSON, help="示例 JSON，用于补充已知标签映射。")
    parser.add_argument("--vis-dir", default=DEFAULT_VIS_DIR, help="可选图片目录，用于补 image_name 和 image_size。")
    parser.add_argument("--userid", type=int, default=DEFAULT_USER_ID, help="输出 annotations.userid。")
    parser.add_argument("--source", default=DEFAULT_SOURCE, help="输出 annotations.source。")
    parser.add_argument(
        "--unknown-occlusion",
        default="no",
        choices=["no", "slight", "partial", "heavy"],
        help="结果里 unknown_occlusion 的回填值。",
    )
    return parser.parse_args()


def load_example_label_map(example_json: Path) -> dict[str, tuple[list[str], list[str]]]:
    if not example_json.is_file():
        return {}

    with example_json.open("r", encoding="utf-8") as file_obj:
        records = json.load(file_obj)

    label_map: dict[str, tuple[list[str], list[str]]] = {}
    for record in records:
        for annotation_block in record.get("annotations", []):
            for annotation in annotation_block.get("anno", []):
                label = annotation.get("label") or []
                label_desc = annotation.get("labelDesc") or []
                if len(label) < 2 or len(label_desc) < 2:
                    continue
                child_label = str(label[1])
                label_map[child_label] = ([str(item) for item in label], [str(item) for item in label_desc])
    return label_map


def resolve_label_mapping(name: str, example_label_map: dict[str, tuple[list[str], list[str]]]) -> tuple[list[str], list[str]]:
    if name in LABEL_OVERRIDES:
        return LABEL_OVERRIDES[name]
    if name in example_label_map:
        return example_label_map[name]
    return (["vehicle", "unknown"], ["机动车", "未知车辆"])


def normalize_occlusion(raw_occlusion: Any, unknown_occlusion: str) -> str:
    mapping = {
        "no_occlusion": "no",
        "slight_occlusion": "slight",
        "partial_occlusion": "partial",
        "heavy_occlusion": "heavy",
        "unknown_occlusion": unknown_occlusion,
        None: "no",
    }
    return mapping.get(raw_occlusion, unknown_occlusion)


def stable_uuid(prefix: str, value: str) -> str:
    return str(uuid5(NAMESPACE_URL, f"{prefix}:{value}"))


def build_vis_image_map(vis_dir: Path) -> dict[str, Path]:
    if not vis_dir.is_dir():
        raise FileNotFoundError(f"vis directory not found: {vis_dir}")

    vis_image_map: dict[str, Path] = {}
    for image_path in sorted(vis_dir.rglob("*")):
        if not image_path.is_file() or image_path.suffix.lower() not in IMAGE_SUFFIXES:
            continue
        vis_image_map.setdefault(image_path.stem, image_path)
    return vis_image_map


def detect_image_info(stem: str, vis_image_map: dict[str, Path]) -> tuple[str, list[int]]:
    image_path = vis_image_map.get(stem)
    if image_path is None:
        raise FileNotFoundError(f"vis image not found for result frame: {stem}")
    return image_path.name, list(DEFAULT_IMAGE_SIZE)


def convert_detection(
    detection: dict[str, Any],
    stem: str,
    index: int,
    example_label_map: dict[str, tuple[list[str], list[str]]],
    unknown_occlusion: str,
) -> dict[str, Any]:
    bbox_xyxy = detection.get("bbox_xyxy")
    if not isinstance(bbox_xyxy, list) or len(bbox_xyxy) != 4:
        raise ValueError(f"invalid bbox_xyxy in {stem}: {detection}")

    label, label_desc = resolve_label_mapping(str(detection.get("name", "unknown")), example_label_map)
    fake_cls_id = detection.get("fake_cls_id", -1)
    annotation_attr = {
        "is_fake": "true" if fake_cls_id not in (-1, None) else "false",
        "is_inferred": "false",
        "occlusion": normalize_occlusion(detection.get("occlusion"), unknown_occlusion),
        "truncation": "no",
    }

    return {
        "attr": annotation_attr,
        "data": {
            "point2d": [
                [float(bbox_xyxy[0]), float(bbox_xyxy[1])],
                [float(bbox_xyxy[2]), float(bbox_xyxy[3])],
            ]
        },
        "id": index,
        "index": stable_uuid("annotation-index", f"{stem}:{index}"),
        "label": label,
        "labelDesc": label_desc,
        "type": "rect",
    }


def convert_result_file(
    result_path: Path,
    example_label_map: dict[str, tuple[list[str], list[str]]],
    vis_image_map: dict[str, Path],
    userid: int,
    source: str,
    unknown_occlusion: str,
) -> dict[str, Any]:
    with result_path.open("r", encoding="utf-8") as file_obj:
        detections = json.load(file_obj)

    if not isinstance(detections, list):
        raise ValueError(f"result file must be a list: {result_path}")

    stem = result_path.stem
    image_name, image_size = detect_image_info(stem, vis_image_map)
    anno = [
        convert_detection(detection, stem, index, example_label_map, unknown_occlusion)
        for index, detection in enumerate(detections)
    ]

    return {
        "annotations": [
            {
                "anno": anno,
                "anno_uuid": stable_uuid("anno", stem),
                "source": source,
                "userid": userid,
            }
        ],
        "image_attr": None,
        "image_name": image_name,
        "image_size": image_size,
        "image_uuid": stable_uuid("image", stem),
        "package_image_uuid": stable_uuid("package-image", stem),
    }


def main() -> int:
    args = parse_args()
    results_dir = Path(args.results_dir).resolve()
    output_json = Path(args.output_json).resolve()
    example_json = Path(args.example_json).resolve()
    vis_dir = Path(args.vis_dir).resolve()

    if not results_dir.is_dir():
        raise FileNotFoundError(f"results directory not found: {results_dir}")

    example_label_map = load_example_label_map(example_json)
    vis_image_map = build_vis_image_map(vis_dir)
    result_files = sorted(results_dir.glob("*.json"))
    if not result_files:
        raise ValueError(f"no json files found under {results_dir}")

    selected_result_files = [result_path for result_path in result_files if result_path.stem in vis_image_map]
    skipped_result_files = [result_path for result_path in result_files if result_path.stem not in vis_image_map]
    if not selected_result_files:
        raise ValueError(f"no result frames matched vis images under {vis_dir}")

    converted = [
        convert_result_file(
            result_path,
            example_label_map,
            vis_image_map,
            userid=args.userid,
            source=args.source,
            unknown_occlusion=args.unknown_occlusion,
        )
        for result_path in selected_result_files
    ]

    output_json.parent.mkdir(parents=True, exist_ok=True)
    with output_json.open("w", encoding="utf-8") as file_obj:
        json.dump(converted, file_obj, ensure_ascii=False, indent=2)

    print(
        f"result_files={len(result_files)} matched_vis_files={len(selected_result_files)} "
        f"skipped_files={len(skipped_result_files)} output={output_json}"
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())