import argparse
from omegaconf import OmegaConf
import json
import os
from typing import List, Dict
from PIL import Image as PILImage

import hashlib
from datetime import datetime

import cv2
import numpy as np

def generate_image_name(image_path, prompt, gen_num=0):
    # 提取图片名称（不带扩展名）
    image_name, _ = os.path.splitext(os.path.basename(image_path))
    
    # 对 prompt 进行哈希处理，生成唯一标识
    prompt_hash = hashlib.md5(prompt.encode()).hexdigest()
    
    
    # 组合生成文件名
    return f"{image_name}_{prompt_hash}_{gen_num}.jpg"

def extract_subject(image_path, mask_path):
    """
    保留mask中黑色主体区域，将背景变为白色
    参数:
        image (numpy.ndarray): 输入图像（BGR格式）
        mask (numpy.ndarray): 掩膜图像（单通道或三通道，主体为黑/背景为白）
    返回:
        numpy.ndarray: 处理后只保留主体的图像
    """
    # 确保mask是单通道
    image = PILImage.open(image_path).convert("RGB")
    mask = PILImage.open(mask_path).convert("L")
    # mask = ImageOps.invert(mask)
    white = PILImage.new("RGB", image.size, (255, 255, 255))
    masked_image = PILImage.composite(image, white, mask)
    return masked_image


def overwrite_config(config, opts):
    if not opts:
        return config

    assert len(opts) % 2 == 0
    for i in range(0, len(opts), 2):
        k = opts[i]
        v = opts[i+1]

        # format v based on suffix, e.g. True@bool
        if "@" in v:
            _v, dtype = v.split("@")
            if dtype == "int":
                v = int(_v)
            elif dtype == "float":
                v = float(_v)
            elif dtype == "bool":
                _v = _v.strip().lower()
                assert _v in ["true", "false"]
                v = (_v == "true")
            else:
                raise NotImplementedError
        # guass the type
        else:
            try:
                v = int(v)
            except ValueError:
                try:
                    v = float(v)
                except ValueError:
                    if v.lower() in ["true", "false"]:
                        v = (v.lower() == "true")
                    else:
                        pass

        keys = k.split(".")
        parent = config
        for key in keys[:-1]:
            parent = parent[key]
        final_key = keys[-1]
        print(f"overwrite {k}: {parent[final_key]} -> {v} (type: {type(v)})")
        parent[final_key] = v 
    return config


def parse_args_and_config(add_custom_args=None):
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", type=str, default="./configs/train.yaml")
    parser.add_argument('--opts', type=str, nargs="+")
    if add_custom_args is not None:
        add_custom_args(parser)
    args = parser.parse_args()

    if args.config[-5:] == ".yaml":
        config = OmegaConf.load(args.config)
    else:
        raise ValueError("Do not support this format config file")

    config = overwrite_config(config, args.opts)
    config = OmegaConf.to_container(config, resolve=True)
    return args, config

