#!/usr/bin/env python
"""把 generate_boxes.py 的产出整理成 Label Studio 复核材料，输出放在 runs/ 下。

    uv run python scripts/prepare_label_studio.py --source runs/boxes_shap

generate_boxes.py 生成的是弱监督框，需要人工在 Label Studio 里修正之后才能拿去训练 YOLO。
它的产出按**类别**分目录（便于人工核对），而 label-studio-converter 的 YOLO 导入只认扁平的
`<input>/{images,labels,classes.txt}`，且 os.listdir 不递归。这一步做三件事：

1. 把按类别分的目录**拍平**；
2. 图片**统一成 PNG**：本来就是 PNG 的硬链接过去，不占额外磁盘；BMP 等其他格式
   （pearlescent / blue / uv 原始数据都是 BMP）无损转成 PNG。不缩放像素，
   归一化的框坐标原样有效；
3. 文件**按类别重新编号**：`<类别>_<序号>.png`，序号在类内按源路径排序，重跑结果一致。
   原文件名里的正反面、位置、时间戳等信息记在 `name_map.csv` 里，不会丢。
   想保留原名加 `--keep-names`。

然后在输出目录上调用 label-studio-converter，生成可直接导入的任务：

```
runs/<source>_label_studio/      # 同时也是 Label Studio 的 DOCUMENT_ROOT
├── images/<类别>_<序号>.png
├── labels/<类别>_<序号>.txt
├── classes.txt                  # 每行一个类别名，行号即 class_id
├── name_map.csv                 # 新名 ↔ 源文件对照
├── summary.json
├── tasks.json                   # 每张图一个任务，框作为 predictions
└── tasks.label_config.xml
```

**图片 URL 相对 Label Studio 的 DOCUMENT_ROOT。** 任务里写的是
`/data/local-files/?d=<相对路径>/images/<图名>.png`，Label Studio 把 `d` 拼到
`LABEL_STUDIO_LOCAL_FILES_DOCUMENT_ROOT` 后面找文件。所以 URL 必须按实际的 DOCUMENT_ROOT 生成：

    # 缺省读环境变量 LABEL_STUDIO_LOCAL_FILES_DOCUMENT_ROOT；也可显式指定
    --document-root /home/deep/labeling

输出目录不在 DOCUMENT_ROOT 下时，会在 DOCUMENT_ROOT 里建一个同名软链接指过去，
产出本身仍留在 runs/。没有设置 DOCUMENT_ROOT 时，输出目录自己就当作 DOCUMENT_ROOT。

这里**不写 data.yaml、不划分 train/val**：这份材料是给人复核的，不是训练集。
修好标签后，再从 Label Studio 导出、划分到 data/ 下给 YOLO 训练。
复核期间 Label Studio 直接从这个目录读图，**不要删除或 --overwrite 它**。

没有框的图（标签为空）默认保留，导入后是一个没有预标注框的任务。但**空标签在 YOLO 里的语义是
"这张图什么都没有"**，不是"有缺陷但没定位出来"。来自缺陷类别的空标签如果原样走到训练，
等于在教模型抑制真实缺陷，所以可以用 `--empty-policy` 按来源类别分别处置：

    --empty-policy drop normal=keep black=full
    # 缺陷类没定位出来的直接丢弃；normal 留作负样本；整图属性的类退化成整图框
"""

from __future__ import annotations

import argparse
import contextlib
import csv
import io
import json
import logging
import os
import shutil
import sys
import urllib.parse
from collections import Counter, defaultdict
from collections.abc import Sequence
from pathlib import Path
from typing import NamedTuple

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import yaml  # noqa: E402

from datasets.yolo import IMAGE_SUFFIXES, write_classes_txt  # noqa: E402
from utils import link_or_copy, setup_logging  # noqa: E402

logger = None

EMPTY_POLICIES = ("keep", "drop", "full")
EMPTY_POLICY_DESC = {
    "keep": "保留为负样本",
    "drop": "丢弃",
    "full": "退化成整图框",
}


PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
# PNG 能原样存下的像素模式；其余（CMYK、YCbCr 等）先转 RGB
PNG_MODES = frozenset({"1", "L", "LA", "I", "I;16", "P", "RGB", "RGBA"})


