from pathlib import Path
from typing import Dict

import cv2
import numpy as np

from app.config import settings
from app.utils import (
    draw_results,
    now_timestamp_minute,
    safe_stem,
    save_json,
    write_image_xml,
)


def make_unique_result_dir(parent_dir: Path, timestamp: str, original_filename: str) -> tuple[Path, str, str]:
    """
    创建单次图片识别结果目录。

    目录名：
        时间戳_原文件名stem

    例如：
        runs/images/202605181431_DJI_20250820151801_0006_Z/

    返回：
        result_dir, result_base_name, original_stem
    """
    original_stem = safe_stem(original_filename)
    base_name = f"{timestamp}_{original_stem}"

    result_dir = parent_dir / base_name

    if not result_dir.exists():
        result_dir.mkdir(parents=True, exist_ok=True)
        return result_dir, base_name, original_stem

    idx = 1
    while True:
        candidate_base = f"{base_name}_{idx:02d}"
        candidate_dir = parent_dir / candidate_base
        if not candidate_dir.exists():
            candidate_dir.mkdir(parents=True, exist_ok=True)
            return candidate_dir, candidate_base, original_stem
        idx += 1


def save_image_outputs(
    img_bgr: np.ndarray,
    original_filename: str,
    predict_output: Dict,
) -> Dict:
    """
    图片结果自动保存到：
        runs/images/时间戳_原图片名/

    目录内部生成：
        原图片名.jpg
        时间戳_原图片名.json
        时间戳_原图片名.xml

    示例：
        runs/images/202605181431_DJI_20250820151801_0006_Z/
        ├── DJI_20250820151801_0006_Z.jpg
        ├── 202605181431_DJI_20250820151801_0006_Z.json
        └── 202605181431_DJI_20250820151801_0006_Z.xml
    """
    images_root = Path(settings.output_dir) / settings.image_output_subdir
    images_root.mkdir(parents=True, exist_ok=True)

    timestamp = now_timestamp_minute()
    result_dir, result_base_name, original_stem = make_unique_result_dir(
        parent_dir=images_root,
        timestamp=timestamp,
        original_filename=original_filename,
    )

    # 图片文件名不带时间戳，按你的要求保留原图片名 stem
    image_path = result_dir / f"{original_stem}.jpg"

    # JSON 和 XML 带时间戳
    json_path = result_dir / f"{result_base_name}.json"
    xml_path = result_dir / f"{result_base_name}.xml"

    results = predict_output.get("results", [])

    annotated = draw_results(img_bgr, results)
    cv2.imwrite(str(image_path), annotated)

    json_data = {
        "success": True,
        "type": "image",
        "original_filename": original_filename,
        "result_dir": str(result_dir),
        "saved_image": str(image_path),
        "saved_json": str(json_path),
        "saved_xml": str(xml_path),
        "image_width": predict_output.get("image_width"),
        "image_height": predict_output.get("image_height"),
        "num_objects": predict_output.get("num_objects"),
        "results": results,
    }

    save_json(json_path, json_data)

    write_image_xml(
        xml_path=xml_path,
        image_filename=image_path.name,
        image_path=str(image_path),
        width=int(predict_output.get("image_width", img_bgr.shape[1])),
        height=int(predict_output.get("image_height", img_bgr.shape[0])),
        results=results,
    )

    return {
        "result_dir": str(result_dir),
        "image": str(image_path),
        "json": str(json_path),
        "xml": str(xml_path),
    }