"""
이미지 처리 유틸리티
iopaint와 호환되는 이미지 변환 및 처리 함수들을 제공합니다.
"""
import base64
import io
import logging
from typing import Tuple, Optional, Union
import numpy as np
import cv2
from PIL import Image

logger = logging.getLogger(__name__)


def encode_image_to_base64(image: Union[np.ndarray, Image.Image], format: str = "PNG") -> str:
    """이미지를 base64로 인코딩합니다."""
    try:
        if isinstance(image, np.ndarray):
            # numpy 배열을 PIL Image로 변환
            if image.dtype != np.uint8:
                image = (image * 255).astype(np.uint8)
            
            if len(image.shape) == 3 and image.shape[2] == 3:
                # BGR to RGB
                image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            
            pil_image = Image.fromarray(image)
        else:
            pil_image = image
        
        # PIL Image를 base64로 인코딩
        buffer = io.BytesIO()
        pil_image.save(buffer, format=format)
        img_str = base64.b64encode(buffer.getvalue()).decode('utf-8')
        
        return img_str
    except Exception as e:
        logger.error(f"이미지 base64 인코딩 실패: {e}")
        raise


def decode_base64_to_image(base64_string: str, gray: bool = False) -> Tuple[np.ndarray, Optional[np.ndarray], dict, str]:
    """base64 문자열을 이미지로 디코딩합니다 (iopaint 호환)."""
    alpha_channel = None  # 변수 초기화
    
    try:
        # 데이터 URL 형식인 경우, 접두사 제거 (e.g., "data:image/png;base64,")
        if "," in base64_string:
            base64_string = base64_string.split(',', 1)[1]

        # base64 디코딩
        image_data = base64.b64decode(base64_string)
        
        # PIL Image로 로드
        pil_image = Image.open(io.BytesIO(image_data))
        
        # 이미지 정보 추출
        info = {
            "format": pil_image.format,
            "mode": pil_image.mode,
            "size": pil_image.size,
            "parameters": ""
        }
        
        # 확장자 결정
        ext = pil_image.format.lower() if pil_image.format else "png"
        
        # numpy 배열로 변환
        if gray:
            # 그레이스케일로 변환
            if pil_image.mode in ('RGBA', 'LA', 'P'):
                pil_image = pil_image.convert('L')
            image_array = np.array(pil_image)
        else:
            # RGB로 변환
            if pil_image.mode == 'RGBA':
                # RGBA를 RGB로 변환하고 알파 채널 분리
                rgb_image = pil_image.convert('RGB')
                alpha_channel = np.array(pil_image.split()[-1])
                image_array = np.array(rgb_image)
            elif pil_image.mode == 'P':
                # 팔레트 이미지를 RGB로 변환
                image_array = np.array(pil_image.convert('RGB'))
                alpha_channel = None
            else:
                # RGB 이미지
                image_array = np.array(pil_image)
                alpha_channel = None
        
        # BGR로 변환 (OpenCV 호환)
        if len(image_array.shape) == 3 and image_array.shape[2] == 3:
            image_array = cv2.cvtColor(image_array, cv2.COLOR_RGB2BGR)
        
        return image_array, alpha_channel, info, ext
        
    except Exception as e:
        logger.error(f"base64 이미지 디코딩 실패: {e}")
        raise


def concat_alpha_channel(image: np.ndarray, alpha_channel: Optional[np.ndarray]) -> np.ndarray:
    """이미지와 알파 채널을 결합합니다."""
    if alpha_channel is None:
        return image
    
    try:
        # BGR to RGB
        if len(image.shape) == 3 and image.shape[2] == 3:
            rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        else:
            rgb_image = image
        
        # RGBA 이미지 생성
        if len(rgb_image.shape) == 3:
            rgba_image = np.dstack((rgb_image, alpha_channel))
        else:
            rgba_image = np.dstack((rgb_image, rgb_image, rgb_image, alpha_channel))
        
        return rgba_image
        
    except Exception as e:
        logger.error(f"알파 채널 결합 실패: {e}")
        return image