def is_png(path: Path) -> bool:
    """看文件头而不是后缀：后缀写着 .png、内容却是 BMP/JPEG 的也要转。"""
    with path.open("rb") as handle:
        return handle.read(len(PNG_SIGNATURE)) == PNG_SIGNATURE


def materialize_png(source: Path, destination: Path) -> str:
    """把图片落成 PNG：本来就是 PNG 的硬链接，其余格式无损转换。

    返回 "hard" / "copy" / "convert"。转换不缩放像素，归一化的框坐标原样有效。
    """
    if is_png(source):
        return link_or_copy(source, destination)

    from PIL import Image

    if destination.exists() or destination.is_symlink():
        destination.unlink()
    destination.parent.mkdir(parents=True, exist_ok=True)
    with Image.open(source) as image:
        if image.mode not in PNG_MODES:
            image = image.convert("RGB")
        image.save(destination, format="PNG")
    return "convert"


class Sample(NamedTuple):
    """一张图及其标签。

    Attributes:
        image: 图片路径。
        label: 标签路径。
        class_id: 标签里第一个框的类别下标；无框时为 None。
        group: 来源的类别目录名——也就是人工归类给这张图的类别。
            标签为空时，这是唯一还能知道"这张图本该是什么"的线索。
        full_box: 整图框的类别下标。非 None 时 _emit 会**重写**标签而不是链接原文件。
    """

    image: Path
    label: Path
    class_id: int | None
    group: str
    full_box: int | None = None


def first_class_id(label_path: Path) -> int | None:
    """取标签里第一个框的类别下标；空标签（无框）返回 None。"""
    for line in label_path.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if line:
            return int(float(line.split()[0]))
    return None


def collect(source: Path) -> list[Sample]:
    """配对 images/ 与 labels/。"""
    images_root, labels_root = source / "images", source / "labels"
    if not images_root.is_dir() or not labels_root.is_dir():
        raise FileNotFoundError(f"{source} 下缺少 images/ 或 labels/，先跑 scripts/generate_boxes.py")

    samples = []
    missing = []
    for image in sorted(p for p in images_root.rglob("*") if p.suffix.lower() in IMAGE_SUFFIXES):
        relative = image.relative_to(images_root)
        label = labels_root / relative.with_suffix(".txt")
        if not label.exists():
            missing.append(image)
            continue
        group = relative.parent.name  # 顶层就是图片时为 ""
        samples.append(Sample(image, label, first_class_id(label), group))
    if missing:
        raise FileNotFoundError(
            f"有 {len(missing)} 张图没有对应的标签，例如 {missing[0]}；"
            "产出目录可能不完整，重新跑一次 generate_boxes.py"
        )
    if not samples:
        raise ValueError(f"{source} 下没有找到任何图片")
    return samples


def class_of(sample: Sample, names: Sequence[str]) -> str:
    """样本在输出里归属的类别名：优先看标签内容，空标签才退回来源目录名。

    以标签为准，重命名后的文件名就永远不会与标签里的 class_id 矛盾。
    """
    if sample.class_id is not None:
        return names[sample.class_id]
    if sample.full_box is not None:
        return names[sample.full_box]
    return sample.group or "background"


def assign_names(samples: Sequence[Sample], names: Sequence[str]) -> dict[Path, str]:
    """按类别重新编号，返回 {源图片: 新文件名主干}，如 `blackdot_0042`。

    序号在类内按源路径排序：同一份来源重跑，同一张图拿到的名字不变，
    已经做过的人工复核记录还能对得上。
    """
    by_class: dict[str, list[Sample]] = defaultdict(list)
    for sample in samples:
        by_class[class_of(sample, names)].append(sample)
    width = max(4, len(str(max((len(v) for v in by_class.values()), default=0))))
    mapping = {}
    for name, items in by_class.items():
        for number, sample in enumerate(sorted(items, key=lambda s: str(s.image)), start=1):
            mapping[sample.image] = f"{name}_{number:0{width}d}"
    return mapping


