#!/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

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

**预测是追加的，不会覆盖。** 同一个任务可以有多版预测，Label Studio 在项目设置里按
model_version 选用哪一版。重复跑同一个版本会堆重复预测，用 `--replace` 先删掉同名版本。
"""

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.session.headers["Authorization"] = f"Token {token}"

    def tasks(self, project: int, page_size: int = 500) -> list[dict]:
        """列出项目里的全部任务（只取 id 与 data，分页拉完）。"""
        collected, page = [], 1
        while True:
            response = self.session.get(
                f"{self.url}/api/tasks/",
                params={"project": project, "page": page, "page_size": page_size,
                        "fields": "task_only"},
                timeout=120,
            )
            response.raise_for_status()
            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:
        """删掉该项目下指定版本的预测，返回删除条数。"""
        removed = 0
        page = 1
        while True:
            response = self.session.get(
                f"{self.url}/api/predictions/",
                params={"project": project, "page": page, "page_size": 500},
                timeout=120,
            )
            if response.status_code == 404:      # 翻到尾页，某些版本直接返回 404
                break
            response.raise_for_status()
            payload = response.json()
            items = payload.get("results", payload) if isinstance(payload, dict) else payload
            if not items:
                break
            for item in items:
                if item.get("model_version") == model_version:
                    self.session.delete(
                        f"{self.url}/api/predictions/{item['id']}/", timeout=60
                    ).raise_for_status()
                    removed += 1
            if isinstance(payload, dict) and not payload.get("next"):
                break
            page += 1
        return removed

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


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("--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

    client = LabelStudio(args.url, args.token)
    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 and not args.dry_run:
        removed = client.delete_predictions(args.project, model_version)
        logger.info("删除同版本旧预测 %d 条", 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:
            flush()
        processed = start + len(chunk)
        if processed % 1000 < args.batch:
            logger.info("处理 [%d/%d]，已上传 %d", processed, len(pairs), uploaded)
    flush()

    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())