"""用 LocalAI qwen3.5-9b 多模态模型自动标注图片。

用法：python auto_label.py
输入：datasets/images/<class>/*.jpg
输出：datasets/dataset/{train,val}/{images,labels} (YOLO 格式) + data.yaml
"""
import os
import json
import re
import base64
import shutil
import sys
import urllib.request
from PIL import Image, ImageDraw

# LocalAI 服务器配置
LOCAL_AI_URL = "http://192.168.3.210:18080/v1/chat/completions"
MODEL_NAME = "qwen3.5-9b"

# 标注配置
KEEP_PER_CLASS = 100
MIN_CONF = 0.5    # 自动标注最低置信度
MIN_BOX = 0.03    # 框面积占比下限

BASE = os.path.dirname(__file__)
IMG_DIR = os.path.join(BASE, "datasets", "images")
OUT_DIR = os.path.join(BASE, "datasets", "dataset")
PREVIEW_DIR = os.path.join(BASE, "datasets", "preview")

CLASSES = ["pheasant"]
PROMPTS = {
    "pheasant": "pheasant",
}

# 标注提示词模板
ANNOTATION_PROMPT = """你是一个目标检测助手。请检测图片中【{target}】动物。

输出格式（严格只输出这一行，不要其他文字）：
pheasant 0.92 0.1 0.2 0.3 0.4

字段说明：
- 第一个数字：置信度（0-1）
- 后四个数字：边界框坐标 x1 y1 x2 y2（归一化 0-1）

如果没找到目标，只输出：NOT_FOUND
"""


def encode_image_to_base64(path: str) -> str:
    """将图片路径转换为 base64 数据"""
    with open(path, "rb") as f:
        return base64.b64encode(f.read()).decode("utf-8")


def call_localai(image_path: str, prompt: str) -> dict:
    """调用 LocalAI qwen3.5-9b 进行标注"""
    image_b64 = encode_image_to_base64(image_path)

    data = {
        "model": MODEL_NAME,
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt},
                    {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_b64}"}},
                ],
            }
        ],
        "temperature": 0.1,
    }

    try:
        with urllib.request.urlopen(LOCAL_AI_URL, data=json.dumps(data).encode("utf-8")) as resp:
            if resp.status == 200:
                result = json.loads(resp.read().decode("utf-8"))
                return {"success": True, "data": result}
            else:
                return {"success": False, "error": f"HTTP {resp.status}"}
    except Exception as e:
        return {"success": False, "error": str(e)}


def parse_annotation(response: dict) -> list:
    """解析 LocalAI qwen3.5-9b 返回的标注结果（文本格式）"""
    if not response.get("success"):
        return []

    content = response.get("choices", [{}])[0].get("message", {}).get("content", "")

    # 如果没找到目标
    if "NOT_FOUND" in content or "未找到" in content:
        return []

    # 解析文本格式：类别 置信度 x1 y1 x2 y2
    # 例如：pheasant 0.92 0.1 0.2 0.3 0.4
    parts = content.strip().split()
    if len(parts) < 5:
        return []

    cls_name = parts[0].lower()
    try:
        confidence = float(parts[1])
        x1, y1, x2, y2 = [float(v) for v in parts[2:6]]
    except ValueError:
        return []

    # 确保坐标顺序
    if x1 > x2 or y1 > y2:
        x1, x2 = x2, x1
        y1, y2 = y2, y1

    return [{
        "class": cls_name,
        "confidence": confidence,
        "bbox": [x1, y1, x2, y2],
    }]


def main() -> None:
    for split in ("train", "val"):
        for sub in ("images", "labels"):
            os.makedirs(os.path.join(OUT_DIR, split, sub), exist_ok=True)
    os.makedirs(PREVIEW_DIR, exist_ok=True)

    for cls in CLASSES:
        src_dir = os.path.join(IMG_DIR, cls)
        files = sorted(f for f in os.listdir(src_dir) if f.endswith(".jpg"))
        print(f"[label] {cls}: {len(files)} 张", flush=True)

        kept = []
        for f in files:
            path = os.path.join(src_dir, f)
            image = Image.open(path).convert("RGB")
            w, h = image.size

            # 调用 LocalAI 进行标注
            prompt = ANNOTATION_PROMPT.format(target=cls)
            result = call_localai(path, prompt)
            annotations = parse_annotation(result)

            # 取置信度最高的标注
            best = None
            best_score = 0
            for ann in annotations:
                if ann.get("class") != cls:
                    continue
                score = ann.get("confidence", 0)
                if score > best_score:
                    best_score = score
                    best = ann

            if best is None:
                print(f"  {f}: ✗ (未找到目标)", flush=True)
                continue

            if best_score < MIN_CONF:
                print(f"  {f}: ✗ (置信度 {best_score:.2f} < {MIN_CONF})", flush=True)
                continue

            x1, y1, x2, y2 = best.get("bbox", [0, 0, 1, 1])
            kept.append((path, (best_score, x1, y1, x2, y2)))
            print(f"  {f}: ✓ ({best_score:.2f})", flush=True)

        kept.sort(key=lambda t: t[1][0], reverse=True)
        kept = kept[:KEEP_PER_CLASS]
        avg = sum(k[1][0] for k in kept) / max(len(kept), 1)
        print(f"  → 保留 {len(kept)} 张（平均置信度 {avg:.2f}）", flush=True)

        for i, (path, (score, x1, y1, x2, y2)) in enumerate(kept):
            split = "train" if i % 10 else "val"
            dst_img = os.path.join(OUT_DIR, split, "images", f"{cls}_{i:03d}.jpg")
            shutil.copy(path, dst_img)

            # 计算 YOLO 格式坐标
            w, h = Image.open(dst_img).size
            cx = (x1 + x2) / 2 / w
            cy = (y1 + y2) / 2 / h
            bw = (x2 - x1) / w
            bh = (y2 - y1) / h

            label_file = os.path.join(OUT_DIR, split, "labels", f"{cls}_{i:03d}.txt")
            with open(label_file, "w") as f:
                f.write(f"{CLASSES.index(cls)} {cx:.6f} {cy:.6f} {bw:.6f} {bh:.6f}\n")

            # 生成预览图
            img = Image.open(dst_img).convert("RGB")
            draw = ImageDraw.Draw(img)
            draw.rectangle([x1, y1, x2, y2], outline="#E53935", width=3)
            draw.text((x1, max(y1 - 14, 0)), f"{cls} {score:.0%}", fill="#E53935")
            img.save(os.path.join(PREVIEW_DIR, f"{cls}_{i:03d}.jpg"), "JPEG", quality=85)

    # 生成 data.yaml
    with open(os.path.join(OUT_DIR, "data.yaml"), "w") as f:
        f.write(f"path: {OUT_DIR}\n")
        f.write("train: train/images\nval: val/images\n")
        f.write(f"names: {CLASSES}\n")
    print("[done] 数据集就绪:", OUT_DIR, flush=True)


if __name__ == "__main__":
    sys.exit(main())