DOCUMENT_ROOT_ENV = "LABEL_STUDIO_LOCAL_FILES_DOCUMENT_ROOT"


def expose_under_document_root(out_dir: Path, document_root: Path | None) -> tuple[Path, Path]:
    """确定 Label Studio 从 DOCUMENT_ROOT 下哪个路径访问输出目录，返回 (document_root, 访问路径)。

    - 没有 DOCUMENT_ROOT：输出目录自己当 DOCUMENT_ROOT，访问路径就是它；
    - 输出目录已在 DOCUMENT_ROOT 之下：直接用；
    - 否则在 DOCUMENT_ROOT 下建一个同名软链接指向输出目录。

    软链接可用：Label Studio（1.23.0 核对过）拼路径用 abspath/normpath，存储路径校验也是
    纯字符串的前缀比较，都不解析软链接。所有比较都要在"未解析"的路径上做，否则软链接
    会被解析到 DOCUMENT_ROOT 外面，反倒对不上。
    """
    out_dir = Path(os.path.abspath(out_dir))
    if document_root is None:
        return out_dir, out_dir
    root = Path(os.path.abspath(document_root))
    if not root.is_dir():
        raise RuntimeError(f"DOCUMENT_ROOT {root} 不存在或不是目录")
    if root == out_dir:
        # Label Studio 拒绝与 DOCUMENT_ROOT 相同的存储路径，但存储填的是 <root>/images，没问题
        return root, out_dir
    if root in out_dir.parents:
        return root, out_dir

    link = root / out_dir.name
    if link.is_symlink():
        if Path(os.path.abspath(os.path.join(link.parent, os.readlink(link)))) != out_dir:
            raise RuntimeError(f"{link} 已是指向别处的软链接（{os.readlink(link)}），请先确认后删除")
    elif link.exists():
        raise RuntimeError(
            f"{link} 已存在且不是软链接（可能是之前手动复制的一份）。"
            f"删除它后重跑，或者用 --out {link} 直接把材料生成到 DOCUMENT_ROOT 里"
        )
    else:
        link.symlink_to(out_dir, target_is_directory=True)
    return root, link


def build_label_studio_tasks(
    out_dir: Path, num_images: int, document_root: Path, served_dir: Path
) -> dict:
    """在输出目录上调用 label-studio-converter，生成 tasks.json 与 labeling config。

    converter 的 YOLO 导入只认 `<input>/{images,labels,classes.txt}`，且 os.listdir 不递归——
    这正是本脚本产出的扁平结构。

    图片 URL 是 `/data/local-files/?d=<served_dir 相对 document_root>/images/<图名>`：
    Label Studio 把 d 拼在 DOCUMENT_ROOT 后面找文件，相对路径错了就是满屏 404。

    框以 predictions 导入：它们是 SHAP / Grad-CAM 生成的，不是人工标注。审核员修改提交
    之后才成为 annotation，Label Studio 导出时也只导出 annotation，没审过的框不会混进去。
    """
    try:
        from label_studio_sdk.converter.imports.yolo import convert_yolo_to_ls
    except ImportError as exc:
        raise RuntimeError("需要 label-studio（uv sync 后可用）") from exc

    tasks_file = out_dir / "tasks.json"
    relative = served_dir.relative_to(document_root).as_posix()
    image_dir = "images" if relative == "." else f"{relative}/images"
    image_root_url = "/data/local-files/?d=" + urllib.parse.quote(image_dir, safe="/")
    # converter 会往 stdout 打一份缺了环境变量的导入说明，还会刷一串 root logger 的 INFO；
    # 静音它们，由 main() 打印带实际路径的完整步骤
    root_logger = logging.getLogger()
    previous_level = root_logger.level
    root_logger.setLevel(logging.WARNING)
    try:
        with contextlib.redirect_stdout(io.StringIO()):
            # 不传 image_dims：图片尺寸不统一，传了会把同一个宽高写进每一条标注
            convert_yolo_to_ls(
                input_dir=str(out_dir),
                out_file=str(tasks_file),
                out_type="predictions",
                image_root_url=image_root_url,
                image_ext=".png",
            )
    finally:
        root_logger.setLevel(previous_level)
    tasks = json.loads(tasks_file.read_text(encoding="utf-8"))
    if len(tasks) != num_images:
        raise RuntimeError(
            f"label-studio-converter 只生成了 {len(tasks)} 个任务，应为 {num_images}；"
            f"检查 {out_dir / 'images'}"
        )
    return {
        "tasks": len(tasks),
        "document_root": str(document_root),
        "storage_path": str(served_dir / "images"),
        "image_root_url": image_root_url,
        "tasks_file": str(tasks_file.resolve()),
        "label_config": str((out_dir / "tasks.label_config.xml").resolve()),
    }