def pil_to_bytes(image: Image.Image, ext: str = "PNG", quality: int = 100, infos: dict = None) -> bytes:
    """PIL Image를 바이트로 변환합니다."""
    try:
        buffer = io.BytesIO()
        
        if ext.lower() in ['jpg', 'jpeg']:
            image.save(buffer, format='JPEG', quality=quality, optimize=True)
        else:
            image.save(buffer, format=ext.upper(), optimize=True)
        
        return buffer.getvalue()
        
    except Exception as e:
        logger.error(f"PIL Image를 바이트로 변환 실패: {e}")
        raise


def numpy_to_bytes(image: np.ndarray, ext: str = "PNG", quality: int = 100) -> bytes:
    """numpy 배열을 바이트로 변환합니다."""
    try:
        # numpy 배열을 PIL Image로 변환
        if image.dtype != np.uint8:
            image = (image * 255).astype(np.uint8)
        
        if len(image.shape) == 3 and image.shape[2] == 3:
            # BGR to RGB
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        
        pil_image = Image.fromarray(image)
        return pil_to_bytes(pil_image, ext, quality)
        
    except Exception as e:
        logger.error(f"numpy 배열을 바이트로 변환 실패: {e}")
        raise


def adjust_mask(mask: np.ndarray, kernel_size: int = 5, operate: str = "dilate") -> np.ndarray:
    """마스크를 조정합니다 (dilate/erode)."""
    try:
        if kernel_size <= 0:
            return mask
        
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size))
        
        if operate.lower() == "dilate":
            result = cv2.dilate(mask, kernel, iterations=1)
        elif operate.lower() == "erode":
            result = cv2.erode(mask, kernel, iterations=1)
        else:
            logger.warning(f"알 수 없는 마스크 연산: {operate}")
            return mask
        
        return result
        
    except Exception as e:
        logger.error(f"마스크 조정 실패: {e}")
        return mask


def gen_frontend_mask(mask: np.ndarray) -> np.ndarray:
    """프론트엔드용 마스크를 생성합니다."""
    try:
        # 마스크를 0-255 범위로 정규화
        if mask.dtype != np.uint8:
            mask = (mask * 255).astype(np.uint8)
        
        # 이진화
        _, binary_mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)
        
        return binary_mask
        
    except Exception as e:
        logger.error(f"Error in gen_frontend_mask: {e}")
        return image # fallback to original mask


def fill_transparent_background_with_white(image: np.ndarray) -> np.ndarray:
    """
    RGBA 이미지의 투명 배경을 흰색으로 채웁니다.
    """
    if image.shape[2] == 4:
        # 알파 채널을 마스크로 사용
        alpha_channel = image[:, :, 3]
        mask = alpha_channel > 0

        # RGB 채널과 흰색 배경 생성
        rgb_image = image[:, :, :3]
        white_background = np.full(rgb_image.shape, 255, dtype=np.uint8)

        # 마스크를 사용하여 전경과 배경을 조합
        result = np.where(mask[:, :, np.newaxis], rgb_image, white_background)
        return result
    else:
        # 이미 RGBA가 아니면 그대로 반환
        return image


