import cv2
import numpy as np
from typing import List, Dict, Any
from shapely.geometry import Polygon
import logging

class MaskModule:
    def __init__(self, logger, base_dir):
        self.logger = logger
        self.base_dir = base_dir
        self.logger.log("마스크 모듈 초기화 완료", level=logging.INFO)

    def create_masks(self, image_path: str, ocr_results: List[Dict], expansion_size: int = 6, blur_size: int = 7, mask_option: str = "basic") -> np.ndarray:
        image = cv2.imread(image_path)
        if image is None:
            self.logger.log(f"이미지를 읽을 수 없습니다: {image_path}", level=logging.ERROR)
            return None
        height, width = image.shape[:2]
        mask = np.zeros((height, width), dtype=np.uint8)
        for i, result in enumerate(ocr_results, 1):
            polygon = result['polygon']
            try:
                # 좌표가 float일 수 있으므로 안전 캐스팅 + 범위 클램프
                poly_np = np.array(polygon, dtype=np.float32)
                # 이미지 경계 밖 좌표를 안전하게 클램프
                poly_np[:, 0] = np.clip(poly_np[:, 0], 0, width - 1)
                poly_np[:, 1] = np.clip(poly_np[:, 1], 0, height - 1)
                # 확장
                expanded_poly = self.expand_polygon(poly_np.tolist(), offset=5)
                expanded_poly = expanded_poly.astype(np.int32)
                cv2.fillPoly(mask, [expanded_poly], 255)
            except Exception as e:
                self.logger.log(f"마스크 폴리곤 처리 오류({i}): {e}", level=logging.WARNING)
                continue
        processed_mask = self.process_mask(mask, expansion_size, blur_size)

        # # 디버깅용 마스크 저장 (항상 0과 255만 가지는 표준 흑백 마스크로 저장)
        # try:
        #     import os
        #     base_dir = os.path.dirname(image_path)
        #     base_name = os.path.splitext(os.path.basename(image_path))[0]
        #     debug_mask_path = os.path.join(base_dir, f"debug_mask_{base_name}.png")
        #     # 마스크가 0~255 사이의 값이 섞여 있을 수 있으니, 128 기준으로 이진화
        #     mask_to_save = ((processed_mask > 128) * 255).astype('uint8')
        #     cv2.imwrite(debug_mask_path, mask_to_save)
        #     self.logger.log(f"디버깅용 마스크 저장: {debug_mask_path}", level=20)
        # except Exception as e:
        #     self.logger.log(f"디버깅용 마스크 저장 실패: {e}", level=40)
        # return processed_mask

        return processed_mask
    
    def expand_polygon(self, polygon, offset=15):
        poly = Polygon(polygon)
        expanded = poly.buffer(offset)
        if expanded.is_empty:
            return np.array(polygon, dtype=np.int32)
        return np.array(expanded.exterior.coords, dtype=np.int32)
    
    def process_mask(self, mask: np.ndarray, expansion_size: int = 5, blur_size: int = 3) -> np.ndarray:
        processed_mask = mask.copy()
        if expansion_size > 0:
            kernel = np.ones((expansion_size, expansion_size), np.uint8)
            processed_mask = cv2.dilate(processed_mask, kernel, iterations=1)
        if blur_size > 0:
            blur_size = blur_size if blur_size % 2 == 1 else blur_size + 1
            processed_mask = cv2.GaussianBlur(processed_mask, (blur_size, blur_size), 0)
        return processed_mask