#!/usr/bin/env python
"""用指定权重给 Label Studio 项目里的所有任务批量生成预测。

    uv run python scripts/push_predictions.py --project 39 \\
        --weights runs/detect/current.pt

适合重训完成后，一次性把全部样本的预标注刷新成新模型的结果。

与网页端的 **Retrieve Predictions** 相比：那条路要经过 ML backend，Label Studio 会把选中的
任务一股脑发过去；本脚本在本地按 `--batch` 分批推理，再用批量导入接口
（`POST /api/projects/<id>/import/predictions`）一次提交一批，显存和请求体都可控，
也不依赖后端服务在跑。

需要一个 API token（网页端 Account & Settings → Access Token）：

    export LABEL_STUDIO_API_KEY=xxxx

新版 Label Studio 默认只认 **JWT 个人访问令牌**（组织设置里 `legacy_api_tokens_enabled=false`）。
脚本会自动识别：令牌形如 JWT（两个点分三段）时先用 `/api/token/refresh/` 换取 access token
并以 `Bearer` 发送，否则按旧式 `Token` 发送。

图片路径从任务的 `/data/local-files/?d=...` 还原，所以 `--document-root` 必须和
Label Studio 的 `LABEL_STUDIO_LOCAL_FILES_DOCUMENT_ROOT` 一致。

**预测是追加的，不会覆盖。** 同一个任务可以有多版预测，Label Studio 在项目设置里按
model_version 选用哪一版。重复跑同一个版本会堆重复预测：`--replace` 先删掉同名版本
（换了权重时一般删不到东西，很快），`--replace-all` 则一次删光项目里的所有预测。
"""

from __future__ import annotations

import argparse
import os
import sys
import urllib.parse
from pathlib import Path

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

from convert import to_prediction  # noqa: E402

from utils import setup_logging  # noqa: E402
from utils.device import DeviceNotAvailableError, describe_device, resolve_device  # noqa: E402

logger = None


class LabelStudio:
    """Label Studio REST API 的最小封装。"""

    def __init__(self, url: str, token: str) -> None:
        import requests

        self.url = url.rstrip("/")
        self.session = requests.Session()
        self.token = token
        self._authorize()

    def _authorize(self) -> None:
        """设置 Authorization 头；JWT 令牌会换成 access token。

        新版默认关闭了旧式 `Token xxx`（组织设置 legacy_api_tokens_enabled=false），
        个人访问令牌是一个 JWT refresh token，要先换成 access token 再以 Bearer 发送。
        """
        if self.token.count(".") != 2:                 # 不是 JWT，按旧式 Token 处理
            self.session.headers["Authorization"] = f"Token {self.token}"
            return
        response = self.session.post(
            f"{self.url}/api/token/refresh/", json={"refresh": self.token}, timeout=60
        )
        if response.status_code >= 400:
            raise RuntimeError(
                f"用个人访问令牌换取 access token 失败（HTTP {response.status_code}）："
                f"{response.text[:200]}"
            )
        self.session.headers["Authorization"] = f"Bearer {response.json()['access']}"

    def request(self, method: str, path: str, **kwargs):
        """带自动续期的请求。

        access token 默认只有 5 分钟有效期（Label Studio 没有覆盖 SimpleJWT 的默认值），
        而推完上万条预测远不止 5 分钟，所以碰到 401 就用 refresh token 再换一个再试一次。
        """
        kwargs.setdefault("timeout", 120)
        response = self.session.request(method, f"{self.url}{path}", **kwargs)
        if response.status_code == 401:
            self._authorize()
            response = self.session.request(method, f"{self.url}{path}", **kwargs)
        if response.status_code >= 400:
            # 带上响应体：Label Studio 的 4xx/5xx 常把真正原因写在里面，
            # raise_for_status 只给一行 URL，排查时毫无用处
            raise RuntimeError(
                f"{method} {path} 失败 HTTP {response.status_code}：{response.text[:500]}"
            )
        return response

    def tasks(self, project: int, page_size: int = 500) -> list[dict]:
        """列出项目里的全部任务（只取 id 与 data，分页拉完）。"""
        collected, page = [], 1
        while True:
            response = self.request(
                "GET", "/api/tasks/",
                params={"project": project, "page": page, "page_size": page_size,
                        "fields": "task_only"},
            )
            payload = response.json()
            # 不同版本的分页字段不一样：data manager 用 tasks/total，DRF 默认用 results/count
            batch = payload.get("tasks", payload.get("results", payload)) if isinstance(
                payload, dict
            ) else payload
            if not batch:
                break
            collected.extend(batch)
            total = payload.get("total", payload.get("count")) if isinstance(payload, dict) else None
            if total is not None and len(collected) >= int(total):
                break
            page += 1
        return collected

    def delete_predictions(self, project: int, model_version: str) -> int:
        """删掉该项目下指定版本的预测，返回删除条数。

        `/api/predictions/` **不分页**：Label Studio 的 REST_FRAMEWORK 没设
        DEFAULT_PAGINATION_CLASS，PAGE_SIZE 形同虚设，它一次返回全部预测并忽略 page 参数。
        早先按页翻会翻不到头（实测翻到第 1798 页才因令牌过期报错），所以这里只请求一次。
        """
        payload = self.request("GET", "/api/predictions/", params={"project": project}).json()
        items = payload.get("results", []) if isinstance(payload, dict) else payload
        removed = 0
        for item in items:
            if item.get("model_version") == model_version:
                self.request("DELETE", f"/api/predictions/{item['id']}/")
                removed += 1
        return removed

    def delete_all_predictions(self, project: int) -> dict:
        """用 Data Manager 的批量操作一次删光项目里的所有预测（不分版本）。"""
        return self.request(
            "POST", "/api/dm/actions/",
            params={"id": "delete_tasks_predictions", "project": project},
            json={"selectedItems": {"all": True, "excluded": []}},
        ).json()

    def import_predictions(self, project: int, predictions: list[dict]) -> None:
        """批量导入，服务端自己再分批入库。"""
        self.request(
            "POST", f"/api/projects/{project}/import/predictions",
            json=predictions, timeout=600,
        )