def parse_empty_policy(tokens: Sequence[str]) -> tuple[str, dict[str, str]]:
    """解析 --empty-policy：一个可选的全局默认，加若干 `类别=策略`。

        --empty-policy drop normal=keep black=full
        -> ("drop", {"normal": "keep", "black": "full"})
    """
    default, per_class = "keep", {}
    for token in tokens:
        if "=" in token:
            name, _, policy = token.partition("=")
            if policy not in EMPTY_POLICIES:
                raise ValueError(f"未知策略 {policy!r}（类别 {name}）；可选 {EMPTY_POLICIES}")
            per_class[name] = policy
        elif token in EMPTY_POLICIES:
            default = token
        else:
            raise ValueError(
                f"看不懂 --empty-policy 的 {token!r}；"
                f"要么是全局默认（{'/'.join(EMPTY_POLICIES)}），要么是 `类别=策略`"
            )
    return default, per_class


def apply_empty_policy(
    samples: Sequence[Sample], default: str, per_class: dict[str, str], names: Sequence[str]
) -> tuple[list[Sample], dict[str, tuple[str, int]]]:
    """按来源类别处置空标签，返回 (保留下来的样本, {类别: (策略, 数量)})。

    空标签在 YOLO 里的语义是"这张图什么都没有"——一个明确的负样本断言，而不是
    "有缺陷但没定位出来"。所以来自缺陷类别的空标签如果照搬，等于在教模型抑制真实缺陷。
    """
    index = {name: i for i, name in enumerate(names)}
    unknown = {
        name for name in per_class if name not in index
    }
    if unknown:
        raise ValueError(
            f"--empty-policy 里的类别 {sorted(unknown)} 不在 data.yaml 中；可选 {list(names)}"
        )

    kept: list[Sample] = []
    stats: dict[str, tuple[str, int]] = {}
    for sample in samples:
        if sample.class_id is not None:
            kept.append(sample)
            continue
        policy = per_class.get(sample.group, default)
        counted = stats.get(sample.group, (policy, 0))
        stats[sample.group] = (policy, counted[1] + 1)
        if policy == "drop":
            continue
        if policy == "full":
            if sample.group not in index:
                raise ValueError(
                    f"{sample.group!r} 要用整图框，但它不在 data.yaml 的类别里；可选 {list(names)}"
                )
            kept.append(sample._replace(full_box=index[sample.group]))
        else:
            kept.append(sample)
    return kept, stats


