import h5py
import numpy as np
import pandas as pd
from pathlib import Path
from typing import Optional
from PIL import Image
import base64
from io import BytesIO

class HDF5Reader:
    def __init__(self):
        self.file_path: Optional[str] = None
        self.hdf5_file: Optional[h5py.File] = None
        self.cached_dataframe: Optional[pd.DataFrame] = None
        
        self.patch_data_array = None  # patch/__data__/data
        self.patch_info = None  # patch/data
        self.patch_rects = None  # patch/rect
        
        self.particle_patch_ids = None  # particle/patch_id
        self.particle_rects = None  # particle/rect
    
    @staticmethod
    def _format_column_name(col_name: str) -> str:
        if col_name.isupper():
            return col_name
        
        words = col_name.split('_')
        formatted = ' '.join(word.capitalize() for word in words)
        return formatted
        
    def open_file(self, file_path: str) -> bool:
        if self.hdf5_file:
            self.hdf5_file.close()
        
        self.cached_dataframe = None
        self.patch_data_array = None
        self.patch_info = None
        self.patch_rects = None 
        self.particle_patch_ids = None
        self.particle_rects = None
        
        self.file_path = file_path
        self.hdf5_file = h5py.File(file_path, 'r')
        
        self._load_patch_info()
        self._load_particle_info()
        
        return True

    
    def _load_patch_info(self):
        try:
            if 'patch' in self.hdf5_file:
                if '__data__' in self.hdf5_file['patch'] and 'data' in self.hdf5_file['patch/__data__']:
                    self.patch_data_array = self.hdf5_file['patch/__data__/data']
                
                if 'data' in self.hdf5_file['patch']:
                    self.patch_info = self.hdf5_file['patch/data']

                if 'rect' in self.hdf5_file['patch']:
                    self.patch_rects = self.hdf5_file['patch/rect']
                        
        except Exception as e:
            import traceback
            traceback.print_exc()
    
    def _load_particle_info(self):
        try:
            if 'particle' in self.hdf5_file:
                if 'patch_id' in self.hdf5_file['particle']:
                    self.particle_patch_ids = self.hdf5_file['particle/patch_id'][:]
                
                if 'rect' in self.hdf5_file['particle']:
                    self.particle_rects = self.hdf5_file['particle/rect']
                    
        except Exception as e:
            import traceback
            traceback.print_exc()
    
    def _get_patch_image(self, patch_id: int) -> Optional[np.ndarray]:
        if self.patch_data_array is None or self.patch_info is None or self.patch_rects is None:
            return None
        
        if patch_id >= len(self.patch_info):
            return None
        
        try:
            patch_data = self.patch_info[patch_id]
            offset = int(patch_data['offset'])
            size = int(patch_data['size'])
            
            patch_rect = self.patch_rects[patch_id]
            width = int(patch_rect['width'])
            height = int(patch_rect['height'])
            
            image_data = self.patch_data_array[offset:offset + size]
            
            image_2d = image_data.reshape(height, width)
            
            return image_2d
            
        except Exception as e:
            import traceback
            traceback.print_exc()
            return None
    
    def get_particle_image(self, particle_index: int) -> Optional[str]:
        if self.particle_patch_ids is None or self.particle_rects is None:
            return None
        
        if particle_index >= len(self.particle_patch_ids):
            return None
        
        try:
            patch_id = int(self.particle_patch_ids[particle_index])
            patch_image = self._get_patch_image(patch_id)
            if patch_image is None:
                return None
            
            rect_data = self.particle_rects[particle_index]
            
            x = int(rect_data['x'])
            y = int(rect_data['y'])
            w = int(rect_data['width'])
            h = int(rect_data['height'])
            
            if x < 0 or y < 0 or w <= 0 or h <= 0:
                return None
            
            if y + h > patch_image.shape[0]:
                h = patch_image.shape[0] - y
            if x + w > patch_image.shape[1]:
                w = patch_image.shape[1] - x
            
            if w <= 0 or h <= 0:
                return None
            
            cropped = patch_image[y:y+h, x:x+w]
            
            if cropped.size == 0:
                return None
            
            if cropped.dtype != np.uint8:
                img_min, img_max = cropped.min(), cropped.max()
                if img_max > img_min:
                    cropped = ((cropped - img_min) / (img_max - img_min) * 255).astype(np.uint8)
                else:
                    cropped = cropped.astype(np.uint8)
            
            img = Image.fromarray(cropped)
            
            buffered = BytesIO()
            img.save(buffered, format="PNG")
            img_str = base64.b64encode(buffered.getvalue()).decode()
            
            return f'''
<div style="width:{w}px; height:{h}px;">
    <img src="data:image/png;base64,{img_str}"
         style="
            width:{w}px;
            height:{h}px;
            image-rendering: pixelated;
            display:block;
         ">
</div>
'''
            
        except Exception as e:
            import traceback
            traceback.print_exc()
            return None
    
    def has_images(self) -> bool:
        """이미지 데이터가 있는지 확인"""
        return (self.patch_data_array is not None and 
                self.patch_info is not None and
                self.patch_rects is not None and
                self.particle_patch_ids is not None and
                self.particle_rects is not None)
    
    def get_particle_dataframe(self) -> pd.DataFrame:
        if self.cached_dataframe is not None:
            return self.cached_dataframe

        data_dict = {}
            
        exclude_columns = {'cx', 'cy', 'frame_id', 'patch_id', "id"}
        
        if 'particle' not in self.hdf5_file:
            return pd.DataFrame()
        
        particle_group = self.hdf5_file['particle']
        particle_keys = list(particle_group.keys())
        
        # tag 컬럼 먼저 처리 (문자열 타입)
        if 'class' in particle_keys:
            try:
                dataset = particle_group['class']
                if len(dataset.shape) == 1:
                    class_data = dataset[:]
                    # 문자열로 변환
                    converted = []
                    for x in class_data:
                        if isinstance(x, bytes):
                            converted.append(x.decode('utf-8', errors='ignore'))
                        elif isinstance(x, np.bytes_):
                            converted.append(x.decode('utf-8', errors='ignore'))
                        else:
                            converted.append(str(x))
                    data_dict['Tag'] = converted
            except Exception as e:
                print(f"  class : {e}")
        
        for key in particle_keys:
            if key in exclude_columns:
                continue
            
            if key == 'rect' or key == 'class':  # class는 이미 처리함
                continue
                
            try:
                dataset = particle_group[key]
                
                if len(dataset.shape) != 1:
                    continue

                if dataset.dtype.names is not None:
                    continue
                
                if np.issubdtype(dataset.dtype, np.number):
                    formatted_key = self._format_column_name(key)
                    data_dict[formatted_key] = dataset[:]
                    
            except Exception as e:
                print(f"  {key} : {e}")
        
        try:
            self.cached_dataframe = pd.DataFrame(data_dict)
            return self.cached_dataframe
        except Exception as e:
            return pd.DataFrame()
    
    def get_file_attributes(self) -> dict:
        """HDF5 파일의 attributes 반환"""
        if not self.hdf5_file:
            return {}
        
        attrs = {}
        try:
            for key in self.hdf5_file.attrs.keys():
                value = self.hdf5_file.attrs[key]
                # bytes를 string으로 변환
                if isinstance(value, bytes):
                    value = value.decode('utf-8', errors='ignore')
                # numpy 타입을 Python 타입으로 변환
                elif isinstance(value, (np.integer, np.floating)):
                    value = value.item()
                attrs[key] = value
        except Exception as e:
            print(f"Error reading attributes: {e}")
        
        return attrs
    
    def get_total_particle_count(self) -> int:
        if not self.hdf5_file:
            return 0
        
        if 'particle' in self.hdf5_file:
            particle_group = self.hdf5_file['particle']
            keys = list(particle_group.keys())
            for key in keys:
                dataset = particle_group[key]
                if len(dataset.shape) == 1:
                    return len(dataset)
        
        return 0
    
    def save_tags(self, tag_data: list) -> tuple:
        """Tag 데이터를 HDF5 파일의 particle/class에 저장"""
        if not self.file_path or not self.hdf5_file:
            return False, "No file is open"
        
        try:
            # 현재 파일 닫기 (read 모드이므로)
            self.hdf5_file.close()
            self.hdf5_file = None
            
            # read/write 모드로 다시 열기
            with h5py.File(self.file_path, 'r+') as f:
                if 'particle' not in f:
                    return False, "'particle' group not found in HDF5 file"
                
                particle_group = f['particle']
                
                # 문자열을 bytes로 변환
                encoded = []
                for s in tag_data:
                    if isinstance(s, str):
                        encoded.append(s.encode('utf-8'))
                    elif s is None or (isinstance(s, float) and np.isnan(s)):
                        encoded.append(b'')
                    else:
                        encoded.append(str(s).encode('utf-8'))
                
                if 'class' in particle_group:
                    # 기존 dataset 삭제 후 재생성 (크기/타입 변경 대응)
                    del particle_group['class']
                
                # 가변 길이 문자열 타입으로 dataset 생성
                dt = h5py.special_dtype(vlen=bytes)
                particle_group.create_dataset('class', data=encoded, dtype=dt)
            
            # 다시 read 모드로 열기
            self.hdf5_file = h5py.File(self.file_path, 'r')
            self._load_patch_info()
            self._load_particle_info()
            
            return True, "Tags saved successfully"
            
        except Exception as e:
            import traceback
            traceback.print_exc()
            # 파일 다시 열기 시도
            try:
                self.hdf5_file = h5py.File(self.file_path, 'r')
                self._load_patch_info()
                self._load_particle_info()
            except:
                pass
            return False, str(e)
    
    def close(self):
        if self.hdf5_file:
            self.hdf5_file.close()
            self.hdf5_file = None
        self.cached_dataframe = None
        self.patch_data_array = None
        self.patch_info = None
        self.patch_rects = None 
        self.particle_patch_ids = None
        self.particle_rects = None
    
    def __del__(self):
        self.close()