# -*- coding: utf-8 -*-
"""
마스크 모듈 - OCR 결과를 기반으로 마스크 생성 및 처리 (라이브러리화)
확장, 블러 등 basic 마스크 처리만 지원
"""

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

class MaskModule:
    def __init__(self):
        print("마스크 모듈 초기화 완료")
    def create_masks(self, image_path: str, ocr_results: List[Dict], expansion_size: int = 10, blur_size: int = 15, mask_option: str = "basic") -> np.ndarray:
        image = cv2.imread(image_path)
        if image is None:
            print(f"이미지를 읽을 수 없습니다: {image_path}")
            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']
            expanded_poly = self.expand_polygon(polygon, offset=5)
            cv2.fillPoly(mask, [expanded_poly], 255)
        processed_mask = self.process_mask(mask, expansion_size, blur_size)
        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
    def create_masks_with_save(self, image_path: str, ocr_results: List[Dict], output_dir: str = None, expansion_size: int = 10, blur_size: int = 15, mask_option: str = "basic") -> tuple:
        mask = self.create_masks(image_path, ocr_results, expansion_size, blur_size, mask_option)
        mask_image_path = None
        if output_dir and mask is not None:
            os.makedirs(output_dir, exist_ok=True)
            base_name = os.path.splitext(os.path.basename(image_path))[0]
            mask_filename = f"{base_name}_mask_combined.jpg"
            mask_image_path = os.path.join(output_dir, mask_filename)
            cv2.imwrite(mask_image_path, mask)
        return mask, mask_image_path