import os
import cv2
import numpy as np


def create_output_folder(source_folder, suffix="_processed"):
    """创建输出文件夹"""
    base_name = os.path.basename(source_folder)
    parent_dir = os.path.dirname(source_folder)
    output_folder = os.path.join(parent_dir, f"{base_name}{suffix}")

    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    return output_folder


def is_image_file(file_path):
    """检查文件是否为图片"""
    valid_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.gif']
    ext = os.path.splitext(file_path)[1].lower()
    return ext in valid_extensions


def resize_to_fit(image, max_width, max_height):
    """调整图片大小以适应指定尺寸"""
    height, width = image.shape[:2]

    # 计算宽高比
    ratio = min(max_width / width, max_height / height)

    # 计算新尺寸
    new_width = int(width * ratio)
    new_height = int(height * ratio)

    # 调整大小
    return cv2.resize(image, (new_width, new_height))


def draw_watermark_regions(image, regions):
    """在图片上绘制水印区域"""
    result = image.copy()

    for region in regions:
        x, y, w, h = region['position']

        # 根据区域类型选择颜色
        if region['type'] == 'text':
            color = (255, 0, 0)  # 蓝色
        elif region['type'] == 'stamp':
            color = (0, 0, 255)  # 红色
        else:
            color = (0, 255, 0)  # 绿色

        # 绘制矩形
        cv2.rectangle(result, (x, y), (x + w, y + h), color, 2)

        # 添加标签
        label = f"{region['type']} ({region['confidence']:.2f})"
        cv2.putText(result, label, (x, y - 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)

    return result