def local_path(image_url: str, document_root: Path) -> Path:
    """把任务里的图片 URL 还原成磁盘路径。

    只处理本地存储（`/data/local-files/?d=相对路径`）——上传或云存储的图得走下载，
    那种情况用网页端的 Retrieve Predictions 更省事。
    """
    if "?d=" not in image_url:
        raise ValueError(
            f"不是本地存储的图片 URL：{image_url}；本脚本只支持 /data/local-files/?d=..."
        )
    relative = urllib.parse.unquote(image_url.split("?d=", 1)[1])
    return document_root / relative


def main() -> int:
    global logger
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument("--project", type=int, required=True, help="Label Studio 项目 id")
    parser.add_argument(
        "--weights", default="runs/detect/current.pt",
        help="权重；缺省用重训脚本维护的 current.pt 软链接",
    )
    parser.add_argument("--url", default=os.environ.get("LABEL_STUDIO_URL", "http://localhost:8080"))
    parser.add_argument(
        "--token", default=os.environ.get("LABEL_STUDIO_API_KEY", ""),
        help="API token，缺省读环境变量 LABEL_STUDIO_API_KEY",
    )
    parser.add_argument(
        "--document-root",
        default=os.environ.get("LABEL_STUDIO_LOCAL_FILES_DOCUMENT_ROOT", "/home/deep/labeling"),
        help="必须与 Label Studio 的同名环境变量一致",
    )
    parser.add_argument("--model-version", default=None, help="缺省按权重文件名与修改时间生成")
    parser.add_argument("--conf", type=float, default=0.25)
    parser.add_argument("--imgsz", type=int, default=640)
    parser.add_argument("--device", default="auto")
    parser.add_argument("--batch", type=int, default=16, help="每次推理几张")
    parser.add_argument("--chunk", type=int, default=500, help="每次上传几条预测")
    parser.add_argument("--limit", type=int, default=0, help="只处理前 N 个任务（试运行用）")
    parser.add_argument(
        "--replace", action="store_true",
        help="先删掉同一 model_version 的旧预测；否则会追加，同版本堆两份",
    )
    parser.add_argument(
        "--replace-all", action="store_true",
        help="先删光项目里的**所有**预测（不分版本），用 Data Manager 的批量操作，一次请求",
    )
    parser.add_argument("--dry-run", action="store_true", help="只推理不上传")
    args = parser.parse_args()

    logger = setup_logging()
    if not args.token:
        logger.error("缺少 API token：export LABEL_STUDIO_API_KEY=... 或用 --token 传入")
        return 1

    weights = Path(args.weights)
    if not weights.exists():
        logger.error("找不到权重 %s", weights)
        return 1
    model_version = args.model_version or f"{weights.resolve().stem}-{int(weights.stat().st_mtime)}"

    try:
        device = resolve_device(args.device)
    except DeviceNotAvailableError as exc:
        logger.error("%s", exc)
        return 1

    try:
        client = LabelStudio(args.url, args.token)
    except RuntimeError as exc:
        logger.error("%s", exc)
        return 1
    logger.info("拉取项目 %d 的任务列表…", args.project)
    tasks = client.tasks(args.project)
    if args.limit:
        tasks = tasks[: args.limit]
    logger.info(
        "任务 %d 个 | 权重 %s | 版本 %s | 设备 %s",
        len(tasks), weights, model_version, describe_device(device),
    )
    if not tasks:
        return 0

    document_root = Path(args.document_root)
    try:
        pairs = [(task["id"], local_path(task["data"]["image"], document_root)) for task in tasks]
    except (KeyError, ValueError) as exc:
        logger.error("%s", exc)
        return 1
    missing = [path for _, path in pairs if not path.exists()]
    if missing:
        logger.error(
            "有 %d 张图在磁盘上找不到，例如 %s；--document-root 是否与 Label Studio 一致（当前 %s）",
            len(missing), missing[0], document_root,
        )
        return 1

    if args.replace_all and not args.dry_run:
        detail = client.delete_all_predictions(args.project)
        logger.info("删除项目里的全部旧预测：%s", detail.get("detail", detail))
    elif args.replace and not args.dry_run:
        removed = client.delete_predictions(args.project, model_version)
        logger.info("删除同版本（%s）旧预测 %d 条", model_version, removed)

    from PIL import Image
    from ultralytics import YOLO

    model = YOLO(str(weights))
    classes = model.names
    device_arg = device.index or 0 if device.type == "cuda" else device.type

    pending: list[dict] = []
    uploaded = with_boxes = total_boxes = 0

    def flush() -> None:
        nonlocal pending, uploaded
        if not pending or args.dry_run:
            pending = []
            return
        client.import_predictions(args.project, pending)
        uploaded += len(pending)
        pending = []

    for start in range(0, len(pairs), args.batch):
        chunk = pairs[start:start + args.batch]
        results = model.predict(
            source=[str(path) for _, path in chunk], imgsz=args.imgsz, conf=args.conf,
            device=device_arg, verbose=False,
        )
        for (task_id, path), result in zip(chunk, results, strict=True):
            with Image.open(path) as image:
                width, height = image.size
            boxes = [
                (classes[int(class_id)], xyxyn, float(score))
                for xyxyn, class_id, score in zip(
                    result.boxes.xyxyn.tolist(), result.boxes.cls.tolist(),
                    result.boxes.conf.tolist(), strict=True,
                )
            ]
            prediction = to_prediction(boxes, width, height, "label", "image", model_version)
            prediction["task"] = task_id
            pending.append(prediction)
            total_boxes += len(boxes)
            with_boxes += bool(boxes)
        if len(pending) >= args.chunk:
            try:
                flush()
            except RuntimeError as exc:
                logger.error("%s", exc)
                return 1
        processed = start + len(chunk)
        if processed % 1000 < args.batch:
            logger.info("处理 [%d/%d]，已上传 %d", processed, len(pairs), uploaded)
    try:
        flush()
    except RuntimeError as exc:
        logger.error("%s", exc)
        return 1

    logger.info(
        "完成：%d 个任务，%d 个有框（%d 个框），上传 %d 条预测，版本 %s",
        len(pairs), with_boxes, total_boxes, uploaded, model_version,
    )
    if args.dry_run:
        logger.info("--dry-run：未上传")
    else:
        logger.info(
            "在项目设置 Annotation → Prelabeling 里把预测版本选成 %s 即可用于预标注", model_version
        )
    return 0


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