import os
import base64
import pandas as pd
import numpy as np
from pathlib import Path
from io import BytesIO
from PIL import Image

class DataExporter:
    """CSV + Particle PNG (Tag별 폴더) + Graph PNG Export 담당 클래스"""
    
    def __init__(self, reader):
        """
        Args:
            reader: HDF5Reader 인스턴스
        """
        self.reader = reader
        
        # 그래프 이미지 리스트 [{'name': str, 'image': PIL.Image}, ...]
        self.graph_images = []
    
    def set_graph_images(self, graph_images):
        """그래프 이미지 설정
        
        Args:
            graph_images: [{'name': str, 'image': PIL.Image}, ...] 형태의 리스트
        """
        self.graph_images = graph_images or []
    
    def _make_safe_folder_name(self, name):
        """폴더명에 사용 불가능한 문자 제거"""
        # Windows/Linux 공통 사용 불가 문자 제거
        invalid_chars = '<>:"/\\|?*'
        safe_name = name
        for ch in invalid_chars:
            safe_name = safe_name.replace(ch, '_')
        # 앞뒤 공백/점 제거
        safe_name = safe_name.strip().strip('.')
        return safe_name if safe_name else '_unnamed_'
    
    def _generate_unique_folder(self, base_path):
        """고유한 폴더 경로 생성 (이미 존재하면 _1, _2 ... 붙임)"""
        path = Path(base_path)
        if not path.exists():
            return path
        
        n = 1
        while True:
            new_path = Path(f"{base_path}_{n}")
            if not new_path.exists():
                return new_path
            n += 1
    
    def export(self, df, folder_name, folder_location):
        """전체 데이터를 Tag별 폴더 구조로 Export
        
        Args:
            df: 전체 데이터프레임
            folder_name: Export 폴더명
            folder_location: Export 위치 (부모 디렉토리)
        
        Returns:
            (success: bool, result_path_or_error: str)
            
        Export 구조:
            <folder_name>/
            ├── data.csv
            ├── graphs/           (그래프가 있을 때만)
            │   ├── graph_1_bar_AED.png
            │   └── ...
            ├── <Tag1>/
            │   ├── particle_001.png
            │   └── ...
            ├── <Tag2>/
            │   └── ...
            └── (No Tag)/
                └── ...
        """
        try:
            # 루트 폴더 생성
            root_path = self._generate_unique_folder(
                os.path.join(folder_location, folder_name)
            )
            root_path.mkdir(parents=True, exist_ok=True)
            
            # ── 1. CSV 파일 생성 ──
            csv_path = root_path / "data.csv"
            
            # Image 및 Tag 컬럼 제외한 데이터
            columns_to_export = [col for col in df.columns if col not in ['Image', 'Tag']]
            
            # Tag 컬럼은 포함하되 CSV에 넣기
            if 'Tag' in df.columns:
                export_df = df[['Tag'] + columns_to_export].copy()
            else:
                export_df = df[columns_to_export].copy()
            
            export_df.to_csv(csv_path, index=True, index_label='Id')
            print(f"[Export] CSV saved: {csv_path}")
            
            # ── 2. 그래프 PNG 저장 ──
            if len(self.graph_images) > 0:
                graphs_folder = root_path / "graphs"
                graphs_folder.mkdir(exist_ok=True)
                
                for i, graph_data in enumerate(self.graph_images, start=1):
                    graph_name = graph_data.get('name', f'graph_{i}')
                    graph_img = graph_data.get('image', None)
                    
                    if graph_img is not None:
                        safe_name = self._make_safe_folder_name(graph_name)
                        png_path = graphs_folder / f"{safe_name}.png"
                        
                        if isinstance(graph_img, np.ndarray):
                            img = Image.fromarray(graph_img)
                        else:
                            img = graph_img
                        
                        img.save(png_path, 'PNG')
                        print(f"[Export] Graph saved: {png_path}")
            
            # ── 3. Tag별 Particle PNG 저장 ──
            all_indices = df.index.tolist()
            
            # Tag별로 인덱스 그룹화
            tag_groups = {}
            
            if 'Tag' in df.columns:
                for idx in all_indices:
                    tag_value = df.at[idx, 'Tag']
                    
                    # 빈 Tag 처리
                    if pd.isna(tag_value) or str(tag_value).strip() == '':
                        tag_key = '(No Tag)'
                    else:
                        tag_key = str(tag_value).strip()
                    
                    if tag_key not in tag_groups:
                        tag_groups[tag_key] = []
                    tag_groups[tag_key].append(idx)
            else:
                # Tag 컬럼이 없으면 모두 (No Tag)
                tag_groups['(No Tag)'] = all_indices
            
            # 각 Tag별 폴더 생성 및 이미지 저장
            for tag_name, indices in tag_groups.items():
                safe_tag_name = self._make_safe_folder_name(tag_name)
                tag_folder = root_path / safe_tag_name
                tag_folder.mkdir(exist_ok=True)
                
                for i, idx in enumerate(indices, start=1):
                    if idx >= len(df):
                        continue
                    
                    # 이미지 가져오기
                    img_html = self.reader.get_particle_image(idx)
                    
                    if img_html:
                        # base64에서 이미지 추출
                        import re
                        match = re.search(r'src="data:image/png;base64,([^"]+)"', img_html)
                        
                        if match:
                            base64_data = match.group(1)
                            img_bytes = base64.b64decode(base64_data)
                            img = Image.open(BytesIO(img_bytes))
                            
                            png_filename = f"particle_{idx:05d}.png"
                            png_path = tag_folder / png_filename
                            img.save(png_path, 'PNG')
                
                print(f"[Export] Tag '{tag_name}': {len(indices)} particles saved to {tag_folder}")
            
            return True, str(root_path)
        
        except Exception as e:
            import traceback
            traceback.print_exc()
            print(f"Error during export: {e}")
            return False, str(e)