"""
3D 모델 및 이미지 데이터 로더
GLB 파일과 이미지 파일을 로드하고 전처리하는 기능을 제공합니다.
"""

import numpy as np
import trimesh
import open3d as o3d
from PIL import Image
import cv2
from typing import Dict, List, Tuple, Optional
import os


class DataLoader:
    """3D 모델과 이미지 데이터를 로드하고 전처리하는 클래스"""
    
    def __init__(self):
        self.supported_3d_formats = ['.glb', '.gltf', '.obj', '.ply', '.stl']
        self.supported_image_formats = ['.png', '.jpg', '.jpeg', '.bmp', '.tiff']
    
    def load_glb_model(self, path: str) -> Dict:
        """
        GLB 파일을 로드하고 메시 정보를 추출합니다.
        
        Args:
            path (str): GLB 파일 경로
            
        Returns:
            Dict: 3D 모델 정보 (vertices, faces, textures, materials)
        """
        try:
            # trimesh를 사용하여 GLB 파일 로드
            scene = trimesh.load(path)
            
            # 메시 정보 추출
            if hasattr(scene, 'geometry'):
                # Scene인 경우 첫 번째 메시 사용
                mesh = list(scene.geometry.values())[0]
            else:
                # 단일 메시인 경우
                mesh = scene
            
            model_info = {
                'vertices': np.array(mesh.vertices),
                'faces': np.array(mesh.faces),
                'textures': None,
                'materials': None,
                'bounding_box': mesh.bounds,
                'center': mesh.centroid,
                'scale': mesh.scale,
                'file_path': path
            }
            
            # 텍스처 정보가 있는 경우
            if hasattr(mesh.visual, 'material'):
                model_info['materials'] = mesh.visual.material
            
            # UV 좌표가 있는 경우
            if hasattr(mesh.visual, 'uv'):
                model_info['uv_coords'] = mesh.visual.uv
            
            return model_info
            
        except Exception as e:
            raise ValueError(f"GLB 파일 로드 실패: {path}, 오류: {str(e)}")
    
    def load_reference_image(self, path: str) -> np.ndarray:
        """
        참조 이미지를 로드하고 전처리합니다.
        
        Args:
            path (str): 이미지 파일 경로
            
        Returns:
            np.ndarray: 전처리된 이미지 배열 (H, W, C)
        """
        try:
            # PIL을 사용하여 이미지 로드
            image = Image.open(path)
            
            # RGB로 변환 (RGBA인 경우)
            if image.mode == 'RGBA':
                image = image.convert('RGB')
            elif image.mode != 'RGB':
                image = image.convert('RGB')
            
            # numpy 배열로 변환
            image_array = np.array(image)
            
            return image_array
            
        except Exception as e:
            raise ValueError(f"이미지 로드 실패: {path}, 오류: {str(e)}")
    
    def extract_mesh_info(self, model: Dict) -> Dict:
        """
        3D 모델에서 메시 정보를 추출합니다.
        
        Args:
            model (Dict): 3D 모델 정보
            
        Returns:
            Dict: 추출된 메시 정보
        """
        vertices = model['vertices']
        faces = model['faces']
        
        # 메시 통계 계산
        mesh_info = {
            'num_vertices': len(vertices),
            'num_faces': len(faces),
            'bounding_box_size': np.max(vertices, axis=0) - np.min(vertices, axis=0),
            'volume': self._calculate_mesh_volume(vertices, faces),
            'surface_area': self._calculate_surface_area(vertices, faces),
            'center_of_mass': np.mean(vertices, axis=0)
        }
        
        return mesh_info
    
    def create_ground_truth(self, image: np.ndarray) -> Dict:
        """
        이미지에서 Ground Truth 정보를 생성합니다.
        
        Args:
            image (np.ndarray): 입력 이미지
            
        Returns:
            Dict: Ground Truth 정보
        """
        # 이미지 크기 정보
        height, width = image.shape[:2]
        
        # 간단한 객체 감지를 위한 전처리
        gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
        
        # 엣지 검출
        edges = cv2.Canny(gray, 50, 150)
        
        # 윤곽선 찾기
        contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        # 바운딩 박스 생성
        bounding_boxes = []
        for contour in contours:
            if cv2.contourArea(contour) > 100:  # 최소 면적 필터
                x, y, w, h = cv2.boundingRect(contour)
                bounding_boxes.append({
                    'bbox': [x, y, x + w, y + h],
                    'area': w * h,
                    'confidence': 1.0
                })
        
        ground_truth = {
            'image_size': (width, height),
            'bounding_boxes': bounding_boxes,
            'num_objects': len(bounding_boxes),
            'image_histogram': np.histogram(gray, bins=256)[0]
        }
        
        return ground_truth
    
    def mesh_to_pointcloud(self, vertices: np.ndarray, num_points: int = 10000) -> np.ndarray:
        """
        메시를 점군으로 변환합니다.
        
        Args:
            vertices (np.ndarray): 메시 정점
            num_points (int): 생성할 점의 개수
            
        Returns:
            np.ndarray: 점군 데이터
        """
        if len(vertices) >= num_points:
            # 정점이 충분한 경우 랜덤 샘플링
            indices = np.random.choice(len(vertices), num_points, replace=False)
            return vertices[indices]
        else:
            # 정점이 부족한 경우 중복 허용하여 샘플링
            indices = np.random.choice(len(vertices), num_points, replace=True)
            return vertices[indices]
    
    def _calculate_mesh_volume(self, vertices: np.ndarray, faces: np.ndarray) -> float:
        """메시의 부피를 계산합니다."""
        try:
            mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
            return mesh.volume
        except:
            return 0.0
    
    def _calculate_surface_area(self, vertices: np.ndarray, faces: np.ndarray) -> float:
        """메시의 표면적을 계산합니다."""
        try:
            mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
            return mesh.surface_area
        except:
            return 0.0
    
    def validate_file(self, file_path: str, file_type: str = 'auto') -> bool:
        """
        파일 유효성을 검사합니다.
        
        Args:
            file_path (str): 파일 경로
            file_type (str): 파일 타입 ('3d', 'image', 'auto')
            
        Returns:
            bool: 파일 유효성
        """
        if not os.path.exists(file_path):
            return False
        
        _, ext = os.path.splitext(file_path.lower())
        
        if file_type == 'auto':
            return (ext in self.supported_3d_formats or 
                   ext in self.supported_image_formats)
        elif file_type == '3d':
            return ext in self.supported_3d_formats
        elif file_type == 'image':
            return ext in self.supported_image_formats
        
        return False