def preprocess_for_bria_rmbg(image: np.ndarray) -> np.ndarray:
    """
    Bria RMBG 모델을 위한 전처리 (허깅페이스 공식 utilities.py 기반)
    
    preprocessor_config.json:
    - do_rescale: true, rescale_factor: 0.00392156862745098 (1/255)
    - do_normalize: true, image_mean: [0.5, 0.5, 0.5], image_std: [1, 1, 1]  
    - do_resize: true, size: {"width": 1024, "height": 1024}
    """
    # 그레이스케일 처리 (채널이 부족한 경우)
    if len(image.shape) < 3:
        image = image[:, :, np.newaxis]
    
    # BGR to RGB 변환 (OpenCV는 BGR 순서)
    if len(image.shape) == 3 and image.shape[2] >= 3:
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    
    # 1024x1024로 bilinear 리사이즈 후 uint8로 변환
    image = cv2.resize(image, (1024, 1024), interpolation=cv2.INTER_LINEAR).astype(np.uint8)
    
    # rescale: 0-255 -> 0-1 (rescale_factor = 1/255)
    image = image.astype(np.float32) / 255.0
    
    # normalize: (image - mean) / std
    # mean=[0.5, 0.5, 0.5], std=[1.0, 1.0, 1.0]
    image = (image - 0.5) / 1.0
    
    # (H, W, C) -> (1, C, H, W) 변환 (배치 차원 추가 + 채널 우선)
    image = np.transpose(image, (2, 0, 1))  # (C, H, W)
    image = np.expand_dims(image, axis=0)   # (1, C, H, W)
    
    return image


def postprocess_for_bria_rmbg(result_tensor: np.ndarray, original_size: Tuple[int, int]) -> np.ndarray:
    """
    Bria RMBG 모델 출력을 후처리 (허깅페이스 공식 utilities.py 기반)
    
    torch 코드:
    result = torch.squeeze(F.interpolate(result, size=im_size, mode='bilinear'), 0)
    ma = torch.max(result)
    mi = torch.min(result)
    result = (result-mi)/(ma-mi)
    im_array = (result*255).permute(1,2,0).cpu().data.numpy().astype(np.uint8)
    """
    # (1, 1, H, W) -> (1, H, W) squeeze(0차원 제거)
    result = np.squeeze(result_tensor, axis=0)  # (1, H, W)
    
    # bilinear interpolation으로 원본 크기 복원
    # cv2.resize는 (W, H) 순서 주의
    result = cv2.resize(result, (original_size[1], original_size[0]), interpolation=cv2.INTER_LINEAR)
    
    # Min-max 정규화: (result-mi)/(ma-mi)
    ma = result.max()
    mi = result.min()
    if ma > mi:
        result = (result - mi) / (ma - mi)
    else:
        # 모든 값이 같은 경우 (드물지만 방어코드)
        result = np.zeros_like(result)
    
    # 0-255 범위로 변환 후 uint8
    result = (result * 255).astype(np.uint8)
    
    return result


def resize_image(image: np.ndarray, target_size: Tuple[int, int], keep_aspect: bool = True) -> np.ndarray:
    """이미지 크기를 조정합니다."""
    try:
        if keep_aspect:
            # 종횡비 유지하면서 크기 조정
            h, w = image.shape[:2]
            target_h, target_w = target_size
            
            # 종횡비 계산
            aspect = w / h
            target_aspect = target_w / target_h
            
            if aspect > target_aspect:
                # 너비에 맞춤
                new_w = target_w
                new_h = int(target_w / aspect)
            else:
                # 높이에 맞춤
                new_h = target_h
                new_w = int(target_h * aspect)
            
            resized = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LANCZOS4)
            
            # 패딩으로 목표 크기 맞춤
            result = np.zeros((target_h, target_w, 3) if len(image.shape) == 3 else (target_h, target_w), dtype=image.dtype)
            y_offset = (target_h - new_h) // 2
            x_offset = (target_w - new_w) // 2
            result[y_offset:y_offset+new_h, x_offset:x_offset+new_w] = resized
            
            return result
        else:
            # 종횡비 무시하고 크기 조정
            return cv2.resize(image, target_size, interpolation=cv2.INTER_LANCZOS4)
            
    except Exception as e:
        logger.error(f"이미지 크기 조정 실패: {e}")
        return image


def validate_image_size(image: np.ndarray, max_size: int) -> bool:
    """이미지 크기가 제한을 초과하지 않는지 확인합니다."""
    try:
        h, w = image.shape[:2]
        pixels = h * w
        max_pixels = max_size * max_size
        
        return pixels <= max_pixels
        
    except Exception as e:
        logger.error(f"이미지 크기 검증 실패: {e}")
        return False