import xml.etree.ElementTree as ET
import cv2
import os
import numpy as np
from typing import List, Dict, Optional, Tuple
from utils.Stomach_Config import Loading_Config

class XMLAnnotationProcessor:
    """
    XML標註檔案處理器
    專門處理包含bounding box資訊的XML檔案，並在對應圖片上繪製邊界框
    """
    
    def __init__(self, dataset_root: str):
        """
        初始化XML處理器
        
        Args:
            dataset_root: 圖片資料集根目錄
            output_folder: 輸出資料夾
        """
        self.dataset_root = dataset_root
        self.box_color = (0, 255, 0)  # 綠色邊界框
        self.text_color = (0, 255, 0)  # 綠色文字
        self.box_thickness = 2
        self.font_scale = 0.5
        self.font = cv2.FONT_HERSHEY_SIMPLEX
    
    def _ensure_output_folder(self, Save_Root: str) -> None:
        """確保輸出資料夾存在"""
        if not os.path.exists(Save_Root):
            os.makedirs(Save_Root)
    
    def parse_xml(self, xml_file_path: str, Label: str) -> Optional[Dict]:
        """
        解析XML檔案並提取所有相關資訊
        
        Args:
            xml_file_path: XML檔案路徑
            
        Returns:
            Dict: 包含檔案資訊和bounding box的字典，解析失敗時返回None
        """
        try:
            tree = ET.parse(xml_file_path)
            root = tree.getroot()
            
            # 提取基本資訊
            filename_element = root.find('filename')
            
            if filename_element is None:
                print(f"找不到path元素在 {xml_file_path}")
                return None
            
            filename = filename_element.text if filename_element is not None else "Unknown"
            Original_Image_Data_Root = os.path.join(self.dataset_root, Label)
            Original_Image_Data_Root = os.path.join(Original_Image_Data_Root, filename)
            
            # 提取圖片尺寸
            size_element = root.find('size')
            width = int(size_element.find('width').text) if size_element is not None else 0
            height = int(size_element.find('height').text) if size_element is not None else 0
            depth = int(size_element.find('depth').text) if size_element is not None else 3
            
            # 提取所有bounding box
            bounding_boxes = []
            objects = root.findall('object')
            
            for obj in objects:
                bndbox = obj.find('bndbox')
                if bndbox is not None:
                    bbox_info = {
                        'name': obj.find('name').text if obj.find('name') is not None else "Unknown",
                        'pose': obj.find('pose').text if obj.find('pose') is not None else "Unspecified",
                        'truncated': int(obj.find('truncated').text) if obj.find('truncated') is not None else 0,
                        'difficult': int(obj.find('difficult').text) if obj.find('difficult') is not None else 0,
                        'xmin': int(bndbox.find('xmin').text),
                        'ymin': int(bndbox.find('ymin').text),
                        'xmax': int(bndbox.find('xmax').text),
                        'ymax': int(bndbox.find('ymax').text)
                    }
                    bounding_boxes.append(bbox_info)
            
            return {
                'filename': filename,
                'image_path': Original_Image_Data_Root,
                'width': width,
                'height': height,
                'depth': depth,
                'bounding_boxes': bounding_boxes
            }
            
        except Exception as e:
            print(f"解析XML檔案 {xml_file_path} 時發生錯誤: {str(e)}")
            return None
    
    def load_image(self, image_path: str) -> Optional[np.ndarray]:
        """
        載入圖片檔案
        
        Args:
            image_path: 圖片檔案路徑
            
        Returns:
            np.ndarray: 圖片陣列，載入失敗時返回None
        """
        if not os.path.exists(image_path):
            print(f"圖片檔案不存在: {image_path}")
            return None
        
        image = cv2.imread(image_path)
        if image is None:
            print(f"無法讀取圖片: {image_path}")
            return None
        
        return image
    
    def draw_bounding_boxes(self, image: np.ndarray, bounding_boxes: List[Dict]) -> np.ndarray:
        """
        創建遮罩圖片：bounding box內保持原圖，外部為黑色
        
        Args:
            image: 圖片陣列
            bounding_boxes: bounding box資訊列表
            
        Returns:
            np.ndarray: 處理後的遮罩圖片陣列
        """
        # 創建黑色背景圖片
        height, width = image.shape[:2]
        result_image = np.zeros((height, width, 3), dtype=np.uint8)
        
        for i, bbox in enumerate(bounding_boxes):
            xmin, ymin = bbox['xmin'], bbox['ymin']
            xmax, ymax = bbox['xmax'], bbox['ymax']
            object_name = bbox['name']
            
            # 確保座標在圖片範圍內
            xmin = max(0, min(xmin, width-1))
            ymin = max(0, min(ymin, height-1))
            xmax = max(0, min(xmax, width-1))
            ymax = max(0, min(ymax, height-1))
            
            # 將bounding box範圍內的原圖複製到結果圖像中
            result_image[ymin:ymax, xmin:xmax] = image[ymin:ymax, xmin:xmax]
            
            print(f"Object {i+1}: {object_name} - 座標: ({xmin}, {ymin}, {xmax}, {ymax})")
        
        return result_image
    
    def save_annotated_image(self, image: np.ndarray, original_filename: str, Annotation_Root : str, Label : str) -> str:
        """
        儲存標註後的圖片
        
        Args:
            image: 標註後的圖片陣列
            original_filename: 原始檔案名稱
            
        Returns:
            str: 儲存的檔案路徑
        """
        output_filename = f"annotated_{original_filename}"
        output_path = os.path.join(Annotation_Root, Label)
        Save_Image_Roots = os.path.join(output_path, output_filename)
        # 確保輸出資料夾存在
        self._ensure_output_folder(output_path)

        cv2.imwrite(Save_Image_Roots, image)
        print(f"已儲存標註圖片至: {Save_Image_Roots}")
        return Save_Image_Roots
    
    def process_single_xml(self, xml_file_path: str, Annotation_Root : str, Label : str) -> Optional[Tuple[np.ndarray, str]]:
        """
        處理單一XML檔案
        
        Args:
            xml_file_path: XML檔案路徑
            
        Returns:
            Tuple[np.ndarray, str]: (標註後的圖片, 輸出路徑)，處理失敗時返回None
        """
        # 解析XML
        xml_data = self.parse_xml(xml_file_path, Label)
        if xml_data is None:
            return None
        
        # 載入圖片
        image = self.load_image(xml_data['image_path'])
        if image is None:
            return None
        
        # 繪製bounding box
        annotated_image = self.draw_bounding_boxes(image, xml_data['bounding_boxes'])
        
        # 儲存結果
        output_path = self.save_annotated_image(annotated_image, xml_data['filename'], Annotation_Root, Label)
        
        return annotated_image, output_path
    
    def process_multiple_xml(self, xml_folder_path: str, Annotation_Root : str, Label : str) -> List[Tuple[str, bool]]:
        """
        批量處理多個XML檔案
        
        Args:
            xml_folder_path: 包含XML檔案的資料夾路徑
            
        Returns:
            List[Tuple[str, bool]]: [(檔案名稱, 處理成功與否), ...]
        """
        if not os.path.exists(xml_folder_path):
            print(f"XML資料夾不存在: {xml_folder_path}")
            return []
        
        xml_files = [f for f in os.listdir(xml_folder_path) if f.endswith('.xml')]
        
        if not xml_files:
            print(f"在 {xml_folder_path} 中找不到XML檔案")
            return []
        
        print(f"找到 {len(xml_files)} 個XML檔案")
        for xml_file in xml_files:            
            try:
                Read_XML_File = os.path.join(xml_folder_path, xml_file)
                self.process_single_xml(Read_XML_File, Annotation_Root, Label)
                print(f"\n處理檔案: {xml_file}")
            except Exception as e:
                print(f"處理 {xml_file} 時發生錯誤: {str(e)}")
                return
    
    def get_bounding_boxes_info(self, xml_file_path: str) -> Optional[Dict]:
        """
        僅提取XML中的bounding box資訊，不進行圖片處理
        
        Args:
            xml_file_path: XML檔案路徑
            
        Returns:
            Dict: 包含檔案資訊和bounding box座標的字典
        """
        return self.parse_xml(xml_file_path)
    
    def set_drawing_style(self, box_color: Tuple[int, int, int] = None, 
                         text_color: Tuple[int, int, int] = None,
                         box_thickness: int = None, 
                         font_scale: float = None) -> None:
        """
        設定繪圖樣式
        
        Args:
            box_color: 邊界框顏色 (B, G, R)
            text_color: 文字顏色 (B, G, R)
            box_thickness: 邊界框粗細
            font_scale: 字體大小
        """
        if box_color is not None:
            self.box_color = box_color
        if text_color is not None:
            self.text_color = text_color
        if box_thickness is not None:
            self.box_thickness = box_thickness
        if font_scale is not None:
            self.font_scale = font_scale
