import json
import math
from pathlib import Path
from ultralytics import YOLO
import requests
import argparse

from clearml import Task  # 1. ClearML 임포트

# 방법 A) 가장 최신 트렌드인 NMS-Free 임베디드 특화 모델 로드
model = YOLO("./weights/yolo26m.pt")  # Medium 크기 가중치 자동 다운로드

task = Task.init(
    project_name="Normal_Object_Detection", 
    task_name="model-yolo26-human"
)

def xyxy_to_xywh(box: dict) -> dict:
    x1 = math.floor(box["x1"])
    y1 = math.floor(box["y1"])
    x2 = math.ceil(box["x2"])
    y2 = math.ceil(box["y2"])
    return {"x": x1, "y": y1, "w": x2 - x1, "h": y2 - y1}


# 방법 B) 복잡한 구도나 정밀 식별에 강한 어텐션 기반 모델 로드
# model = YOLO("yolo12m.pt")

def main(image_path: str):
    # uv run main.py --image_url "https://acai.ketidev.kr:20443/detect/image/202606/20260619_145116_image.jpg"
    # 이미지 원격 추론 테스트 (COCO 80종 기본 탐지 가능)
    # results = model("https://acai.ketidev.kr:20443/detect/image/202606/20260619_145116_image.jpg")
    results = model(image_path)

    result = results[0]

    image_path = Path(result.path)
    marked_path = image_path.with_name(f"{image_path.stem}_marked{image_path.suffix}")
    result.save(filename=str(marked_path))

    detections = []
    for det in result.summary():
        detections.append({**det, "box": xyxy_to_xywh(det["box"])})

    output = {
        "path": result.path,
        "marked_path": str(marked_path),
        "shape": {"height": result.orig_shape[0], "width": result.orig_shape[1]},
        "speed_ms": result.speed,
        "detections": detections,
    }
    print(json.dumps(output, indent=2, ensure_ascii=False))

    result_data = {"output": output, "status": "PASS"}
    task.upload_artifact(name="final_result", artifact_object=result_data)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--image_url", type=str)
    args = parser.parse_args()

    image_url = args.image_url

    if image_url is None:
        print("Image path is required")
        exit(1)

    response = requests.get(image_url)
    response.raise_for_status()
    image_data = response.content
    with open("image.jpg", "wb") as f:
        f.write(image_data)
    image_path = "image.jpg"

    main(image_path)