import cv2
import numpy as np
import easyocr
from typing import List, Tuple, Dict, Optional
import re
import torch

class OCRModule2:
    """EasyOCR을 활용한 중국어 텍스트 인식 모듈"""
    
    def __init__(self, use_gpu: bool = False, ocr_backend: str = 'easyocr'):
        """
        OCR 모듈 초기화
        
        Args:
            use_gpu: GPU 사용 여부
            ocr_backend: 'easyocr' 또는 'paddleocr' 중 선택
        """
        self.ocr_backend = ocr_backend.lower()
        if self.ocr_backend == 'easyocr':
            import easyocr
            self.ocr = easyocr.Reader(
                ['ch_sim', 'en'],  # 중국어 간체, 영어 지원
                gpu=use_gpu
            )
        elif self.ocr_backend == 'paddleocr':
            try:
                from paddleocr import PaddleOCR
            except ImportError:
                raise ImportError('PaddleOCR가 설치되어 있지 않습니다. 설치: pip install paddleocr==2.7.4')
            self.ocr = PaddleOCR(use_angle_cls=True, lang='ch', use_gpu=use_gpu)
        else:
            raise ValueError("ocr_backend는 'easyocr' 또는 'paddleocr'만 지원합니다.")
        
    def has_chinese_text(self, image: np.ndarray) -> bool:
        """
        이미지에 중국어 텍스트가 포함되어 있는지 확인
        
        Args:
            image: 입력 이미지 (numpy array)
            
        Returns:
            중국어 텍스트 포함 여부 (bool)
        """
        if self.ocr_backend == 'easyocr':
            results = self.ocr.readtext(image)
            for result in results:
                text = result[1]
                if self._is_chinese(text):
                    return True
            return False
        elif self.ocr_backend == 'paddleocr':
            results = self.ocr.ocr(image, cls=True)
            for line in results:
                for res in line:
                    text = res[1][0]
                    if self._is_chinese(text):
                        return True
            return False
        else:
            raise ValueError("ocr_backend는 'easyocr' 또는 'paddleocr'만 지원합니다.")
    
    def extract_chinese_text(self, image: np.ndarray) -> Dict:
        """
        이미지에서 중국어 텍스트 추출
        
        Args:
            image: 입력 이미지 (numpy array)
            
        Returns:
            Dict containing:
            - texts: 추출된 중국어 텍스트 리스트
            - polygons: 각 텍스트의 폴리곤 좌표
            - masks: 텍스트 영역 마스크
            - confidences: 인식 신뢰도
        """
        try:
            if self.ocr_backend == 'easyocr':
                results = self.ocr.readtext(image)
                texts = []
                polygons = []
                masks = []
                confidences = []
                for result in results:
                    polygon = result[0]
                    text = result[1]
                    confidence = result[2]
                    if self._is_chinese(text):
                        texts.append(text)
                        polygons.append(polygon)
                        confidences.append(confidence)
                        mask = self._create_mask_from_polygon(image.shape, polygon)
                        masks.append(mask)
                return {
                    'texts': texts,
                    'polygons': polygons,
                    'masks': masks,
                    'confidences': confidences
                }
            elif self.ocr_backend == 'paddleocr':
                results = self.ocr.ocr(image, cls=True)
                texts = []
                polygons = []
                masks = []
                confidences = []
                for line in results:
                    for res in line:
                        polygon = res[0]  # 4점 좌표
                        text = res[1][0]  # 텍스트
                        confidence = res[1][1]  # 신뢰도
                        if self._is_chinese(text):
                            texts.append(text)
                            polygons.append(polygon)
                            confidences.append(confidence)
                            mask = self._create_mask_from_polygon(image.shape, polygon)
                            masks.append(mask)
                return {
                    'texts': texts,
                    'polygons': polygons,
                    'masks': masks,
                    'confidences': confidences
                }
            else:
                raise ValueError("ocr_backend는 'easyocr' 또는 'paddleocr'만 지원합니다.")
        except Exception as e:
            print(f"중국어 텍스트 추출 중 오류: {e}")
            return {
                'texts': [],
                'polygons': [],
                'masks': [],
                'confidences': []
            }
    
    def _is_chinese(self, text: str) -> bool:
        """
        텍스트가 중국어인지 확인
        
        Args:
            text: 확인할 텍스트
            
        Returns:
            중국어 여부
        """
        chinese_pattern = re.compile(r'[\u4e00-\u9fff]+')
        return bool(chinese_pattern.search(text))
    
    def _create_mask_from_polygon(self, image_shape: Tuple[int, int, int], polygon: List[List[int]]) -> np.ndarray:
        """
        폴리곤으로부터 마스크 생성
        
        Args:
            image_shape: 이미지 크기 (H, W, C)
            polygon: 폴리곤 좌표
            
        Returns:
            마스크 배열
        """
        mask = np.zeros(image_shape[:2], dtype=np.uint8)
        polygon_array = np.array(polygon, dtype=np.int32)
        cv2.fillPoly(mask, [polygon_array], 255)
        return mask
    
    def get_text_properties(self, polygon: List[List[int]]) -> Dict:
        """
        텍스트 영역의 속성 정보 추출
        
        Args:
            polygon: 폴리곤 좌표
            
        Returns:
            텍스트 속성 정보 (위치, 크기, 각도 등)
        """
        polygon_array = np.array(polygon)
        
        # 바운딩 박스 계산
        x_coords = polygon_array[:, 0]
        y_coords = polygon_array[:, 1]
        
        x_min, x_max = int(np.min(x_coords)), int(np.max(x_coords))
        y_min, y_max = int(np.min(y_coords)), int(np.max(y_coords))
        
        width = x_max - x_min
        height = y_max - y_min
        
        # 중심점 계산
        center_x = (x_min + x_max) // 2
        center_y = (y_min + y_max) // 2
        
        # 회전 각도 계산 (첫 번째와 두 번째 점 기준)
        if len(polygon_array) >= 2:
            dx = polygon_array[1][0] - polygon_array[0][0]
            dy = polygon_array[1][1] - polygon_array[0][1]
            angle = np.arctan2(dy, dx) * 180 / np.pi
        else:
            angle = 0
        
        return {
            'bbox': (x_min, y_min, x_max, y_max),
            'center': (center_x, center_y),
            'width': width,
            'height': height,
            'angle': angle
        }