def main() -> int:
    global logger
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument("--source", required=True, help="generate_boxes.py 的产出目录")
    parser.add_argument(
        "--out", default=None, help="输出目录；缺省 runs/<source 目录名>_label_studio",
    )
    parser.add_argument(
        "--empty-policy", nargs="*", default=[], metavar="POLICY|类别=POLICY",
        help=(
            "怎么处置没有框的图（标签为空）。可写一个全局默认，再按类别覆盖，如 "
            "`--empty-policy drop normal=keep black=full`。"
            "keep=当负样本保留（默认）、drop=丢弃、full=退化成整图框。"
            "注意空标签在 YOLO 里的意思是『此图无缺陷』，缺陷类的空标签走到训练会教模型抑制真实缺陷"
        ),
    )
    parser.add_argument(
        "--document-root", default=os.environ.get(DOCUMENT_ROOT_ENV),
        help=(
            f"Label Studio 的 DOCUMENT_ROOT，缺省读环境变量 {DOCUMENT_ROOT_ENV}。"
            "图片 URL 按它生成；输出目录不在它下面时会在里面建同名软链接"
        ),
    )
    parser.add_argument(
        "--keep-names", action="store_true",
        help="保留源文件名，不按 <类别>_<序号> 重新编号（源文件名需去掉后缀后全局唯一）",
    )
    parser.add_argument(
        "--overwrite", action="store_true",
        help="输出目录已存在时先清空。Label Studio 复核期间从这里读图，复核中不要用",
    )
    parser.add_argument("--dry-run", action="store_true", help="只打印计划")
    args = parser.parse_args()

    source = Path(args.source)
    out_dir = Path(args.out) if args.out else Path("runs") / f"{source.name}_label_studio"
    logger = setup_logging()

    try:
        samples = collect(source)
    except (FileNotFoundError, ValueError) as exc:
        logger.error("%s", exc)
        return 1

    meta_path = source / "data.yaml"
    if not meta_path.exists():
        logger.error("找不到 %s", meta_path)
        return 1
    meta = yaml.safe_load(meta_path.read_text(encoding="utf-8"))
    names = [meta["names"][i] for i in sorted(meta["names"])]

    logger.info("来源: %s（%d 张，%d 类）", source, len(samples), len(names))
    try:
        default_policy, per_class_policy = parse_empty_policy(args.empty_policy)
        samples, empty_stats = apply_empty_policy(samples, default_policy, per_class_policy, names)
    except ValueError as exc:
        logger.error("%s", exc)
        return 1

    total_empty = sum(count for _, count in empty_stats.values())
    if total_empty:
        logger.info("无框图片 %d 张，按来源类别处置：", total_empty)
        for group, (policy, count) in sorted(empty_stats.items(), key=lambda kv: -kv[1][1]):
            logger.info("  %-14s %5d 张 -> %s", group or "(顶层)", count, EMPTY_POLICY_DESC[policy])
        # 只对"默认落到 keep"的类别告警——显式写了 `类别=keep` 的是用户的判断，不再唠叨
        as_negative = sorted(
            g for g, (policy, _) in empty_stats.items()
            if policy == "keep" and g not in per_class_policy
        )
        if as_negative:
            logger.warning(
                "以上 %s 的空标签会以『此图无缺陷』的身份导入，复核时不补框的话会这样进入训练。"
                "如果其中某些类其实是有缺陷、只是没能定位出来，它们会教模型抑制真实缺陷——"
                "用 --empty-policy 逐类改成 drop 或 full。",
                "/".join(as_negative),
            )

    counts = Counter(
        names[s.class_id] if s.class_id is not None
        else names[s.full_box] if s.full_box is not None
        else "(无框)"
        for s in samples
    )
    logger.info("%-14s %8s", "类别", "数量")
    for name in names + ["(无框)"]:
        if counts[name]:
            logger.info("%-14s %8d", name, counts[name])

    to_convert = sum(not is_png(sample.image) for sample in samples)
    if to_convert:
        logger.info(
            "其中 %d 张不是 PNG，会无损转成 PNG（无法硬链接，占额外磁盘）；其余 %d 张硬链接",
            to_convert, len(samples) - to_convert,
        )

    if args.dry_run:
        logger.info("输出: %s", out_dir)
        logger.info(
            "DOCUMENT_ROOT: %s", args.document_root or f"未设置 {DOCUMENT_ROOT_ENV}，用输出目录本身"
        )
        logger.info("--dry-run：未写入任何文件")
        return 0

    if out_dir.exists() and any(out_dir.iterdir()):
        if not args.overwrite:
            logger.error("%s 已存在且非空；确认要覆盖请加 --overwrite", out_dir)
            return 1
        shutil.rmtree(out_dir)

    renamed = None if args.keep_names else assign_names(samples, names)
    modes: Counter[str] = Counter()
    seen: dict[str, Path] = {}
    name_rows = []
    for sample in samples:
        image = sample.image
        stem = image.stem if renamed is None else renamed[image]
        output_name = f"{stem}.png"
        # 保留原名时，拍平并统一后缀后不同目录下的 a.png 与 a.bmp 会撞成同一个 a.png；
        # 重新编号时只有类别名仅大小写不同才会撞。按小写比较：macOS 默认文件系统
        # 不区分大小写，A.png 与 a.png 也会互相覆盖。
        key = output_name.lower()
        if key in seen:
            hint = (
                "源数据集的文件名（去掉后缀后）需要全局唯一，"
                "去掉 --keep-names 按类别重新编号，"
                "或用 scripts/export_clusters.py --prefix-parent 重新导出"
                if renamed is None
                else "data.yaml 里有只差大小写的类别名，请先改成互不相同"
            )
            logger.error("输出文件名冲突：%s 与 %s 都会落成 %s。%s", seen[key], image, output_name, hint)
            return 1
        seen[key] = image
        modes[materialize_png(image, out_dir / "images" / output_name)] += 1
        name_rows.append((output_name, class_of(sample, names), str(image.resolve())))
        destination = out_dir / "labels" / f"{stem}.txt"
        if sample.full_box is None:
            link_or_copy(sample.label, destination)
        else:
            # 整图框：内容与源标签不同，只能新写一份，不能链接
            destination.parent.mkdir(parents=True, exist_ok=True)
            destination.write_text(f"{sample.full_box} 0.5 0.5 1.0 1.0\n", encoding="utf-8")

    write_classes_txt(out_dir / "classes.txt", names)
    try:
        document_root, served_dir = expose_under_document_root(
            out_dir, Path(args.document_root) if args.document_root else None
        )
        if served_dir != Path(os.path.abspath(out_dir)):
            logger.info("DOCUMENT_ROOT 下的访问路径: %s -> %s", served_dir, out_dir)
        label_studio = build_label_studio_tasks(out_dir, len(samples), document_root, served_dir)
    except RuntimeError as exc:
        logger.error("%s", exc)
        return 1
    # utf-8-sig：带 BOM，Excel / Numbers 直接打开中文源路径不乱码
    with (out_dir / "name_map.csv").open("w", encoding="utf-8-sig", newline="") as handle:
        writer = csv.writer(handle)
        writer.writerow(["name", "class", "source"])
        writer.writerows(sorted(name_rows, key=lambda row: (row[1], row[0])))
    (out_dir / "summary.json").write_text(
        json.dumps(
            {
                "source": str(source.resolve()),
                "num_images": len(samples),
                "empty_labels": {
                    group: {"policy": policy, "count": count}
                    for group, (policy, count) in sorted(empty_stats.items())
                },
                "empty_policy": {"default": default_policy, **per_class_policy},
                "renamed": renamed is not None,
                "label_studio": label_studio,
                "per_class": dict(counts),
                "link_modes": dict(modes),
            },
            ensure_ascii=False,
            indent=2,
        ),
        encoding="utf-8",
    )
    logger.info("图片落盘方式: %s（hard=硬链接, copy=跨盘复制, convert=转成 PNG）", dict(modes))
    logger.info("完成 -> %s", out_dir)
    # label-studio-converter 自带的提示漏了环境变量（它在循环里改写了 image_root_url，
    # 与默认值比较时就不相等了），这里给出带实际路径的完整步骤
    logger.info(
        "Label Studio 任务 %d 个，导入步骤：\n"
        "  1) export LABEL_STUDIO_LOCAL_FILES_SERVING_ENABLED=true\n"
        "     export LABEL_STUDIO_LOCAL_FILES_DOCUMENT_ROOT=%s\n"
        "     uv run label-studio start\n"
        "  2) 新建项目 -> Labeling Setup -> Code，粘贴 %s\n"
        "  3) Settings -> Cloud Storage -> Add Source Storage -> Local files，\n"
        "     绝对路径填 %s。只保存，不要点 Sync（会给每张图再建一个空任务）\n"
        "  4) 导入 %s\n"
        "图片 URL 形如 %s/<图名>.png；复核期间 Label Studio 从 %s 读图，不要删除或 --overwrite",
        label_studio["tasks"], label_studio["document_root"], label_studio["label_config"],
        label_studio["storage_path"], label_studio["tasks_file"],
        label_studio["image_root_url"], out_dir.resolve(),
    )
    return 0


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