#!/usr/bin/env python
"""把训练好的权重按统一命名归档，并导出 OpenVINO 格式。

    uv run python scripts/export_openvino.py --weights runs/detect/film_20260924_0817/weights/best.pt

产出：

```
runs/detect/models/
├── film_<边长>_<时间戳>.pt
└── film_<边长>_<时间戳>_openvino_model/    # .xml / .bin / metadata.yaml
```

**边长从权重里读**（ultralytics 把训练参数存在 checkpoint 里），不用手填，也就不会出现
文件名写 128、实际是 160 训练的情况。

**时间戳默认取训练目录名里的那一串**（`film_20260924_0817` → `20260924_0817`），
这样同一次训练无论何时归档，名字都一样、可重复；取不到才退回当前时刻。
"""

from __future__ import annotations

import argparse
import re
import shutil
import sys
from datetime import datetime
from pathlib import Path

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

from utils import setup_logging  # noqa: E402

logger = None

#: 训练目录名里的时间戳，如 film_20260924_0817 / film128_20260930_0415
STAMP_PATTERN = re.compile(r"(\d{8}_\d{4})")


def stamp_for(weights: Path) -> str:
    """从权重路径里找训练时间戳；找不到就用当前时刻。"""
    for part in reversed(weights.resolve().parts):
        found = STAMP_PATTERN.search(part)
        if found:
            return found.group(1)
    return datetime.now().strftime("%Y%m%d_%H%M")


def main() -> int:
    global logger
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument("--weights", required=True, help="训练产出的 .pt")
    parser.add_argument(
        "--out", default="runs/detect/models", help="归档目录，默认 runs/detect/models",
    )
    parser.add_argument("--stamp", default=None, help="覆盖时间戳，缺省从权重路径推断")
    parser.add_argument("--size", type=int, default=0, help="覆盖输入边长，缺省从权重里读")
    parser.add_argument("--prefix", default="film", help="文件名前缀")
    parser.add_argument(
        "--no-openvino", action="store_true", help="只归档 .pt，不导出 OpenVINO",
    )
    parser.add_argument("--overwrite", action="store_true")
    args = parser.parse_args()

    logger = setup_logging()
    weights = Path(args.weights)
    if not weights.exists():
        logger.error("找不到权重 %s", weights)
        return 1

    from ultralytics import YOLO

    model = YOLO(str(weights))
    size = args.size or int(model.overrides.get("imgsz") or 0)
    if not size:
        logger.error("权重里读不到 imgsz，请用 --size 指定")
        return 1

    stamp = args.stamp or stamp_for(weights)
    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)
    target = out_dir / f"{args.prefix}_{size}_{stamp}.pt"

    if target.exists() and not args.overwrite:
        logger.error("%s 已存在；确认要覆盖请加 --overwrite", target)
        return 1
    shutil.copy2(weights, target)
    logger.info("归档权重: %s（%.1f MB）", target, target.stat().st_size / 1e6)

    if args.no_openvino:
        return 0

    # 从归档后的文件导出，这样 OpenVINO 目录名与 .pt 同名
    exported = YOLO(str(target)).export(format="openvino", imgsz=size, batch=1, dynamic=False)
    logger.info("OpenVINO: %s", exported)
    logger.info("输入固定为 %dx%d；部署时切片的 overlap 需与训练时一致", size, size)
    return 0


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