#!/usr/bin/env python3
"""校验真实采集数据集索引是否满足算法输入合同。"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path
from typing import Any

from dataset_index_to_site_data_input import (
    SERVICE_INPUT_FIELDS,
    extract_file_uri,
    is_uri,
    load_yaml,
    normalize_dataset_root,
    resolve_file_uri,
)


TASK_REQUIREMENTS = {
    "external": {
        "section": "external",
        "any_of": ["external_observation_files", "marker_observation_files"],
    },
    "chassis": {
        "section": "chassis",
        "all_of": ["chassis_motion_data_files"],
    },
    "control": {
        "section": "control",
        "all_of": ["control_evaluation_data_files"],
    },
    "sensor_intrinsic": {
        "section": "sensor",
        "all_of": ["image_sample_files"],
        "warn_if_missing": ["target_detection_files"],
    },
    "sensor_extrinsic": {
        "section": "sensor",
        "all_of": ["image_sample_files", "robot_pose_sample_files"],
        "warn_if_missing": ["target_detection_files"],
    },
    "hand_eye": {
        "section": "sensor",
        "all_of": ["image_sample_files", "robot_pose_sample_files"],
    },
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="校验 dataset_index.yaml。")
    parser.add_argument("dataset_index", help="待校验的数据集索引文件。")
    parser.add_argument(
        "--tasks",
        default="external,chassis,control,sensor_intrinsic",
        help="逗号分隔的任务列表，用于检查每个任务的最小数据输入。",
    )
    parser.add_argument(
        "--dataset-root",
        help="覆盖索引文件中的 session.dataset_root，用于现场挂载路径重定向。",
    )
    parser.add_argument(
        "--require-existing-files",
        action="store_true",
        help="要求所有本地文件引用都真实存在；URI 引用只检查格式。",
    )
    parser.add_argument(
        "--no-strict-window",
        action="store_true",
        help="不强制要求数据窗口大于 0 且结束时间晚于开始时间。",
    )
    return parser.parse_args()


def parse_tasks(raw_tasks: str) -> list[str]:
    tasks: list[str] = []
    for item in raw_tasks.split(","):
        task = item.strip()
        if not task:
            continue
        if task not in TASK_REQUIREMENTS:
            raise ValueError(f"未知任务 {task!r}，可选任务：{', '.join(TASK_REQUIREMENTS)}")
        tasks.append(task)
    if not tasks:
        raise ValueError("至少需要指定一个任务。")
    return tasks


def known_sections() -> set[str]:
    return {str(config["index_section"]) for config in SERVICE_INPUT_FIELDS.values()}


def known_fields_by_section() -> dict[str, set[str]]:
    result: dict[str, set[str]] = {}
    for config in SERVICE_INPUT_FIELDS.values():
        section = str(config["index_section"])
        result.setdefault(section, set()).update(str(field) for field in config["fields"])
    return result


def require_map(value: Any, field_name: str) -> dict[str, Any]:
    if not isinstance(value, dict):
        raise ValueError(f"{field_name} 必须是 YAML map。")
    return value


def read_int(value: Any, field_name: str) -> int:
    try:
        return int(value)
    except (TypeError, ValueError) as exc:
        raise ValueError(f"{field_name} 必须是整数，当前值为 {value!r}。") from exc


def file_entries(value: Any) -> list[Any]:
    if value is None:
        return []
    return value if isinstance(value, list) else [value]


def resolve_entries(value: Any, dataset_root: Path) -> list[str]:
    result: list[str] = []
    for entry in file_entries(value):
        uri = resolve_file_uri(extract_file_uri(entry), dataset_root)
        if uri and uri not in result:
            result.append(uri)
    return result


def field_has_files(section: dict[str, Any], field_name: str) -> bool:
    return bool(file_entries(section.get(field_name)))


def validate_session(
    index: dict[str, Any],
    strict_window: bool,
) -> dict[str, Any]:
    if read_int(index.get("schema_version"), "schema_version") != 1:
        raise ValueError("schema_version 必须为 1。")
    session = require_map(index.get("session"), "session")
    for field_name in ("session_id", "vehicle_id", "dataset_root"):
        value = session.get(field_name)
        if not isinstance(value, str) or not value.strip():
            raise ValueError(f"session.{field_name} 必须是非空字符串。")

    start_us = read_int(
        session.get("data_window_start_timestamp_us"),
        "session.data_window_start_timestamp_us",
    )
    end_us = read_int(
        session.get("data_window_end_timestamp_us"),
        "session.data_window_end_timestamp_us",
    )
    if strict_window and (start_us <= 0 or end_us <= 0 or end_us <= start_us):
        raise ValueError("数据窗口无效，要求 start_us > 0 且 end_us > start_us。")
    return session


def validate_sections(index: dict[str, Any]) -> dict[str, dict[str, Any]]:
    data_inputs = require_map(index.get("data_inputs"), "data_inputs")
    valid_sections = known_sections()
    valid_fields = known_fields_by_section()
    result: dict[str, dict[str, Any]] = {}

    for section_name, section_value in data_inputs.items():
        if section_name not in valid_sections:
            raise ValueError(f"未知数据输入模块 data_inputs.{section_name}。")
        section = require_map(section_value, f"data_inputs.{section_name}")
        for field_name in section:
            if field_name not in valid_fields[section_name]:
                raise ValueError(f"未知数据输入字段 data_inputs.{section_name}.{field_name}。")
        result[section_name] = section

    for section_name in valid_sections:
        result.setdefault(section_name, {})
    return result


def validate_task_requirements(
    sections: dict[str, dict[str, Any]],
    tasks: list[str],
) -> list[str]:
    warnings: list[str] = []
    for task in tasks:
        requirement = TASK_REQUIREMENTS[task]
        section_name = str(requirement["section"])
        section = sections.get(section_name, {})

        for field_name in requirement.get("all_of", []):
            if not field_has_files(section, str(field_name)):
                raise ValueError(f"任务 {task} 缺少 data_inputs.{section_name}.{field_name}。")

        any_of = [str(field) for field in requirement.get("any_of", [])]
        if any_of and not any(field_has_files(section, field_name) for field_name in any_of):
            fields = ", ".join(f"data_inputs.{section_name}.{field_name}" for field_name in any_of)
            raise ValueError(f"任务 {task} 至少需要以下字段之一：{fields}。")

        for field_name in requirement.get("warn_if_missing", []):
            if not field_has_files(section, str(field_name)):
                warnings.append(f"任务 {task} 未提供 data_inputs.{section_name}.{field_name}。")
    return warnings


def validate_file_references(
    sections: dict[str, dict[str, Any]],
    dataset_root: Path,
    require_existing_files: bool,
) -> list[str]:
    warnings: list[str] = []
    for section_name, section in sections.items():
        for field_name, value in section.items():
            files = resolve_entries(value, dataset_root)
            if not files:
                warnings.append(f"data_inputs.{section_name}.{field_name} 为空。")
                continue
            if not require_existing_files:
                continue
            for file_uri in files:
                if is_uri(file_uri):
                    continue
                if not Path(file_uri).exists():
                    raise ValueError(f"文件不存在：{file_uri}")
    return warnings


def main() -> int:
    args = parse_args()
    index_path = Path(args.dataset_index).expanduser().resolve(strict=False)
    index = load_yaml(index_path)
    tasks = parse_tasks(args.tasks)
    session = validate_session(index, strict_window=not args.no_strict_window)
    sections = validate_sections(index)
    dataset_root = normalize_dataset_root(index_path, args.dataset_root or session.get("dataset_root"))

    warnings = []
    warnings.extend(validate_task_requirements(sections, tasks))
    warnings.extend(validate_file_references(sections, dataset_root, args.require_existing_files))

    print(f"[OK] dataset_index 校验通过: {index_path}")
    print(f"[OK] 数据集根目录: {dataset_root}")
    print(f"[OK] 任务列表: {', '.join(tasks)}")
    for warning in warnings:
        print(f"[WARN] {warning}")
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except Exception as exc:
        print(f"[错误] {exc}", file=sys.stderr)
        raise SystemExit(1)
